detector.py

"""
Person & Car detector via YOLOv8 (ultralytics).
Falls back to density-map peak detection if YOLO not available.
"""
import os
import cv2
import numpy as np

MODEL_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models")
MODEL_PATH = os.path.join(MODEL_DIR, "yolov8s.pt")

# COCO class IDs we care about
PERSON_ID = 0
CAR_ID = 2
TARGET_IDS = {PERSON_ID, CAR_ID}

CONF_THRESHOLD = 0.15
IOU_THRESHOLD = 0.35
MODEL_IMGSZ = 1280  # larger input = better for small/distant objects


class PersonCarDetector:
    """Detect persons and cars using YOLOv8 nano."""

    def __init__(self):
        self._model = None
        self._available = False
        self._init_model()

    def _init_model(self):
        try:
            from ultralytics import YOLO

            os.makedirs(MODEL_DIR, exist_ok=True)

            if not os.path.exists(MODEL_PATH):
                print(f"  [detector] Model not found: {MODEL_PATH}")
                print("  [detector] Run download_models.py first")
                self._available = False
                return

            self._model = YOLO(MODEL_PATH)
            self._available = True
            print("  [detector] YOLOv8s person/car detector loaded")
        except Exception as e:
            print(f"  [detector] YOLO init failed: {e}")
            print("  [detector] Falling back to density-map peak detection")
            self._available = False

    @property
    def available(self) -> bool:
        return self._available

    def detect(self, image: np.ndarray) -> list:
        """
        Detect persons and cars.
        Returns list of dicts: {bbox: (x,y,w,h), class_name: str, confidence: float}
        """
        if not self._available:
            return []

        results = self._model(image, imgsz=MODEL_IMGSZ, conf=CONF_THRESHOLD,
                              iou=IOU_THRESHOLD, verbose=False)
        detections = []

        for r in results:
            boxes = r.boxes
            if boxes is None:
                continue
            for box in boxes:
                cls_id = int(box.cls[0])
                if cls_id not in TARGET_IDS:
                    continue

                conf = float(box.conf[0])
                if conf < CONF_THRESHOLD:
                    continue

                xyxy = box.xyxy[0].cpu().numpy()
                x1, y1, x2, y2 = xyxy
                x, y = int(x1), int(y1)
                w, h = int(x2 - x1), int(y2 - y1)

                class_name = "person" if cls_id == PERSON_ID else "car"

                detections.append({
                    "bbox": (x, y, w, h),
                    "class_name": class_name,
                    "confidence": round(conf, 3),
                })

        return detections


def draw_detections(image: np.ndarray, detections: list) -> np.ndarray:
    """Draw bounding boxes and labels on image. Returns annotated copy."""
    vis = image.copy()
    colors = {"person": (0, 255, 0), "car": (255, 165, 0)}
    for d in detections:
        x, y, w, h = d["bbox"]
        cls = d.get("class_name", "person")
        conf = d.get("confidence", 0)
        color = colors.get(cls, (0, 255, 255))

        cv2.rectangle(vis, (x, y), (x + w, y + h), color, 2)
        label = f"{cls} {conf:.2f}"
        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
        cv2.rectangle(vis, (x, y - th - 6), (x + tw + 4, y), color, -1)
        cv2.putText(vis, label, (x + 2, y - 4), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
    return vis


def find_density_peaks(density_map: np.ndarray, original_shape: tuple,
                       threshold_pct: float = 20) -> list:
    """
    Find person locations from density map peaks.
    Only returns peaks with significant density (filters noise in empty areas).
    """
    h, w = original_shape[:2]
    dh, dw = density_map.shape[:2]

    if density_map.max() <= density_map.min():
        return []

    d_norm = (density_map - density_map.min()) / (density_map.max() - density_map.min() + 1e-8)
    d_u8 = (d_norm * 255).astype(np.uint8)

    # Keep only the top threshold_pct% brightest pixels
    thresh_val = max(int(np.percentile(d_u8, 100 - threshold_pct)), 40)
    _, binary = cv2.threshold(d_u8, thresh_val, 255, cv2.THRESH_BINARY)

    contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    scale_y = h / dh
    scale_x = w / dw

    # Dynamic filtering: only regions well above the mean density
    global_mean = d_norm.mean()

    results = []
    for cnt in contours:
        area = cv2.contourArea(cnt)
        if area < 10:  # filter tiny noise blobs
            continue

        rx, ry, rw, rh = cv2.boundingRect(cnt)

        mask = np.zeros_like(d_norm, dtype=np.uint8)
        cv2.drawContours(mask, [cnt], -1, 255, -1)
        region_vals = d_norm[mask > 0]
        if len(region_vals) == 0:
            continue

        peak_val = float(region_vals.max())
        # Must be significantly above the mean density
        if peak_val < global_mean * 3 or peak_val < 0.3:
            continue

        x = int(rx * scale_x)
        y = int(ry * scale_y)
        bw = int(rw * scale_x)
        bh = int(rh * scale_y)

        results.append({
            "bbox": (max(0, x), max(0, y), bw, bh),
            "class_name": "person",
            "confidence": round(min(peak_val, 0.6), 3),
        })

    return results