Computer vision YOLO11 model

Test_Realtime.py

@@ -0,0 +1,94 @@
+import cv2
+from ultralytics import YOLO
+
+# Camera index (default camera is 0)
+camera_index = 0
+
+# Load the YOLO model
+model = YOLO(r"D:\AIM\lemon\runs\detect\train4\weights\best.pt")  # Load custom model
+
+# Initialize the camera
+cap = cv2.VideoCapture(camera_index)
+if not cap.isOpened():
+    print("Unable to open the camera. Please check the device.")
+    exit()
+
+fps = int(cap.get(cv2.CAP_PROP_FPS))
+width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+print(f"Camera resolution: {width}x{height}, FPS: {fps} FPS")
+
+# Dictionary to track the state of each lemon
+lemon_states = {}  # Format: {ID: "State"}
+
+# Define class labels
+class_labels = {
+    0: "Bruised",
+    1: "DefectiveLemon",
+    2: "GoodLemon",
+    3: "NotRipeLemon",
+    4: "Rotten"
+}
+
+# Classes that require ID assignment
+id_tracked_classes = ["DefectiveLemon", "GoodLemon", "NotRipeLemon"]
+
+# Set the window to be resizable
+cv2.namedWindow("Live Detection", cv2.WINDOW_NORMAL)
+
+# Process video stream in real-time
+while cap.isOpened():
+    ret, frame = cap.read()
+    if not ret:
+        print("Unable to read camera input. Terminating program.")
+        break
+
+    # Perform object tracking using BoT-SORT
+    results = model.track(source=frame, conf=0.5, tracker='botsort.yaml', show=False)
+
+    for result in results:
+        frame = result.orig_img  # Current frame
+        detections = result.boxes  # Detection box information
+
+        for box in detections:
+            x1, y1, x2, y2 = map(int, box.xyxy[0])  # Detection box coordinates
+            obj_id = int(box.id) if box.id is not None else -1  # Tracking object ID
+            class_id = int(box.cls)  # Class ID
+            score = box.conf  # Confidence score
+            label = class_labels.get(class_id, "Unknown")  # Get class label
+
+            # Update lemon state and output information for tracked boxes
+            if obj_id != -1 and label in id_tracked_classes:
+                if obj_id not in lemon_states:
+                    lemon_states[obj_id] = label
+                else:
+                    # Once detected as "DefectiveLemon," the state remains "DefectiveLemon"
+                    if lemon_states[obj_id] != "DefectiveLemon":
+                        lemon_states[obj_id] = label
+
+                # Output ID, position, and label
+                position = f"({x1}, {y1}, {x2}, {y2})"
+                print(f"ID: {obj_id}, Position: {position}, Label: {lemon_states[obj_id]}")
+
+            # Draw detection boxes and labels (including untracked ones)
+            if obj_id != -1 and label in id_tracked_classes:
+                display_text = f"ID {obj_id} | {lemon_states[obj_id]}"
+            else:
+                display_text = label  # For untracked labels, only show the class
+
+            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
+            cv2.putText(frame, display_text, (x1, y1 - 10),
+                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
+
+    # Display the processed video stream
+    cv2.imshow("Live Detection", frame)
+
+    # Exit the loop when ESC key is pressed
+    if cv2.waitKey(1) & 0xFF == 27:  # 27 is the ASCII value for ESC key
+        print("ESC key detected. Exiting the program.")
+        break
+
+# Release resources
+cap.release()
+cv2.destroyAllWindows()
+print("Camera video processing complete. Program terminated.")

Test_image.py

@@ -0,0 +1,66 @@
+import os
+import cv2
+import pandas as pd
+from ultralytics import YOLO
+
+# Define input and output directories
+input_dir = r'D:\AIM\pecan\test_image\input_image'
+output_dir = r'D:\AIM\pecan\test_image\output_image'
+
+# Create the output directory if it doesn't exist
+os.makedirs(output_dir, exist_ok=True)
+
+# Load a YOLO model (you can specify your model path here)
+model = YOLO(r"D:\AIM\pecan\runs\detect\train2\weights\best.pt")  # Load a custom-trained model
+
+# Initialize a threshold and a list for classifications
+threshold = 0.5
+all_classifications = []
+
+# Iterate over all images in the input directory
+for image_filename in os.listdir(input_dir):
+    if image_filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff')):  # Check for image files
+        image_path = os.path.join(input_dir, image_filename)
+        image = cv2.imread(image_path)
+        if image is None:
+            continue  # Skip files that cannot be read
+
+        # Perform inference
+        results = model(image)[0]
+        classifications = []
+
+        # Create a semi-transparent overlay for the bounding boxes
+        for result in results.boxes.data.tolist():
+            x1, y1, x2, y2, score, class_id = result
+
+            if score > threshold:
+                # Create a semi-transparent overlay for the bounding box area
+                overlay = image.copy()
+                cv2.rectangle(overlay, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), -1)
+                alpha = 0.15  # Transparency factor (0 - transparent, 1 - opaque)
+                # Apply the overlay
+                image = cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0)
+
+                # Draw bounding box outline
+                cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
+                # Prepare label with class name and confidence score
+                label = f"{results.names[int(class_id)].upper()} {score:.2f}"
+                # Put text (class label with confidence score)
+                cv2.putText(image, label, (int(x1), int(y1 - 10)),
+                            cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2, cv2.LINE_AA)
+
+                # Add classification to the list
+                classifications.append(f"{results.names[int(class_id)]} ({score:.2f})")
+
+        # Save the output image
+        output_path = os.path.join(output_dir, image_filename)
+        cv2.imwrite(output_path, image)
+
+        # Save classifications for each image
+        all_classifications.append({"image": image_filename, "classes": classifications})
+
+# Create a DataFrame from the classifications
+df = pd.DataFrame(all_classifications)
+df.to_excel(os.path.join(output_dir, 'output_classifications.xlsx'), index=False)
+
+cv2.destroyAllWindows()

