Spaces:
Sleeping
Sleeping
File size: 1,838 Bytes
b9024f3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 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 |