core/detector.py

"""YOLOv8 ONNX Runtime inference wrapper.

Loads the exported ONNX model and runs inference on BGR numpy frames.
Uses CUDAExecutionProvider when available, falls back to CPU.
"""
from __future__ import annotations

import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Sequence

import cv2
import numpy as np

try:
    import onnxruntime as ort
except ImportError as e:
    raise ImportError("onnxruntime-gpu is required: pip install onnxruntime-gpu") from e


CLASS_NAMES: list[str] = ["car", "bus", "motorcycle", "truck"]

# Default ONNX model path — override via MODEL_PATH env var
DEFAULT_MODEL_PATH = Path(__file__).parent.parent / "models" / "weights" / "yolov8s_traffic.onnx"


@dataclass
class Detection:
    """Single bounding box detection from one frame."""
    bbox: tuple[float, float, float, float]  # x1, y1, x2, y2 (pixel coords)
    confidence: float
    class_id: int
    class_name: str = field(init=False)

    def __post_init__(self) -> None:
        self.class_name = CLASS_NAMES[self.class_id] if self.class_id < len(CLASS_NAMES) else "unknown"

    @property
    def xyxy(self) -> tuple[float, float, float, float]:
        return self.bbox

    @property
    def xywh(self) -> tuple[float, float, float, float]:
        x1, y1, x2, y2 = self.bbox
        return x1, y1, x2 - x1, y2 - y1

    @property
    def area(self) -> float:
        x1, y1, x2, y2 = self.bbox
        return max(0.0, x2 - x1) * max(0.0, y2 - y1)


class Detector:
    """ONNX Runtime inference wrapper for YOLOv8.

    Usage:
        detector = Detector()                     # loads from env or default path
        detector = Detector("models/weights/x.onnx")
        detections = detector.detect(bgr_frame)
    """

    def __init__(
        self,
        model_path: str | Path | None = None,
        *,
        confidence_threshold: float = 0.35,
        iou_threshold: float = 0.45,
        input_size: int = 640,
    ) -> None:
        path = Path(
            model_path
            or os.environ.get("MODEL_PATH", "")
            or DEFAULT_MODEL_PATH
        )
        if not path.exists():
            raise FileNotFoundError(
                f"ONNX model not found: {path}\n"
                "Run Phase 2 notebook to export the model, then copy it to models/weights/."
            )

        self.model_path = path
        self.confidence_threshold = confidence_threshold
        self.iou_threshold = iou_threshold
        self.input_size = input_size

        self._session = self._load_session(path)
        self._input_name: str = self._session.get_inputs()[0].name

    # ── Setup ─────────────────────────────────────────────────────────────────

    def _load_session(self, path: Path) -> ort.InferenceSession:
        providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
        sess_opts = ort.SessionOptions()
        sess_opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        session = ort.InferenceSession(str(path), sess_options=sess_opts, providers=providers)
        active = session.get_providers()
        using_gpu = "CUDAExecutionProvider" in active
        print(f"[Detector] Loaded {path.name} | {'GPU' if using_gpu else 'CPU'} | providers={active}")
        return session

    # ── Inference ─────────────────────────────────────────────────────────────

    def detect(self, frame: np.ndarray) -> list[Detection]:
        """Run detection on a single BGR frame. Returns list of Detection objects."""
        orig_h, orig_w = frame.shape[:2]
        blob, scale, pad = self._preprocess(frame)
        raw_output = self._session.run(None, {self._input_name: blob})[0]
        detections = self._postprocess(raw_output, orig_w, orig_h, scale, pad)
        return detections

    # ── Pre/post processing ───────────────────────────────────────────────────

    def _preprocess(
        self, frame: np.ndarray
    ) -> tuple[np.ndarray, float, tuple[int, int]]:
        """Letterbox resize → CHW float32 blob. Returns (blob, scale, (pad_w, pad_h))."""
        s = self.input_size
        h, w = frame.shape[:2]

        # Scale keeping aspect ratio
        scale = min(s / w, s / h)
        new_w = int(round(w * scale))
        new_h = int(round(h * scale))

        resized = cv2.resize(frame, (new_w, new_h), interpolation=cv2.INTER_LINEAR)

        # Pad to square
        pad_w = (s - new_w) // 2
        pad_h = (s - new_h) // 2
        padded = cv2.copyMakeBorder(
            resized, pad_h, s - new_h - pad_h, pad_w, s - new_w - pad_w,
            cv2.BORDER_CONSTANT, value=(114, 114, 114),
        )

        # BGR → RGB, HWC → CHW, [0,255] → [0,1]
        rgb = cv2.cvtColor(padded, cv2.COLOR_BGR2RGB)
        blob = (rgb.transpose(2, 0, 1).astype(np.float32) / 255.0)[np.newaxis]
        return blob, scale, (pad_w, pad_h)

    def _postprocess(
        self,
        output: np.ndarray,
        orig_w: int,
        orig_h: int,
        scale: float,
        pad: tuple[int, int],
    ) -> list[Detection]:
        """Parse YOLOv8 output tensor → Detection list with NMS applied.

        YOLOv8 ONNX output shape: (1, 8, num_anchors)
        where 8 = [cx, cy, w, h, cls0_score, cls1_score, cls2_score, cls3_score]
        """
        # output: (1, num_attrs, num_anchors) → transpose to (num_anchors, num_attrs)
        preds = output[0].T  # (num_anchors, 8)

        pad_w, pad_h = pad
        detections: list[Detection] = []
        boxes_for_nms: list[list[float]] = []
        scores_for_nms: list[float] = []
        class_ids: list[int] = []

        for pred in preds:
            cx, cy, w, h = pred[:4]
            class_scores = pred[4:]
            class_id = int(np.argmax(class_scores))
            confidence = float(class_scores[class_id])

            if confidence < self.confidence_threshold:
                continue

            # Convert from padded/scaled coords back to original frame coords
            x1 = (cx - w / 2 - pad_w) / scale
            y1 = (cy - h / 2 - pad_h) / scale
            x2 = (cx + w / 2 - pad_w) / scale
            y2 = (cy + h / 2 - pad_h) / scale

            # Clamp to frame bounds
            x1 = max(0.0, min(float(orig_w), x1))
            y1 = max(0.0, min(float(orig_h), y1))
            x2 = max(0.0, min(float(orig_w), x2))
            y2 = max(0.0, min(float(orig_h), y2))

            boxes_for_nms.append([x1, y1, x2 - x1, y2 - y1])  # xywh for cv2.dnn.NMSBoxes
            scores_for_nms.append(confidence)
            class_ids.append(class_id)

        if not boxes_for_nms:
            return []

        # NMS
        indices = cv2.dnn.NMSBoxes(
            boxes_for_nms, scores_for_nms, self.confidence_threshold, self.iou_threshold
        )
        if indices is None or len(indices) == 0:
            return []

        for i in indices.flatten():
            x1, y1, w, h = boxes_for_nms[i]
            detections.append(
                Detection(
                    bbox=(x1, y1, x1 + w, y1 + h),
                    confidence=scores_for_nms[i],
                    class_id=class_ids[i],
                )
            )

        return detections

    def __repr__(self) -> str:
        return (
            f"Detector(model={self.model_path.name}, "
            f"conf={self.confidence_threshold}, iou={self.iou_threshold})"
        )