src/detection/postprocess.py

"""
Detection post-processing — decode YOLOv8 output tensor and apply NMS.

YOLOv8 ONNX output shape: ``[1, 84, num_anchors]``  (transposed from v5/v7)
  * Row 0-3 : cx, cy, w, h  (normalised to input_size)
  * Row 4-83: class scores   (no objectness score in v8)
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Optional

import cv2
import numpy as np

from src.detection.class_labels import get_color, get_label


@dataclass
class Detection:
    """A single detected object."""
    bbox: tuple[int, int, int, int]   # x1, y1, x2, y2  (pixel coords in original frame)
    conf: float                        # detection confidence
    class_id: int
    label: str = field(init=False)
    color: tuple[int, int, int] = field(init=False)

    def __post_init__(self):
        self.label = get_label(self.class_id)
        self.color = get_color(self.class_id)

    @property
    def x1(self) -> int: return self.bbox[0]
    @property
    def y1(self) -> int: return self.bbox[1]
    @property
    def x2(self) -> int: return self.bbox[2]
    @property
    def y2(self) -> int: return self.bbox[3]

    @property
    def width(self) -> int: return self.x2 - self.x1
    @property
    def height(self) -> int: return self.y2 - self.y1

    @property
    def center(self) -> tuple[int, int]:
        return (self.x1 + self.x2) // 2, (self.y1 + self.y2) // 2

    @property
    def tlwh(self) -> tuple[int, int, int, int]:
        """top-left x, top-left y, width, height"""
        return self.x1, self.y1, self.width, self.height

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

    @property
    def area(self) -> int:
        return max(0, self.width) * max(0, self.height)

    def iou(self, other: "Detection") -> float:
        """IoU with another Detection."""
        ix1 = max(self.x1, other.x1)
        iy1 = max(self.y1, other.y1)
        ix2 = min(self.x2, other.x2)
        iy2 = min(self.y2, other.y2)
        inter = max(0, ix2 - ix1) * max(0, iy2 - iy1)
        union = self.area + other.area - inter
        return inter / (union + 1e-6)

    def crop(self, frame: np.ndarray, pad: float = 0.05) -> np.ndarray:
        """Crop the ROI from *frame* with optional percentage padding."""
        h, w = frame.shape[:2]
        pw = int(self.width * pad)
        ph = int(self.height * pad)
        x1 = max(0, self.x1 - pw)
        y1 = max(0, self.y1 - ph)
        x2 = min(w, self.x2 + pw)
        y2 = min(h, self.y2 + ph)
        return frame[y1:y2, x1:x2].copy()

    def __repr__(self) -> str:
        return (
            f"Detection({self.label!r}, conf={self.conf:.2f}, "
            f"bbox={self.bbox})"
        )


# ── Postprocessing ────────────────────────────────────────────────────────────

def decode_yolov8(
    output: np.ndarray,
    orig_shape: tuple[int, int],
    input_size: int = 640,
    conf_threshold: float = 0.40,
    nms_threshold: float = 0.45,
    class_filter: Optional[list[int]] = None,
) -> list[Detection]:
    """
    Decode a raw YOLOv8 ONNX output into ``Detection`` objects.

    Args:
        output:        Raw net.forward() output, shape ``[1, 84, N]`` or ``[84, N]``.
        orig_shape:    Original frame shape ``(height, width)``.
        input_size:    Square input dimension used during inference.
        conf_threshold: Minimum confidence to keep a detection.
        nms_threshold: IoU threshold for NMS.
        class_filter:  If non-empty, only keep detections of listed class IDs.

    Returns:
        List of ``Detection`` objects in original image coordinates.
    """
    # Ensure shape is [84, N]
    pred = output
    if pred.ndim == 3:
        pred = pred[0]                # [84, N]
    if pred.shape[0] != 84:
        pred = pred.T                 # some exports give [N, 84]

    orig_h, orig_w = orig_shape
    scale_x = orig_w / input_size
    scale_y = orig_h / input_size

    boxes_xywh = pred[:4].T          # [N, 4]  cx,cy,w,h (normalised to input_size)
    class_scores = pred[4:].T        # [N, 80]

    class_ids = np.argmax(class_scores, axis=1)
    confs = class_scores[np.arange(len(class_ids)), class_ids]

    # Confidence threshold
    mask = confs >= conf_threshold
    if class_filter:
        mask &= np.isin(class_ids, class_filter)

    boxes_xywh = boxes_xywh[mask]
    confs = confs[mask]
    class_ids = class_ids[mask]

    if len(boxes_xywh) == 0:
        return []

    # Convert cx,cy,w,h → x1,y1,w,h for cv2.dnn.NMSBoxes
    cx, cy, bw, bh = boxes_xywh[:, 0], boxes_xywh[:, 1], boxes_xywh[:, 2], boxes_xywh[:, 3]
    x1 = (cx - bw / 2) * scale_x
    y1 = (cy - bh / 2) * scale_y
    w_px = bw * scale_x
    h_px = bh * scale_y

    # Clip to frame
    x1 = np.clip(x1, 0, orig_w).astype(int)
    y1 = np.clip(y1, 0, orig_h).astype(int)
    w_px = np.clip(w_px, 1, orig_w - x1).astype(int)
    h_px = np.clip(h_px, 1, orig_h - y1).astype(int)

    # NMS per class
    boxes_list = [[int(x), int(y), int(w), int(h)] for x, y, w, h in zip(x1, y1, w_px, h_px)]
    indices = cv2.dnn.NMSBoxes(
        boxes_list,
        confs.tolist(),
        conf_threshold,
        nms_threshold,
    )
    if isinstance(indices, np.ndarray):
        indices = indices.flatten().tolist()

    detections: list[Detection] = []
    for i in indices:
        x, y, w, h = boxes_list[i]
        detections.append(
            Detection(
                bbox=(x, y, x + w, y + h),
                conf=float(confs[i]),
                class_id=int(class_ids[i]),
            )
        )
    return detections