Computer vision YOLO11 model

2025-01-24 10:04:59 -05:00
commit 2e6f7a8a1f
22 changed files with 476 additions and 0 deletions
--- a/Test_track_putposition.py
+++ b/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.")