Update objec_detect_yolo.py

objec_detect_yolo.py

@@ -5,117 +5,97 @@ from ultralytics import YOLO
 import time
 from typing import Tuple, Set
 
-def detection(path: str) -> Tuple[Set[str], str]:
+def detect_objects_in_video(path: str) -> Tuple[Set[str], str]:
     """
-    Detects and tracks objects in a video using YOLOv8 model, saving an annotated output video.
-    
+    Detects and tracks objects in a video using a YOLOv8 model, saving an annotated output video.
+
     Args:
-        path (str): Path to the input video file. Supports common video formats (mp4, avi, etc.)
-        
+        path (str): Path to the input video file.
+
     Returns:
         Tuple[Set[str], str]: 
             - Set of unique detected object labels (e.g., {'Gun', 'Knife'})
             - Path to the output annotated video with detection boxes and tracking IDs
-            
-    Raises:
-        FileNotFoundError: If input video doesn't exist
-        ValueError: If video cannot be opened/processed
     """
-
-    # Validate input file exists
     if not os.path.exists(path):
         raise FileNotFoundError(f"Video file not found: {path}")
 
-    # Initialize YOLOv8 model with pretrained weights
-    # Model is trained to detect: ['Fire', 'Gun', 'License_Plate', 'Smoke', 'knife']
-    model = YOLO("yolo"/"best.pt"))
-    class_names = model.names  # Get class label mappings
+    # Load YOLOv8 model (adjust path if necessary)
+    model = YOLO("yolo/best.pt")  # Make sure this path is correct
+    class_names = model.names
 
-    # Set up output paths:
-    # 1. Temporary output during processing
-    # 2. Final output with detected objects in filename
+    # Output setup
     input_video_name = os.path.basename(path)
     base_name = os.path.splitext(input_video_name)[0]
     temp_output_name = f"{base_name}_output_temp.mp4"
     output_dir = "results"
-    os.makedirs(output_dir, exist_ok=True)  # Create output dir if needed
-    if not os.path.exists(output_dir):
-        raise ValueError(f"Failed to create output directory: {output_dir}")
+    os.makedirs(output_dir, exist_ok=True)
     temp_output_path = os.path.join(output_dir, temp_output_name)
 
-    # Video processing setup:
-    # - Open input video stream
-    # - Initialize output writer with MP4 codec
+    # Video I/O setup
     cap = cv2.VideoCapture(path)
     if not cap.isOpened():
         raise ValueError(f"Failed to open video file: {path}")
 
-    # Process all frames at 640x640 resolution for consistency
     frame_width, frame_height = 640, 640
     out = cv2.VideoWriter(
-        temp_output_path, 
-        cv2.VideoWriter_fourcc(*'mp4v'),  # MP4 codec
-        30.0,  # Output FPS
+        temp_output_path,
+        cv2.VideoWriter_fourcc(*'mp4v'),
+        30.0,
         (frame_width, frame_height)
     )
 
-    # Main processing loop:
-    # 1. Read each frame
-    # 2. Run object detection + tracking
-    # 3. Annotate frame with boxes and IDs
-    # 4. Collect detected classes
-    crimes = []  # Track all detected objects
+    detected_labels = set()
     start = time.time()
     print(f"[INFO] Processing started at {start:.2f} seconds")
 
     while True:
         ret, frame = cap.read()
-        if not ret:  # End of video
+        if not ret:
             break
 
-        # Resize and run detection + tracking
         frame = cv2.resize(frame, (frame_width, frame_height))
+
+        # Run detection and tracking
         results = model.track(
             source=frame,
-            conf=0.7,  # Minimum confidence threshold
-            persist=True  # Enable tracking across frames
+            conf=0.7,
+            persist=True
         )
 
-        # Annotate frame with boxes and tracking IDs
-        annotated_frame = results[0].plot()
-        
-        # Record detected classes
-        for box in results[0].boxes:
-            cls = int(box.cls)
-            crimes.append(class_names[cls])
+        if results and hasattr(results[0], "plot"):
+            annotated_frame = results[0].plot()
+            out.write(annotated_frame)
 
-        out.write(annotated_frame)
+            # Extract class labels
+            if hasattr(results[0], "boxes"):
+                for box in results[0].boxes:
+                    cls = int(box.cls)
+                    detected_labels.add(class_names[cls])
+        else:
+            out.write(frame)
 
-    # Clean up video resources
     end = time.time()
-    print(f"[INFO] Processing finished at {end:.2f} seconds")
-    print(f"[INFO] Total execution time: {end - start:.2f} seconds")
     cap.release()
     out.release()
 
-    # Generate final output filename containing detected object labels
-    # Format: {original_name}_{detected_objects}_output.mp4
-    unique_crimes = set(crimes)
-    crimes_str = "_".join(sorted(unique_crimes)).replace(" ", "_")[:50]  # truncate if needed
+    # Create final output filename
+    crimes_str = "_".join(sorted(detected_labels)).replace(" ", "_")[:50]
     final_output_name = f"{base_name}_{crimes_str}_output.mp4"
     final_output_path = os.path.join(output_dir, final_output_name)
 
-    # Rename the video file
     os.rename(temp_output_path, final_output_path)
 
-    print(f"[INFO] Detected crimes: {unique_crimes}")
+    print(f"[INFO] Processing finished at {end:.2f} seconds")
+    print(f"[INFO] Total execution time: {end - start:.2f} seconds")
+    print(f"[INFO] Detected crimes: {detected_labels}")
     print(f"[INFO] Annotated video saved at: {final_output_path}")
 
-    return unique_crimes, final_output_path
+    return detected_labels, final_output_path
 
 
-# # Entry point
-# path0 = input("Enter the local path to the video file to detect objects: ")
-# path = path0.strip('"')  # Remove extra quotes if copied from Windows
-# print(f"[INFO] Loading video: {path}")
-# detection(path)
+# Example usage (uncomment to use as standalone script)
+# if __name__ == "__main__":
+#     video_path = input("Enter the path to the video file: ").strip('"')
+#     print(f"[INFO] Loading video: {video_path}")
+#     detect_objects_in_video(video_path)