Test_logic.py

@@ -0,0 +1,52 @@
+import os
+import pandas as pd
+from ultralytics import YOLO
+import cv2
+
+# Input and output video paths
+video_path = r'D:\AIM\pecan\GH014359.mp4'
+video_path_out = r'D:\AIM\pecan\GH014359_out.mp4'
+
+cap = cv2.VideoCapture(video_path)
+ret, frame = cap.read()
+H, W, _ = frame.shape
+out = cv2.VideoWriter(video_path_out, cv2.VideoWriter_fourcc(*'MP4V'), int(cap.get(cv2.CAP_PROP_FPS)), (W, H))
+
+# Load the YOLO model
+model = YOLO(r"D:\AIM\pecan\runs\detect\train2\weights\best.pt")  # Load a custom model
+
+threshold = 0.5
+detected_cracked = False  # Initialize a flag for detecting cracked pecans
+
+while ret:
+    # Perform detection on the current frame
+    results = model(frame)[0]
+
+    for result in results.boxes.data.tolist():
+        x1, y1, x2, y2, score, class_id = result
+
+        if score > threshold:
+            # Draw bounding boxes and labels
+            cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 4)
+            label = results.names[int(class_id)].upper()
+            cv2.putText(frame, f"{label} {score:.2f}", (int(x1), int(y1 - 10)),
+                        cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 255, 0), 3, cv2.LINE_AA)
+
+            # Check for the "cracked" label
+            if label == "CRACKED":
+                detected_cracked = True
+
+    # Write the processed frame to the output video
+    out.write(frame)
+    ret, frame = cap.read()
+
+# Determine the final label based on detections
+final_label = "CRACKED" if detected_cracked else "GOOD"
+
+# Print the final label
+print(f"Final Label: {final_label}")
+
+# Release video resources
+cap.release()
+out.release()
+cv2.destroyAllWindows()

Test_logic_track.py

