Spaces:
Sleeping
Sleeping
| 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 |