Update objec_detect_yolo.py

objec_detect_yolo.py

@@ -1,101 +1,217 @@
+--- START OF FILE objec_detect_yolo.py ---
+
 import cv2
 import numpy as np
 import os
 from ultralytics import YOLO
 import time
-from typing import Tuple, Set
+from typing import Tuple, Set, List
 
-def detect_objects_in_video(path: str) -> Tuple[Set[str], str]:
+def detection(path: str) -> Tuple[Set[str], str]:
     """
-    Detects and tracks objects in a video using a YOLOv8 model, saving an annotated output video.
+    Detects and tracks objects in a video using YOLOv8 model, saving an annotated output video.
 
     Args:
-        path (str): Path to the input video file.
+        path (str): Path to the input video file. Supports common video formats (mp4, avi, etc.)
 
     Returns:
-        Tuple[Set[str], str]: 
+        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 or output dir cannot be created
     """
+
+    # Validate input file exists
     if not os.path.exists(path):
         raise FileNotFoundError(f"Video file not found: {path}")
 
-    # Load YOLOv8 model (adjust path if necessary)
-    model = YOLO("yolo/best.pt")  # Make sure this path is correct
-    class_names = model.names
+    # --- Model Loading ---
+    # Construct path relative to this script file
+    model_path = os.path.join(os.path.dirname(__file__), "yolo", "best.pt")
+    if not os.path.exists(model_path):
+         raise FileNotFoundError(f"YOLO model file not found at: {model_path}")
+    try:
+        model = YOLO(model_path)
+        class_names = model.names  # Get class label mappings
+        print(f"[INFO] YOLO model loaded from {model_path}. Class names: {class_names}")
+    except Exception as e:
+        raise ValueError(f"Failed to load YOLO model: {e}")
 
-    # Output setup
+
+    # --- Output Path 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)
+    # Sanitize basename to prevent issues with weird characters in filenames
+    safe_base_name = "".join(c if c.isalnum() or c in ('-', '_') else '_' for c in base_name)
+
+    # Define output directory relative to this script
+    # In HF Spaces, this will be inside the container's file system
+    output_dir = os.path.join(os.path.dirname(__file__), "results")
+    temp_output_name = f"{safe_base_name}_output_temp.mp4"
+
+    try:
+        os.makedirs(output_dir, exist_ok=True)  # Create output dir if needed
+        if not os.path.isdir(output_dir):
+             raise ValueError(f"Path exists but is not a directory: {output_dir}")
+    except OSError as e:
+        raise ValueError(f"Failed to create or access output directory '{output_dir}': {e}")
+
     temp_output_path = os.path.join(output_dir, temp_output_name)
+    print(f"[INFO] Temporary output will be saved to: {temp_output_path}")
+
 
-    # Video I/O setup
+    # --- Video Processing Setup ---
     cap = cv2.VideoCapture(path)
     if not cap.isOpened():
         raise ValueError(f"Failed to open video file: {path}")
 
+    # Get video properties for output writer
+    # Use source FPS if available and reasonable, otherwise default to 30
+    source_fps = cap.get(cv2.CAP_PROP_FPS)
+    output_fps = source_fps if 10 <= source_fps <= 60 else 30.0
+
+    # Process at a fixed resolution for consistency or use source resolution
+    # Using fixed 640x640 as potentially used during training/fine-tuning
     frame_width, frame_height = 640, 640
-    out = cv2.VideoWriter(
-        temp_output_path,
-        cv2.VideoWriter_fourcc(*'mp4v'),
-        30.0,
-        (frame_width, frame_height)
-    )
-
-    detected_labels = set()
+    # OR use source resolution (might require adjusting YOLO parameters if model expects specific size)
+    # frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    # frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+    try:
+        out = cv2.VideoWriter(
+            temp_output_path,
+            cv2.VideoWriter_fourcc(*'mp4v'),  # Use MP4 codec
+            output_fps,
+            (frame_width, frame_height)
+        )
+        if not out.isOpened():
+            # Attempt alternative codec if mp4v fails (less common)
+            print("[WARNING] mp4v codec failed, trying avc1...")
+            out = cv2.VideoWriter(
+                temp_output_path,
+                cv2.VideoWriter_fourcc(*'avc1'),
+                output_fps,
+                (frame_width, frame_height)
+            )
+            if not out.isOpened():
+                 raise ValueError("Failed to initialize VideoWriter with mp4v or avc1 codec.")
+
+    except Exception as e:
+        cap.release() # Release capture device before raising
+        raise ValueError(f"Failed to create VideoWriter: {e}")
+
+
+    # --- Main Processing Loop ---
+    detected_classes: List[str] = []  # Track detected object class names
     start = time.time()
-    print(f"[INFO] Processing started at {start:.2f} seconds")
+    frame_count = 0
+    print(f"[INFO] Video processing started...")
 
     while True:
         ret, frame = cap.read()
