Merge branch 'main' of http://192.168.10.28:3000/Yaxin/Wafer_Inspection

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)

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}")

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