bantaycam/backend/app/services/vehicle_detector.py

"""
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()