Spaces:
Running
Running
| """YOLOv8 detection wrapper. | |
| Detects road users (COCO), maps them to display categories, and counts | |
| occupants per vehicle by how much each person box sits inside it. | |
| """ | |
| from collections import Counter | |
| import numpy as np | |
| from app.config import settings | |
| # COCO class id -> (kind, label, display category) | |
| COCO_MAP = { | |
| 0: ("person", "person", "Person"), | |
| 1: ("vehicle", "bicycle", "Bicycle"), | |
| 2: ("vehicle", "car", "Car"), | |
| 3: ("vehicle", "motorcycle", "Two-Wheeler"), | |
| 5: ("vehicle", "bus", "Public Transport"), | |
| 7: ("vehicle", "truck", "Heavy Vehicle"), | |
| 9: ("signal", "traffic light", "Signal"), | |
| } | |
| # A person counts as an occupant if at least this fraction sits inside a vehicle. | |
| MIN_CONTAINMENT = 0.2 | |
| def _containment(person: list[int], vehicle: list[int]) -> float: | |
| """Fraction of the person box area that lies inside the vehicle box.""" | |
| ix1, iy1 = max(person[0], vehicle[0]), max(person[1], vehicle[1]) | |
| ix2, iy2 = min(person[2], vehicle[2]), min(person[3], vehicle[3]) | |
| inter = max(0, ix2 - ix1) * max(0, iy2 - iy1) | |
| area = (person[2] - person[0]) * (person[3] - person[1]) | |
| return inter / area if area else 0.0 | |
| def _assign_occupants(detections: list[dict]) -> None: | |
| """Attribute each person to the vehicle that best contains them.""" | |
| vehicles = [d for d in detections if d["kind"] == "vehicle"] | |
| for person in (d for d in detections if d["kind"] == "person"): | |
| best, best_c = None, MIN_CONTAINMENT | |
| for v in vehicles: | |
| c = _containment(person["bbox"], v["bbox"]) | |
| if c > best_c: | |
| best, best_c = v, c | |
| if best is not None: | |
| best["occupants"] += 1 | |
| def summarize(detections: list[dict]) -> list[dict]: | |
| """Road-user counts grouped by category, most common first (no signals).""" | |
| counts = Counter( | |
| d["category"] for d in detections if d["kind"] in ("vehicle", "person") | |
| ) | |
| return [{"category": k, "count": v} for k, v in counts.most_common()] | |
| class Detector: | |
| """Thin YOLOv8 wrapper. The model is loaded lazily on first use.""" | |
| def __init__(self): | |
| self._model = None | |
| def model(self): | |
| if self._model is None: | |
| from ultralytics import YOLO | |
| self._model = YOLO(settings.yolo_weights) | |
| return self._model | |
| def detect(self, image: np.ndarray) -> list[dict]: | |
| """Run inference (conf filter + NMS are built into YOLO).""" | |
| results = self.model(image, conf=settings.confidence_threshold, verbose=False)[0] | |
| detections = [] | |
| for i, box in enumerate(results.boxes): | |
| cls = int(box.cls[0]) | |
| if cls not in COCO_MAP: | |
| continue | |
| kind, label, category = COCO_MAP[cls] | |
| detections.append({ | |
| "id": i, | |
| "kind": kind, | |
| "label": label, | |
| "category": category, | |
| "bbox": [int(v) for v in box.xyxy[0].tolist()], | |
| "confidence": round(float(box.conf[0]), 3), | |
| "occupants": 0, | |
| }) | |
| _assign_occupants(detections) | |
| return detections | |
| detector = Detector() | |