@@ -0,0 +1,80 @@
+import os
+import cv2
+from ultralytics import YOLO
+
+# 输入和输出视频路径
+# video_path = r'D:\AIM\pecan\OneDrive_2_2024-11-7\G5 Flex 01 8-5-2024, 1.20.12pm EDT - 8-5-2024, 1.23.22pm EDT.mp4'
+# video_path_out = r'D:\AIM\pecan\G5 Flex 01 8-5-2024_out.mp4'
+video_path = r'D:\AIM\pecan\GH014359.mp4'
+video_path_out = r'D:\AIM\pecan\GH014359_out.mp4'
+
+# 加载 YOLO 模型
+model = YOLO(r"D:\AIM\pecan\runs\detect\train2\weights\best.pt")  # 加载自定义模型
+
+# 初始化 VideoWriter 用于保存输出视频
+cap = cv2.VideoCapture(video_path)
+fps = int(cap.get(cv2.CAP_PROP_FPS))
+width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+out = cv2.VideoWriter(video_path_out, fourcc, fps, (width, height))
+
+# 字典，用于跟踪每个核桃的状态
+walnut_states = {}  # 格式: {ID: "状态"}
+
+# 定义类别标签
+class_labels = {
+    0: "circumferential",
+    1: "cracked",
+    2: "crushed",
+    3: "longitudinal",
+    4: "open",
+    5: "uncracked"
+}
+
+# 需要分配 ID 的类别
+id_tracked_classes = ["cracked", "uncracked"]
+
+# 使用 BoT-SORT 进行目标跟踪
+results = model.track(source=video_path, conf=0.5, tracker='botsort.yaml', show=False)
+
+for result in results:
+    frame = result.orig_img  # 当前帧
+    detections = result.boxes  # 检测框信息
+
+    # 处理每个检测框
+    for box in detections:
+        x1, y1, x2, y2 = map(int, box.xyxy[0])  # 检测框坐标
+        obj_id = int(box.id) if box.id is not None else -1  # 跟踪目标ID
+        class_id = int(box.cls)  # 类别ID
+        score = box.conf  # 置信度
+
+        # 获取检测框对应的类别标签
+        label = class_labels.get(class_id, "unknown")
+
+        # 仅对需要分配ID的类别更新核桃状态
+        if label in id_tracked_classes:
+            if obj_id not in walnut_states:
+                walnut_states[obj_id] = label
+            else:
+                # 一旦检测到“cracked”，状态保持为“cracked”
+                if walnut_states[obj_id] != "cracked":
+                    walnut_states[obj_id] = label
+
+            display_text = f"ID {obj_id} | {walnut_states[obj_id]}"
+        else:
+            # 非分配ID的类别，仅显示类别标签
+            display_text = label
+
+        # 绘制检测框和标签
+        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
+        cv2.putText(frame, display_text, (x1, y1 - 10),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
+
+    # 将处理好的帧写入输出视频
+    out.write(frame)
+
+# 释放资源
+cap.release()
+out.release()
+print("视频处理完成，结果已保存。")

Test_track_putposition.py

@@ -0,0 +1,80 @@
+import os
+import cv2
+from ultralytics import YOLO
+
+# Input and output video paths
+video_path = r'D:\AIM\lemon\Lemon Videos 1\Lemon Videos\test\7.mp4'
+video_path_out = r'D:\AIM\lemon\Lemon Videos 1\Lemon Videos\test\7_out.mp4'
+# video_path = r'D:\AIM\lemon\Lemon Videos 1\Lemon Videos\Bad Lemons\8.mp4'
+# video_path_out = r'D:\AIM\lemon\Lemon Videos 1\Lemon Videos\Bad Lemons\8_out.mp4'
+
+# Load the YOLO model
+model = YOLO(r"D:\AIM\lemon\runs\detect\train4\weights\best.pt")  # Load the custom model
+
+# Initialize VideoWriter to save the output video
+cap = cv2.VideoCapture(video_path)
+fps = int(cap.get(cv2.CAP_PROP_FPS))
+width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+out = cv2.VideoWriter(video_path_out, fourcc, fps, (width, height))
+
+# Dictionary to track the state of each lemon
+lemon_states = {}  # Format: {ID: "State"}
+
+# Define class labels
+class_labels = {
+    0: "Bruised",
+    1: "DefectiveLemon",
+    2: "GoodLemon",
+    3: "NotRipeLemon",
+    4: "Rotten"
+}
+
+# Classes that require ID assignment
+id_tracked_classes = ["DefectiveLemon", "GoodLemon", "NotRipeLemon"]
+
+# Perform object tracking using BoT-SORT
+results = model.track(source=video_path, conf=0.5, tracker='botsort.yaml', show=False)
+
+for result in results:
+    frame = result.orig_img  # Current frame
+    detections = result.boxes  # Detection box information
+
+    for box in detections:
+        x1, y1, x2, y2 = map(int, box.xyxy[0])  # Detection box coordinates
+        obj_id = int(box.id) if box.id is not None else -1  # Tracking object ID
+        class_id = int(box.cls)  # Class ID
+        score = box.conf  # Confidence score
+        label = class_labels.get(class_id, "Unknown")  # Get class label
+
+        # Update lemon state and output information for tracked boxes
+        if obj_id != -1 and label in id_tracked_classes:
+            if obj_id not in lemon_states:
+                lemon_states[obj_id] = label
+            else:
+                # Once detected as "DefectiveLemon," the state remains "DefectiveLemon"
+                if lemon_states[obj_id] != "DefectiveLemon":
+                    lemon_states[obj_id] = label
+
+            # Output ID, position, and label
+            position = f"({x1}, {y1}, {x2}, {y2})"
+            print(f"ID: {obj_id}, Position: {position}, Label: {lemon_states[obj_id]}")
+
+        # Draw detection boxes and labels (including untracked ones)
+        if obj_id != -1 and label in id_tracked_classes:
+            display_text = f"ID {obj_id} | {lemon_states[obj_id]}"
+        else:
+            display_text = label  # For untracked labels, only show the class
+
+        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
+        cv2.putText(frame, display_text, (x1, y1 - 10),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
+
+    # Write the processed frame to the output video
+    out.write(frame)
+
+# Release resources
+cap.release()
+out.release()
+print("Video processing completed, and the result has been saved.")

Test_video.py

@@ -0,0 +1,52 @@
+import os
+import pandas as pd
+from ultralytics import YOLO
+import cv2
+
+video_path = r'D:\AIM\lemon\test.mp4'
+video_path_out = r'D:\AIM\lemon\test_out.mp4'
+
+cap = cv2.VideoCapture(video_path)
+ret, frame = cap.read()
+H, W, _ = frame.shape
+out = cv2.VideoWriter(video_path_out, cv2.VideoWriter_fourcc(*'MP4V'), int(cap.get(cv2.CAP_PROP_FPS)), (W, H))
+
+model_path = os.path.join('.', 'runs', 'detect', 'train', 'weights', 'last.pt')
+
+# Load a model
+model = YOLO(r"D:\AIM\lemon\YOLO-Training\YOLOv11_finetune\weights\best.pt")  # load a custom model
+
+threshold = 0.5
+classifications = []
+
+while ret:
+
+    results = model(frame)[0]
+
+    for result in results.boxes.data.tolist():
+        x1, y1, x2, y2, score, class_id = result
+
+        if score > threshold:
+            cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 4)
+            cv2.putText(frame, results.names[int(class_id)].upper(), (int(x1), int(y1 - 10)),
+                        cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 255, 0), 3, cv2.LINE_AA)
+            
+            # percentage = 
+            
+            # if results.names[int(class_id)] in classifications:
+                
+            # else:
+            classifications.append(results.names[int(class_id)])
+
+    out.write(frame)
+    ret, frame = cap.read()
+
+# Create a new DataFrame from the classifications
+new_df = pd.DataFrame(classifications)
+
+# Write the new DataFrame to an Excel file
+new_df.to_excel('output_classifications.xlsx', index=False)
+
+cap.release()
+out.release()
+cv2.destroyAllWindows()

