detector.py

"""
detector.py
YOLOv8 person detection wrapper.
Returns raw lists so other modules stay decoupled from supervision version.
"""
from __future__ import annotations
import numpy as np


class PersonDetector:
    """
    Wraps YOLOv8. Returns detections as a plain dict so the rest of the
    pipeline never touches supervision directly from here.
    """

    def __init__(
        self,
        model_path: str = "yolov8n.pt",
        conf: float = 0.30,
        iou: float  = 0.50,
        device: str = "cpu",
    ) -> None:
        from ultralytics import YOLO
        print(f"[Detector] loading {model_path} on {device}")
        self.model  = YOLO(model_path)
        self.conf   = conf
        self.iou    = iou
        self.device = device

    def detect(self, frame: np.ndarray) -> list[dict]:
        """
        Run YOLOv8 on one BGR frame.

        Returns:
            list of {"xyxy": [x1,y1,x2,y2], "conf": float}
        """
        results = self.model(
            frame,
            conf=0.1,         # Pass low conf detections to ByteTrack
            iou=self.iou,
            classes=[0],      # person only
            imgsz=480,        # Faster inference
            verbose=False,
            device=self.device,
        )[0]

        out = []
        boxes = results.boxes
        if boxes is None or len(boxes) == 0:
            return out

        for box in boxes:
            xyxy = box.xyxy[0].cpu().numpy().tolist()   # [x1,y1,x2,y2]
            conf = float(box.conf[0].cpu().numpy())
            out.append({"xyxy": xyxy, "conf": conf})

        return out