-        if not ret:
+        if not ret:  # End of video or read error
             break
 
-        frame = cv2.resize(frame, (frame_width, frame_height))
-
-        # Run detection and tracking
-        results = model.track(
-            source=frame,
-            conf=0.7,
-            persist=True
-        )
-
-        if results and hasattr(results[0], "plot"):
-            annotated_frame = results[0].plot()
+        frame_count += 1
+        # Resize frame BEFORE passing to model
+        resized_frame = cv2.resize(frame, (frame_width, frame_height))
+
+        try:
+            # Run YOLOv8 detection and tracking on the resized frame
+            results = model.track(
+                source=resized_frame, # Use resized frame
+                conf=0.7,              # Confidence threshold
+                persist=True,          # Maintain track IDs across frames
+                verbose=False          # Suppress Ultralytics console output per frame
+            )
+
+            # Check if results are valid and contain boxes
+            if results and results[0] and results[0].boxes:
+                # Annotate the RESIZED frame with bounding boxes and track IDs
+                annotated_frame = results[0].plot() # plot() draws on the source image
+
+                # Record detected class names for this frame
+                for box in results[0].boxes:
+                    if box.cls is not None: # Check if class ID is present
+                        cls_id = int(box.cls[0]) # Get class index
+                        if 0 <= cls_id < len(class_names):
+                            detected_classes.append(class_names[cls_id])
+                        else:
+                            print(f"[WARNING] Detected unknown class ID: {cls_id}")
+            else:
+                 # If no detections, use the original resized frame for the output video
+                 annotated_frame = resized_frame
+
+            # Write the (potentially annotated) frame to the output video
             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)
+        except Exception as e:
+             print(f"[ERROR] Error processing frame {frame_count}: {e}")
+             # Write the unannotated frame to keep video timing consistent
+             out.write(resized_frame)
+
 
+    # --- Clean Up ---
     end = time.time()
+    print(f"[INFO] Video processing finished. Processed {frame_count} frames.")
+    print(f"[INFO] Total processing time: {end - start:.2f} seconds")
     cap.release()
     out.release()
+    cv2.destroyAllWindows() # Close any OpenCV windows if they were opened
+
 
-    # Create final output filename
-    crimes_str = "_".join(sorted(detected_labels)).replace(" ", "_")[:50]
-    final_output_name = f"{base_name}_{crimes_str}_output.mp4"
+    # --- Final Output Renaming ---
+    unique_detected_labels = set(detected_classes)
+    # Create a short string from labels for the filename
+    labels_str = "_".join(sorted(list(unique_detected_labels))).replace(" ", "_")
+    # Limit length to avoid overly long filenames
+    max_label_len = 50
+    if len(labels_str) > max_label_len:
+        labels_str = labels_str[:max_label_len] + "_etc"
+    if not labels_str: # Handle case where nothing was detected
+        labels_str = "no_detections"
+
+    final_output_name = f"{safe_base_name}_{labels_str}_output.mp4"
     final_output_path = os.path.join(output_dir, final_output_name)
 
-    os.rename(temp_output_path, final_output_path)
+    # Ensure final path doesn't already exist (rename might fail otherwise)
+    if os.path.exists(final_output_path):
+        os.remove(final_output_path)
+
+    try:
+        # Rename the temporary file to the final name
+        os.rename(temp_output_path, final_output_path)
+        print(f"[INFO] Detected object labels: {unique_detected_labels}")
+        print(f"[INFO] Annotated video saved successfully at: {final_output_path}")
+    except OSError as e:
+        print(f"[ERROR] Failed to rename {temp_output_path} to {final_output_path}: {e}")
+        # Fallback: return the temp path if rename fails but file exists
+        if os.path.exists(temp_output_path):
+             print(f"[WARNING] Returning path to temporary file: {temp_output_path}")
+             return unique_detected_labels, temp_output_path
+        else:
+             raise ValueError(f"Output video generation failed. No output file found.")
 
-    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 detected_labels, final_output_path
+    return unique_detected_labels, final_output_path
 
 
-# Example usage (uncomment to use as standalone script)
+# # Example usage (commented out for library use)
 # 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)
+#     test_video = input("Enter the local path to the video file: ").strip('"')
+#     if os.path.exists(test_video):
+#         try:
+#             print(f"[INFO] Processing video: {test_video}")
+#             labels, out_path = detection(test_video)
+#             print(f"\nDetection Complete.")
+#             print(f"Detected unique labels: {labels}")
+#             print(f"Output video saved to: {out_path}")
+#         except (FileNotFoundError, ValueError, Exception) as e:
+#             print(f"\nAn error occurred: {e}")
+#     else:
+#         print(f"Error: Input video file not found - {test_video}")
+
+--- END OF FILE objec_detect_yolo.py ---