diff --git a/Code/算法训练与推理/export.py b/Code/算法训练与推理/export.py new file mode 100644 index 0000000..25fd0fc --- /dev/null +++ b/Code/算法训练与推理/export.py @@ -0,0 +1,8 @@ +from ultralytics import YOLO + +# Load a model +model = YOLO('E:/code_files/train/ultralytics-main/model_files/yolov8x.pt') # load an official model +model = YOLO('E:/code_files/jy/20240411_model_test/train/weights/best.pt') # load a custom trained model + +# Export the model +model.export(format='engine',half=True) \ No newline at end of file diff --git a/Code/算法训练与推理/predict.py b/Code/算法训练与推理/predict.py new file mode 100644 index 0000000..81278cd --- /dev/null +++ b/Code/算法训练与推理/predict.py @@ -0,0 +1,70 @@ +from ultralytics import YOLO +from PIL import Image +import cv2 +import os + +# def draw_boxes(image, detections): +# for detection in detections: +# label_name, bbox_list = detection['label'], detection['bbox'] +# confidence = detection['confidence'] + +# # 提取边界框坐标 +# x1, y1, x2, y2 = map(int, bbox_list) + +# # 在图像上绘制边界框 +# cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2) + +# # 在图像上写入标签名和置信度 +# text = f"{label_name}: {confidence}" +# cv2.putText(image, text, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2) + +# return image + +# def analysis_yolov8(frame, model_coco, confidence_set): +# results_coco = model_coco(frame) +# re_list = [] + +# if results_coco: +# for r in results_coco: +# boxes = r.boxes +# for box in boxes: +# b = box.xyxy[0] # 获取边界框坐标 +# c = box.cls +# labels_name = model_coco.names[int(c)] +# confidence = float(box.conf) +# confidence = round(confidence, 2) + +# if confidence < confidence_set: +# continue + + +# re_dict = {'label': labels_name, 'bbox': b, 'confidence': confidence} +# re_list.append(re_dict) + +# return re_list + +# # # 创建新文件夹保存结果 +# output_folder = "E:/code_files/jy/20240411_test_img_file/t_0" +# os.makedirs(output_folder, exist_ok=True) + +# 加载模型和图像 +model = YOLO("E:/code_files/jy/20240411_model_test/train/weights/best.onnx") +img_files = r"E:\code_files\jy\20240423" + +for filename in os.listdir(img_files): + + img_path = os.path.join(img_files, filename) + model.predict(img_path, save=True) + # frame = cv2.imread(img_path) # 替换为你的图像路径 + + # 调用检测函数获取边界框结果 + # detections = analysis_yolov8(frame, model, confidence_set=0.5) + + # # 绘制边界框和标签 + # annotated_image = draw_boxes(frame, detections) + + # # 保存结果图像到新文件夹 + # output_path = os.path.join(output_folder, filename) + # cv2.imwrite(output_path, annotated_image) + + # print(f"Annotated image saved at: {output_path}") diff --git a/Code/算法训练与推理/train.py b/Code/算法训练与推理/train.py new file mode 100644 index 0000000..f1440b5 --- /dev/null +++ b/Code/算法训练与推理/train.py @@ -0,0 +1,11 @@ +from ultralytics import YOLO + +# Load a model +model = YOLO("/home/yaxin/wangying/fd/yolov8/ultralytics/ultralytics/cfg/models/v8/yolov8.yaml") # build a new model from scratch +model = YOLO("/home/yaxin/wangying/fd/yolov8/yolov8n.pt") # load a pretrained model (recommended for training) + +# Use the model +model.train(data="/home/yaxin/wangying/fd/yolov8/person.yaml", epochs=1000) # train the model +metrics = model.val() # evaluate model performance on the validation set +# results = model("https://ultralytics.com/images/bus.jpg") # predict on an image +# path = model.export(format="onnx") # export the model to ONNX format \ No newline at end of file