miner.py

# v8: yolo11s trained on validator-aligned SAM3 labels.
# Pool val backtest: F1=0.862 vs v32 F1=0.742 (+0.125 absolute, smoke recall 89% vs 43%).
from pathlib import Path

import cv2
import numpy as np
import onnxruntime as ort
from numpy import ndarray
from pydantic import BaseModel


class BoundingBox(BaseModel):
    x1: int
    y1: int
    x2: int
    y2: int
    cls_id: int
    conf: float


class TVFrameResult(BaseModel):
    frame_id: int
    boxes: list[BoundingBox]
    keypoints: list[tuple[int, int]]


class Miner:
    """v8: ONNX with built-in NMS → light post-processing.

    Pipeline:
      1. Letterbox to 1280x1280
      2. ONNX inference (returns [1, 300, 6] post-NMS)
      3. Conf filter
      4. Extra per-class dedup (IoU > 0.3 OR ≥80% containment) — catches nested
         duplicates the model inherits from SAM3 training labels that default
         NMS@0.5 doesn't suppress
      5. Top-1 fallback if everything got filtered — empty predictions are heavily
         penalized; even a low-conf best guess scores better than nothing
      6. Un-letterbox coords back to original size + clip
    """

    # validator-visible class output order (what the runner expects in cls_id)
    class_names = ["fire", "smoke", "fire extinguisher"]
    # order the v8 ONNX emits classes (training CLASS_ORDER in v8_build_dataset.py)
    _model_class_order = ["fire", "fire extinguisher", "smoke"]

    input_size = 1280
    conf_thresh = 0.25
    nms_iou_thresh = 0.3
    contain_thresh = 0.80
    fallback_min_conf = 0.05  # top-1 fallback floors at this; never return total junk

    def __init__(self, path_hf_repo: Path) -> None:
        model_path = path_hf_repo / "weights.onnx"

        self.cls_remap = np.array(
            [self.class_names.index(n) for n in self._model_class_order],
            dtype=np.int32,
        )

        try:
            ort.preload_dlls()
        except Exception as e:
            print(f"preload_dlls: {e}")

        sess_options = ort.SessionOptions()
        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

        try:
            self.session = ort.InferenceSession(
                str(model_path),
                sess_options=sess_options,
                providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
            )
        except Exception as e:
            print(f"CUDA failed, CPU: {e}")
            self.session = ort.InferenceSession(
                str(model_path),
                sess_options=sess_options,
                providers=["CPUExecutionProvider"],
            )

        self.input_name = self.session.get_inputs()[0].name
        print(f"v8 ONNX loaded, providers={self.session.get_providers()}")

    def __repr__(self) -> str:
        return f"v8 Miner (providers={self.session.get_providers()})"

    def _letterbox(self, image: ndarray):
        h, w = image.shape[:2]
        s = self.input_size / max(h, w)
        nw, nh = int(round(w * s)), int(round(h * s))
        if (nw, nh) != (w, h):
            interp = cv2.INTER_CUBIC if s > 1.0 else cv2.INTER_LINEAR
            image = cv2.resize(image, (nw, nh), interpolation=interp)
        canvas = np.full((self.input_size, self.input_size, 3), 114, dtype=np.uint8)
        dx = (self.input_size - nw) // 2
        dy = (self.input_size - nh) // 2
        canvas[dy:dy + nh, dx:dx + nw] = image
        return canvas, s, (dx, dy)

    def _preprocess(self, image: ndarray):
        H, W = image.shape[:2]
        padded, scale, (dx, dy) = self._letterbox(image)
        x = padded[:, :, ::-1].astype(np.float32) / 255.0  # BGR→RGB, /255
        x = np.ascontiguousarray(x.transpose(2, 0, 1)[None], dtype=np.float32)
        return x, scale, (dx, dy), (W, H)

    @staticmethod
    def _iou(a, b):
        ix1 = max(a[0], b[0]); iy1 = max(a[1], b[1])
        ix2 = min(a[2], b[2]); iy2 = min(a[3], b[3])
        iw = max(0.0, ix2 - ix1); ih = max(0.0, iy2 - iy1)
        inter = iw * ih
        ua = (a[2]-a[0])*(a[3]-a[1]) + (b[2]-b[0])*(b[3]-b[1]) - inter
        return inter / ua if ua > 0 else 0.0

    @staticmethod
    def _containment(inner, outer):
        ix1 = max(inner[0], outer[0]); iy1 = max(inner[1], outer[1])
        ix2 = min(inner[2], outer[2]); iy2 = min(inner[3], outer[3])
        iw = max(0.0, ix2 - ix1); ih = max(0.0, iy2 - iy1)
        inter = iw * ih
        a_in = (inner[2]-inner[0]) * (inner[3]-inner[1])
        return inter / a_in if a_in > 0 else 0.0

    def _dedup(self, boxes_xyxy, scores, cls_ids):
        """Per-class dedup: drop a box if same-class IoU>nms_iou OR ≥contain_thresh
        contained in a larger same-class box. Keep larger box on ties."""
        n = len(boxes_xyxy)
        if n <= 1:
            return np.arange(n, dtype=np.intp)
        # Sort by area desc (so we always test smaller against larger we already kept)
        areas = (boxes_xyxy[:, 2] - boxes_xyxy[:, 0]) * (boxes_xyxy[:, 3] - boxes_xyxy[:, 1])
        order = np.argsort(-areas)
        keep = []
        suppressed = np.zeros(n, dtype=bool)
        for i in order:
            if suppressed[i]: continue
            keep.append(int(i))
            for j in order:
                if j == i or suppressed[j]: continue
                if cls_ids[i] != cls_ids[j]: continue
                if self._iou(boxes_xyxy[i], boxes_xyxy[j]) > self.nms_iou_thresh:
                    suppressed[j] = True; continue
                if self._containment(boxes_xyxy[j], boxes_xyxy[i]) >= self.contain_thresh:
                    suppressed[j] = True
        keep.sort()
        return np.array(keep, dtype=np.intp)

    def _predict_one(self, frame: ndarray) -> list[BoundingBox]:
        x, scale, (dx, dy), (W, H) = self._preprocess(frame)
        out = self.session.run(None, {self.input_name: x})[0]
        # output shape: [1, 300, 6] — (x1, y1, x2, y2, conf, cls_id)
        raw = out[0]
        if raw.shape[0] == 0:
            return []

        # Apply conf filter (keep raw for fallback)
        primary = raw[raw[:, 4] >= self.conf_thresh]

        # Per-class dedup on the conf-filtered set
        final_dets = []
        if len(primary) > 0:
            xyxy = primary[:, :4].astype(np.float32)
            scores = primary[:, 4].astype(np.float32)
            cls_ids = primary[:, 5].astype(np.int32)
            keep_idx = self._dedup(xyxy, scores, cls_ids)
            primary = primary[keep_idx]
            for det in primary:
                final_dets.append(det)

        # Fallback: nothing left → return single highest-conf raw box (above floor)
        if not final_dets and raw.shape[0] > 0:
            top = raw[np.argmax(raw[:, 4])]
            if top[4] >= self.fallback_min_conf:
                final_dets.append(top)

        # Build BoundingBox list with un-letterbox + cls remap
        boxes_out: list[BoundingBox] = []
        for det in final_dets:
            x1, y1, x2, y2, conf, model_cls_id = det
            x1 = (x1 - dx) / scale; x2 = (x2 - dx) / scale
            y1 = (y1 - dy) / scale; y2 = (y2 - dy) / scale
            x1 = max(0.0, min(W - 1.0, x1)); x2 = max(0.0, min(W - 1.0, x2))
            y1 = max(0.0, min(H - 1.0, y1)); y2 = max(0.0, min(H - 1.0, y2))
            if x2 <= x1 or y2 <= y1:
                continue
            mapped_cls = int(self.cls_remap[int(model_cls_id)])
            boxes_out.append(BoundingBox(
                x1=int(x1), y1=int(y1), x2=int(x2), y2=int(y2),
                cls_id=mapped_cls, conf=float(conf),
            ))
        return boxes_out

    def predict_batch(
        self,
        batch_images: list[ndarray],
        offset: int,
        n_keypoints: int,
    ) -> list[TVFrameResult]:
        """Required interface for chute template (sv_chutes_*.py)."""
        results: list[TVFrameResult] = []
        for frame_number_in_batch, image in enumerate(batch_images):
            try:
                boxes = self._predict_one(image)
            except Exception as e:
                print(f"⚠️ Inference failed for frame "
                      f"{offset + frame_number_in_batch}: {e}")
                boxes = []
            results.append(TVFrameResult(
                frame_id=offset + frame_number_in_batch,
                boxes=boxes,
                keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))],
            ))
        return results

    # Back-compat alias for local sanity testing
    def run(self, frames: list[ndarray]) -> list[TVFrameResult]:
        return self.predict_batch(frames, offset=0, n_keypoints=0)