""" VehicleDetector — YOLOv8-based vehicle detection, classification, and person counting. Also handles license plate OCR via fast-alpr. """ import cv2 import numpy as np import logging from pathlib import Path from typing import Optional from dataclasses import dataclass, field from app.core.config import settings logger = logging.getLogger(__name__) VEHICLE_CLASSES = { 2: "car", 3: "motorcycle", 5: "bus", 7: "truck", } # COCO person class PERSON_CLASS = 0 # Color detection ranges (HSV) COLOR_RANGES = { "red": [(0, 70, 50), (10, 255, 255)], "red2": [(170, 70, 50), (180, 255, 255)], "blue": [(100, 70, 50), (130, 255, 255)], "white": [(0, 0, 180), (180, 30, 255)], "black": [(0, 0, 0), (180, 255, 50)], "silver": [(0, 0, 100), (180, 30, 180)], "yellow": [(20, 70, 50), (35, 255, 255)], "green": [(35, 70, 50), (85, 255, 255)], "orange": [(10, 70, 50), (20, 255, 255)], } @dataclass class VehicleDetection: vehicle_type: str = "unknown" color: Optional[str] = None plate_number: Optional[str] = None plate_confidence: float = 0.0 person_count: int = 1 detection_confidence: float = 0.0 vehicle_bbox: dict = field(default_factory=dict) plate_bbox: dict = field(default_factory=dict) vehicle_crop: Optional[np.ndarray] = None plate_crop: Optional[np.ndarray] = None class VehicleDetector: def __init__(self): self._yolo = None self._alpr = None self._ready = False def load(self): """Load models — call once at startup.""" try: from ultralytics import YOLO self._yolo = YOLO(settings.YOLO_MODEL) logger.info("YOLO model loaded") except Exception as e: logger.error(f"Failed to load YOLO: {e}") return try: from fast_alpr import ALPR self._alpr = ALPR( detector_model="plate-detection-v1-large", ocr_model="global-plates-mobile-vit-v2-model", device="cuda" if settings.GPU_DEVICE >= 0 else "cpu", ) logger.info("ALPR model loaded") except Exception as e: logger.warning(f"ALPR not loaded (will skip plate reading): {e}") self._ready = True def detect(self, frame: np.ndarray) -> list[VehicleDetection]: """ Run vehicle detection on a frame. Returns list of VehicleDetection objects. """ if not self._ready or self._yolo is None: return [] results = self._yolo( frame, conf=settings.DETECTION_CONFIDENCE_THRESHOLD, verbose=False, ) detections = [] h, w = frame.shape[:2] for result in results: boxes = result.boxes if boxes is None: continue # Group: find all vehicles and all persons vehicles = [] persons = [] for box in boxes: cls_id = int(box.cls[0]) conf = float(box.conf[0]) x1, y1, x2, y2 = map(int, box.xyxy[0]) if cls_id in VEHICLE_CLASSES: vehicles.append({ "type": VEHICLE_CLASSES[cls_id], "conf": conf, "bbox": {"x": x1, "y": y1, "w": x2 - x1, "h": y2 - y1}, "crop": frame[y1:y2, x1:x2], }) elif cls_id == PERSON_CLASS: persons.append({"bbox": {"x": x1, "y": y1, "w": x2 - x1, "h": y2 - y1}}) for vehicle in vehicles: det = VehicleDetection( vehicle_type=vehicle["type"], detection_confidence=vehicle["conf"], vehicle_bbox=vehicle["bbox"], vehicle_crop=vehicle["crop"], ) # Detect color det.color = self._detect_color(vehicle["crop"]) # Count persons on/near motorcycle if vehicle["type"] == "motorcycle": det.person_count = self._count_persons_on_motorcycle( vehicle["bbox"], persons ) # Read plate if self._alpr is not None: plate_result = self._read_plate(frame, vehicle["bbox"]) if plate_result: det.plate_number = plate_result["text"] det.plate_confidence = plate_result["confidence"] det.plate_bbox = plate_result["bbox"] bx = plate_result["bbox"] det.plate_crop = frame[ bx["y"]:bx["y"] + bx["h"], bx["x"]:bx["x"] + bx["w"] ] detections.append(det) return detections def _read_plate(self, frame: np.ndarray, vehicle_bbox: dict) -> Optional[dict]: """Run ALPR on vehicle region.""" if self._alpr is None: return None try: # Expand vehicle crop slightly for plate detection x, y, w, h = vehicle_bbox["x"], vehicle_bbox["y"], vehicle_bbox["w"], vehicle_bbox["h"] crop = frame[max(0, y):y + h, max(0, x):x + w] results = self._alpr.run(crop) if results: best = results[0] ocr = best.ocr if ocr and ocr.text: # Offset bbox back to full frame coordinates pb = best.detection.bounding_box return { "text": ocr.text.upper().replace(" ", ""), "confidence": float(ocr.confidence), "bbox": { "x": x + int(pb.x1), "y": y + int(pb.y1), "w": int(pb.x2 - pb.x1), "h": int(pb.y2 - pb.y1), } } except Exception as e: logger.debug(f"ALPR error: {e}") return None def _detect_color(self, crop: np.ndarray) -> str: """Detect dominant vehicle color using HSV histogram.""" if crop is None or crop.size == 0: return "unknown" try: hsv = cv2.cvtColor(crop, cv2.COLOR_BGR2HSV) max_pixels = 0 detected_color = "unknown" for color_name, (lower, upper) in COLOR_RANGES.items(): mask = cv2.inRange(hsv, np.array(lower), np.array(upper)) pixel_count = cv2.countNonZero(mask) if pixel_count > max_pixels: max_pixels = pixel_count detected_color = color_name.replace("2", "") # red2 → red return detected_color except Exception: return "unknown" def _count_persons_on_motorcycle(self, moto_bbox: dict, persons: list) -> int: """Count persons whose center falls within or near the motorcycle bbox.""" count = 0 mx, my, mw, mh = moto_bbox["x"], moto_bbox["y"], moto_bbox["w"], moto_bbox["h"] for person in persons: px, py, pw, ph = person["bbox"]["x"], person["bbox"]["y"], person["bbox"]["w"], person["bbox"]["h"] # Person center cx = px + pw // 2 cy = py + ph // 2 # Check if center is within motorcycle bounding box (with some margin) margin = 30 if (mx - margin <= cx <= mx + mw + margin and my - margin <= cy <= my + mh + margin): count += 1 return max(count, 1) # At least 1 rider assumed # Singleton vehicle_detector = VehicleDetector()