Create object_detection.py

services/object_detection.py

@@ -0,0 +1,54 @@
+import cv2
+import numpy as np
+from ultralytics import YOLO
+from typing import List, Tuple, Dict, Any
+
+# Load YOLOv8 model for general object detection
+model = YOLO("models/yolov8n.pt")
+
+def detect_objects(frame: np.ndarray) -> Tuple[List[Dict[str, Any]], np.ndarray]:
+    """
+    Detect cars, bikes, humans, dogs, and other objects in the frame using YOLOv8.
+    Args:
+        frame: Input frame as a numpy array.
+    Returns:
+        Tuple of (list of detections, annotated frame).
+    """
+    # Perform inference
+    results = model(frame, conf=0.5)  # Detect all classes with confidence > 0.5
+    
+    detections = []
+    for i, r in enumerate(results[0].boxes):
+        x_min, y_min, x_max, y_max = map(int, r.xyxy[0])
+        conf = float(r.conf)
+        cls = int(r.cls)
+        # Map YOLOv8 class IDs to labels (based on COCO dataset)
+        label_map = {
+            0: "person",  # Human
+            1: "bicycle", # Bike (approximation)
+            2: "car",
+            3: "motorcycle", # Bike
+            5: "bus",     # Treat as car
+            7: "truck",   # Treat as car
+            16: "dog",
+        }
+        dtype = label_map.get(cls, "object")  # Default to "object" for unmapped classes
+        if dtype in ["bicycle", "motorcycle"]:
+            dtype = "bike"
+        elif dtype in ["bus", "truck"]:
+            dtype = "car"
+        label = f"{dtype.capitalize()} {i+1}"
+        
+        # Determine severity (not used for objects, but included for consistency)
+        area = (x_max - x_min) * (y_max - y_min)
+        severity = "Moderate"  # Default for objects
+        
+        detections.append({
+            "box": [x_min, y_min, x_max, y_max],
+            "label": label,
+            "type": dtype,
+            "confidence": conf,
+            "severity": severity
+        })
+    
+    return detections, frame