Train.py

@@ -0,0 +1,33 @@
+from ultralytics import YOLO
+import os
+
+os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+def main():
+    # 初始化 Weights & Biases
+    import wandb
+    wandb.login(key='7cfbcf76a18a8441b04eb5d7adb988e69a79705e')  # 替换为你的 API 密钥
+    wandb.init(project="YOLO-Training", name="YOLOv11_finetune", mode="online")
+
+    # 加载 YOLO 模型
+    model = YOLO('yolo11n.pt')
+
+    # 开始训练
+    model.train(
+        data=r"D:\AIM\lemon\lemon_quality_dataset_YOLO11\data.yaml",  # 数据集配置文件路径
+        epochs=2,                              # 训练轮次
+        imgsz=640,                             # 图像大小
+        batch=16,                              # 批量大小
+        lr0=0.005,                              # 初始学习率
+        workers=0,                             # 设置为 0，禁用多进程数据加载
+        device=0,                              # 设备 (0 = 第一块 GPU, 'cpu' = CPU)
+        project="YOLO-Training",               # W&B 项目名称
+        name="YOLOv11_finetune",               # W&B 实验名称
+        exist_ok=True                          # 如果目录存在是否覆盖
+    )
+
+if __name__ == '__main__':
+    import torch.multiprocessing as mp
+    mp.freeze_support()  # 解决 Windows 下的多进程启动问题
+    main()

Train_model.py

@@ -0,0 +1,19 @@
+from ultralytics import YOLO
+import os
+
+os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
+os.environ["CUDA_VISIBLE_DEVICES"]="0"
+
+
+if __name__ == '__main__':
+    model = YOLO('yolo11n.pt')
+
+    results = model.train(
+        data = "D:\AIM\lemon\AIM.v3i.yolov11\data.yaml",       # Path to the dataset (YAML file for COCO128)
+        epochs = 50,                              # Number of epochs for fine-tuning
+        imgsz = 640,                             # Image size for training
+        batch = 16,                              # Batch size
+        lr0 = 0.005,                              # Initial learning rate
+        workers = 4,                             # Number of data loading workers
+        device = 0                               # Specify the device (0 = first GPU, 'cpu' for CPU)
+    )