bus_01 / miner.py
scorevision: push artifact
1209a21 verified about 2 months ago
35.8 kB
import os
import math
from pathlib import Path

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


# =============================================================================
# Vehicle miner (manak0/Detect-Vehicle) — port of the proven carwash001 miner.
#
# Model: yolo26n end2end ONNX @704, trained on vehicle2-2-powerful with the
# name-based class remap, so it emits ELEMENT-CANONICAL indices natively:
#     0=car, 1=bus, 2=truck, 3=motorcycle, 4=van
# (cls_remap is read from ONNX metadata anyway and is the identity here).
#
# Carried over from carwash001 (all validated on the carwash element):
#   * robust model loading, FP16 auto-detect, metadata class order, ORT thread
#     pinning, raw-YOLO decode fallback, NO CLAHE (train/test match)
#   * vectorized hard NMS + cluster-max boost, per-class NMS IoU / min-area,
#     pre_nms_topk, cross-class dedup, square-domain hook
#   * soft-NMS with post-decay re-filter (use_soft_nms) — measured +0.0003..10
#     composite on carwash; RE-SWEEP for vehicle before trusting it
#
# NEW, vehicle-specific lever (OFF by default):
#   * dual-label bus/truck hedging (dual_label_below) — the rival config-driven
#     miners (roadsign1007 / tv-vehicle-cea7014a) duplicate any bus or truck box
#     below conf 0.6 under the sibling class, hedging the bus<->truck confusion
#     against the validator's per-class greedy matcher. The duplicate is a
#     guaranteed FP (-0.1 fp on that frame) whenever the original class was
#     right, so this is NOT a free win — enable only if tune_miner proves it.
#
# Thresholds below are TUNED (tune_miner.py, 2026-09-04, adaptive-IoU matcher
# on vehicle2-2-powerful/val: composite 0.8450 -> 0.8613). Re-run the sweep
# after ANY retrain/re-export. A failing conformity tuple (latency/format)
# removes the miner from scoring entirely — bench latency on 2 pinned CPUs too
# (2026-09-04: mean 55ms, p95 60ms -> PASS).
# =============================================================================


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:
    def __init__(self, path_hf_repo: Path) -> None:
        model_path = self._resolve_model_path(path_hf_repo)
        # Canonical class order every downstream consumer (validator,
        # BoundingBox.cls_id) sees — MUST equal element.objects order:
        self.class_names = ["car", "bus", "truck", "motorcycle"]
        # Fallback model-emit order (used only when ONNX metadata is missing).
        self._model_class_order = ["car", "bus", "truck", "motorcycle"]
        print("ORT version:", ort.__version__)

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

        print("ORT available providers BEFORE session:", ort.get_available_providers())

        sess_options = ort.SessionOptions()
        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        # Pin threads for the CPU latency gate.
        sess_options.intra_op_num_threads = 2
        sess_options.inter_op_num_threads = 1
        sess_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL

        # MINER_FORCE_CPU=1 skips the CUDA attempt — used by bench_latency.py to
        # measure the conformity-gate path (this box has onnxruntime-gpu too, so
        # an unforced "CPU" benchmark silently runs on the GPU).
        force_cpu = os.environ.get("MINER_FORCE_CPU", "") == "1"
        try:
            if force_cpu:
                raise RuntimeError("MINER_FORCE_CPU=1")
            self.session = ort.InferenceSession(
                str(model_path),
                sess_options=sess_options,
                providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
            )
            print("Created ORT session with preferred CUDA provider list")
        except Exception as e:
            print(f"CUDA session creation skipped/failed, using CPU: {e}")
            self.session = ort.InferenceSession(
                str(model_path),
                sess_options=sess_options,
                providers=["CPUExecutionProvider"],
            )
        print("ORT session providers:", self.session.get_providers())

        # cls_remap[i] = self.class_names.index(model_class_order[i]); order comes
        # from ONNX metadata when present, else the static fallback.
        model_class_order = self._read_model_class_order()
        if model_class_order is None:
            model_class_order = list(self._model_class_order)
            print(f"cls order: no usable ONNX metadata, FALLBACK {model_class_order}")
        else:
            print(f"cls order: from ONNX metadata {model_class_order}")
        self.cls_remap = np.array(
            [self.class_names.index(n) for n in model_class_order], dtype=np.int32
        )

        for inp in self.session.get_inputs():
            print("INPUT:", inp.name, inp.shape, inp.type)
        for out in self.session.get_outputs():
            print("OUTPUT:", out.name, out.shape, out.type)

        self.input_name = self.session.get_inputs()[0].name
        self.output_names = [o.name for o in self.session.get_outputs()]
        self.input_shape = self.session.get_inputs()[0].shape

        # Match the ONNX input dtype (FP16 export needs float16 input).
        input_type = self.session.get_inputs()[0].type
        self.np_dtype = np.float16 if "float16" in input_type else np.float32
        print(f"ONNX input dtype: {input_type} -> numpy {self.np_dtype}")

        # The miner MUST run at the ONNX's baked input size (704). Dropping to
        # 640/576 is a latency lever but an EXPORT decision (re-export +
        # re-eval), not a miner-side change.
        self.input_height = self._safe_dim(self.input_shape[2], default=704)
        self.input_width = self._safe_dim(self.input_shape[3], default=704)

        # ── Post-processing config (tune on vehicle2-2-powerful/val) ──────────
        # Per-class NMS IoU (car, bus, truck, motorcycle, van).
        self._iou_thres_array = np.array([0.55] * 4, dtype=np.float32)
        self.cross_iou_thresh = 0.9   # cross-class dedup IoU (same vehicle firing
                                      # 2 classes, e.g. van->car+van, bus->bus+truck)
        self.max_det = 300            # vehicle frames can be dense (car parks)
        self.pre_nms_topk = 1000      # cap candidates before NMS (crowded-frame speed guard)
        self.use_tta = False          # 2nd flipped pass; doubles latency -> off for CPU gate

        # Single-view cluster boost: raise each survivor's conf to its same-class
        # IoU-cluster max. OFF per the 2026-09-04 vehicle sweep (carwash had it
        # on -- element-specific, keep sweeping on new weights).
        self.use_cluster_boost = True

        # Soft-NMS: Gaussian decay of overlapping same-class scores instead of
        # hard deletion, then RE-FILTER decayed scores against the per-class
        # confidence thresholds (the re-filter is what keeps decayed junk away
        # from the FP pillar). Carwash-validated, but the 2026-09-04 vehicle
        # sweep preferred hard NMS (off: 0.8613 vs on: 0.8608) -- element- and
        # weights-specific, re-sweep after every retrain.
        self.use_soft_nms = False
        self.soft_nms_sigma = 0.5

        # Per-class confidence thresholds (car, bus, truck, motorcycle, van).
        # `_conf_thres_array` is the SQUARE / validator-eval set (every scored
        # frame is 1024x1024 squished, so this is the one that matters);
        # `_extra` covers non-square inputs (warmup / any 16:9 frame).
        # Tuned 2026-09-04 (tune_miner.py, 209-img vehicle2-2-powerful/val,
        # adaptive-IoU matcher): comp 0.8450 -> 0.8613, map50 0.8056 fp 0.9450.
        # NOTE the van value is a PLACEHOLDER: train/val have ZERO van labels
        # and the van head emits nothing (>0.01) -- retune once van data lands.
        self._conf_thres_array = np.array([0.3, 0.1, 0.4, 0.17], dtype=np.float32)
        self._extra_conf_thres_array = np.array([0.25, 0.1, 0.4, 0.17], dtype=np.float32)

        # Per-class rescue bonus: if a class has ZERO boxes passing, admit its
        # top-1 when score >= (threshold - bonus). ON per the 2026-09-04 sweep
        # (carwash had it off -- element-specific).
        self._bonus_array = np.full(4, 0.02, dtype=np.float32)

        # Per-class min box area (px^2), loose so nothing is silently dropped
        # before tuning: distant cars/motorcycles are small in 1024x1024 frames.
        self._min_box_area_array = np.array([12.0] * 4, dtype=np.float32)
        self.min_side = 3
        self.max_aspect_ratio = 12.0
        self.max_area_frac = 0.95

        # Square-domain de-squish: widen a square (validator) frame by this
        # fraction before letterboxing, then drop pad-center boxes and un-pad.
        # 0.0 = off — the model was TRAINED on squished 1024 frames, so the
        # domains already match; enabling changes aspect -> validate first.
        self.square_pad_frac = 0.0

        # Dual-label bus/truck hedging (rival tactic, see header). 0.0 = OFF.
        # e.g. 0.6 duplicates every kept bus/truck box with conf < 0.6 under the
        # sibling class at the same conf. Applied AFTER cross-class dedup (a
        # duplicate at IoU 1.0 would otherwise be deduped straight back out).
        self.dual_label_pair = ("bus", "truck")
        self.dual_label_below = 0.0

        self._avg_iou = float(np.mean(self._iou_thres_array))
        print(f"ONNX model loaded from: {model_path}")
        print(f"ONNX input: name={self.input_name}, shape={self.input_shape}")
        self._warmup()

    # ── Setup helpers ────────────────────────────────────────────────────────

    def _warmup(self, iters: int = 3) -> None:
        try:
            dummy = np.zeros((1024, 1024, 3), dtype=np.uint8)  # validator domain
            for _ in range(max(1, iters)):
                self.predict_batch(batch_images=[dummy], offset=0, n_keypoints=0)
            print(f"warmup: {iters} dummy predict_batch call(s) done")
        except Exception as e:
            print(f"warmup skipped: {e}")

    def __repr__(self) -> str:
        return f"VehicleMiner(classes={len(self.class_names)}, providers={self.session.get_providers()})"

    @staticmethod
    def _safe_dim(value, default: int) -> int:
        return value if isinstance(value, int) and value > 0 else default

    @staticmethod
    def _resolve_model_path(repo: Path) -> Path:
        """Prefer weights.onnx, then weights_int8.onnx, then any .onnx."""
        for name in ("weights.onnx", "weights_int8.onnx"):
            p = repo / name
            if p.exists():
                if name != "weights.onnx":
                    print(f"model: weights.onnx not found, using {name}")
                return p
        candidates = sorted(repo.glob("*.onnx"))
        if candidates:
            print(f"model: using {candidates[0].name}")
            return candidates[0]
        return repo / "weights.onnx"  # let session creation raise the error

    def _read_model_class_order(self) -> list[str] | None:
        """Read class order from Ultralytics ONNX `names` metadata, or None."""
        try:
            import ast
            meta = self.session.get_modelmeta().custom_metadata_map
            names = ast.literal_eval(meta["names"])
            order = ([str(names[i]) for i in sorted(names)] if isinstance(names, dict)
                     else [str(n) for n in names])
        except Exception as e:
            print(f"cls order: could not read ONNX names metadata ({e})")
            return None
        if sorted(order) != sorted(self.class_names):
            print(f"cls order: ONNX names {order} != expected {self.class_names}; ignoring")
            return None
        return order

    # ── Preprocessing ────────────────────────────────────────────────────────

    def _letterbox(self, image: ndarray, new_shape: tuple[int, int],
                   color=(114, 114, 114)) -> tuple[ndarray, float, tuple[float, float]]:
        h, w = image.shape[:2]
        new_w, new_h = new_shape
        ratio = min(new_w / w, new_h / h)
        rw, rh = int(round(w * ratio)), int(round(h * ratio))
        if (rw, rh) != (w, h):
            interp = cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR
            image = cv2.resize(image, (rw, rh), interpolation=interp)
        dw = (new_w - rw) / 2.0
        dh = (new_h - rh) / 2.0
        left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
        top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
        padded = cv2.copyMakeBorder(image, top, bottom, left, right,
                                    cv2.BORDER_CONSTANT, value=color)
        return padded, ratio, (dw, dh)

    def _preprocess(self, image: ndarray) -> tuple[np.ndarray, dict]:
        """Letterbox to the ONNX input size; NO CLAHE/denoise/sharpen (any
        enhancement the model wasn't trained on is a train/test mismatch that
        HURTS accuracy and also risks the CPU latency gate).
        Square-domain de-squish: when the frame is square (the validator's
        1024x1024 squished eval image) and square_pad_frac > 0, widen it with
        gray bars first -- objects come out horizontally compressed in those
        frames, and mild widening partly restores their aspect.
        """
        orig_h, orig_w = image.shape[:2]
        extra_left = extra_right = 0
        if self.square_pad_frac > 0.0 and orig_w == orig_h:
            target_w = int(orig_w * (1.0 + self.square_pad_frac))
            if target_w > orig_w:
                total = target_w - orig_w
                extra_left = total // 2
                extra_right = total - extra_left
                image = cv2.copyMakeBorder(image, 0, 0, extra_left, extra_right,
                                           cv2.BORDER_CONSTANT, value=(114, 114, 114))

        img, ratio, pad = self._letterbox(image, (self.input_width, self.input_height))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = img.astype(np.float32) / 255.0
        img = np.transpose(img, (2, 0, 1))[None, ...]
        img = np.ascontiguousarray(img, dtype=self.np_dtype)
        return img, {
            "ratio": ratio, "pad": pad, "orig_size": (orig_w, orig_h),
            "extra_left": extra_left, "extra_right": extra_right,
        }

    # ── Vectorized box ops ───────────────────────────────────────────────────

    @staticmethod
    def _clip_boxes(boxes: np.ndarray, image_size: tuple[int, int]) -> np.ndarray:
        w, h = image_size
        boxes[:, 0] = np.clip(boxes[:, 0], 0, w - 1)
        boxes[:, 1] = np.clip(boxes[:, 1], 0, h - 1)
        boxes[:, 2] = np.clip(boxes[:, 2], 0, w - 1)
        boxes[:, 3] = np.clip(boxes[:, 3], 0, h - 1)
        return boxes

    @staticmethod
    def _xywh_to_xyxy(boxes: np.ndarray) -> np.ndarray:
        out = np.empty_like(boxes)
        out[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.0
        out[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.0
        out[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.0
        out[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
        return out

    @staticmethod
    def _hard_nms(boxes: np.ndarray, scores: np.ndarray, iou_thresh: float) -> np.ndarray:
        """Vectorized greedy NMS: areas precomputed once, each step is a single
        numpy IoU vector -- no inner Python loop."""
        n = len(boxes)
        if n == 0:
            return np.array([], dtype=np.intp)
        x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
        areas = np.maximum(0.0, x2 - x1) * np.maximum(0.0, y2 - y1)
        order = np.argsort(-scores)
        keep = []
        while order.size > 0:
            i = int(order[0])
            keep.append(i)
            if order.size == 1:
                break
            rest = order[1:]
            xx1 = np.maximum(x1[i], x1[rest])
            yy1 = np.maximum(y1[i], y1[rest])
            xx2 = np.minimum(x2[i], x2[rest])
            yy2 = np.minimum(y2[i], y2[rest])
            inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
            iou = inter / (areas[i] + areas[rest] - inter + 1e-7)
            order = rest[iou <= iou_thresh]
        return np.array(keep, dtype=np.intp)

    @staticmethod
    def _soft_nms(boxes: np.ndarray, scores: np.ndarray, sigma: float
                  ) -> tuple[np.ndarray, np.ndarray]:
        """Gaussian soft-NMS: decay overlapping scores instead of deleting.
        Returns (original indices in decay order, decayed scores)."""
        n = len(boxes)
        if n == 0:
            return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
        b = boxes.astype(np.float32, copy=True)
        sc = scores.astype(np.float32, copy=True)
        order = np.arange(n)
        for i in range(n):
            mp = i + int(np.argmax(sc[i:]))
            b[[i, mp]] = b[[mp, i]]; sc[[i, mp]] = sc[[mp, i]]
            order[[i, mp]] = order[[mp, i]]
            if i + 1 >= n:
                break
            xx1 = np.maximum(b[i, 0], b[i+1:, 0]); yy1 = np.maximum(b[i, 1], b[i+1:, 1])
            xx2 = np.minimum(b[i, 2], b[i+1:, 2]); yy2 = np.minimum(b[i, 3], b[i+1:, 3])
            inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
            ai = max(0.0, float((b[i, 2]-b[i, 0]) * (b[i, 3]-b[i, 1])))
            aj = np.maximum(0.0, b[i+1:, 2]-b[i+1:, 0]) * np.maximum(0.0, b[i+1:, 3]-b[i+1:, 1])
            iou = inter / (ai + aj - inter + 1e-7)
            sc[i+1:] *= np.exp(-(iou ** 2) / sigma)
        return order, sc

    def _per_class_soft_nms(self, boxes: np.ndarray, scores: np.ndarray,
                            cls_ids: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        """Soft-NMS per class, then re-filter decayed scores against the
        per-class confidence thresholds so no decayed box reaches the FP pillar
        below its class bar."""
        keep_idx: list[int] = []
        keep_sc: list[float] = []
        for c in np.unique(cls_ids):
            idx = np.where(cls_ids == c)[0]
            order, dec = self._soft_nms(boxes[idx], scores[idx], self.soft_nms_sigma)
            thr = float(self._conf_thres_array[c])
            for k, s_dec in zip(order, dec):
                if s_dec >= thr:
                    keep_idx.append(int(idx[k])); keep_sc.append(float(s_dec))
        return np.array(keep_idx, dtype=np.intp), np.array(keep_sc, dtype=np.float32)

    def _per_class_hard_nms(self, boxes: np.ndarray, scores: np.ndarray,
                            cls_ids: np.ndarray) -> np.ndarray:
        """Per-class NMS using the per-class IoU thresholds."""
        if len(boxes) == 0:
            return np.array([], dtype=np.intp)
        all_keep: list[int] = []
        for c in np.unique(cls_ids):
            idx = np.where(cls_ids == c)[0]
            cls_iou = float(self._iou_thres_array[c])
            keep = self._hard_nms(boxes[idx], scores[idx], cls_iou)
            all_keep.extend(idx[keep].tolist())
        all_keep.sort()
        return np.array(all_keep, dtype=np.intp)

    @staticmethod
    def _max_score_per_cluster(post_boxes: np.ndarray, post_cls: np.ndarray,
                               full_boxes: np.ndarray, full_scores: np.ndarray,
                               full_cls: np.ndarray, iou_thresh: float) -> np.ndarray:
        """Each survivor's confidence -> max score in its SAME-CLASS IoU cluster.
        Vectorized single (n_post x n_full) IoU matrix."""
        n = len(post_boxes)
        if n == 0:
            return np.empty(0, dtype=np.float32)
        m = len(full_boxes)
        if m == 0:
            return np.zeros(n, dtype=np.float32)
        pa = (np.maximum(0.0, post_boxes[:, 2] - post_boxes[:, 0]) *
              np.maximum(0.0, post_boxes[:, 3] - post_boxes[:, 1]))
        fa = (np.maximum(0.0, full_boxes[:, 2] - full_boxes[:, 0]) *
              np.maximum(0.0, full_boxes[:, 3] - full_boxes[:, 1]))
        xx1 = np.maximum(post_boxes[:, 0][:, None], full_boxes[:, 0][None, :])
        yy1 = np.maximum(post_boxes[:, 1][:, None], full_boxes[:, 1][None, :])
        xx2 = np.minimum(post_boxes[:, 2][:, None], full_boxes[:, 2][None, :])
        yy2 = np.minimum(post_boxes[:, 3][:, None], full_boxes[:, 3][None, :])
        inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
        iou = inter / (pa[:, None] + fa[None, :] - inter + 1e-7)
        mask = (iou >= iou_thresh) & (post_cls[:, None] == full_cls[None, :])
        tiled = np.where(mask, full_scores[None, :], -np.inf)
        out = tiled.max(axis=1)
        out[~np.isfinite(out)] = 0.0
        return out.astype(np.float32)

    def _conf_filter_mask(self, scores: np.ndarray, cls_ids: np.ndarray,
                          extra_left: int) -> np.ndarray:
        """Per-class threshold (square vs non-square set) + per-class rescue."""
        if len(scores) == 0:
            return np.zeros(0, dtype=bool)
        thr_arr = self._extra_conf_thres_array if extra_left > 0 else self._conf_thres_array
        keep = scores >= thr_arr[cls_ids]
        for c in np.unique(cls_ids):
            b = float(self._bonus_array[c])
            if b <= 0.0:
                continue
            cm = cls_ids == c
            if keep[cm].any():
                continue
            idx = np.where(cm)[0]
            top = int(idx[int(np.argmax(scores[idx]))])
            if scores[top] >= float(self._conf_thres_array[c]) - b:
                keep[top] = True
        return keep

    def _filter_sane_boxes(self, boxes: np.ndarray, scores: np.ndarray,
                           cls_ids: np.ndarray, orig_size: tuple[int, int]
                           ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Vectorized per-class min-area / max-area / min-side / aspect filter."""
        if len(boxes) == 0:
            return boxes, scores, cls_ids
        orig_w, orig_h = orig_size
        image_area = float(orig_w * orig_h)
        bw = np.maximum(0.0, boxes[:, 2] - boxes[:, 0])
        bh = np.maximum(0.0, boxes[:, 3] - boxes[:, 1])
        area = bw * bh
        ar = np.maximum(bw / np.maximum(bh, 1e-6), bh / np.maximum(bw, 1e-6))
        keep = (
            (bw >= self.min_side) & (bh >= self.min_side) &
            (area >= self._min_box_area_array[cls_ids]) &
            (area <= self.max_area_frac * image_area) &
            (ar <= self.max_aspect_ratio)
        )
        return boxes[keep], scores[keep], cls_ids[keep]

    def _cross_class_dedup_op(self, boxes: np.ndarray, scores: np.ndarray,
                              cls_ids: np.ndarray, iou_thresh: float
                              ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Suppress near-duplicate boxes ACROSS classes (same physical vehicle
        firing 2 classes, e.g. van -> car+van, bus -> bus+truck). Order by
        conf-margin then area; keep highest, drop IoU>thresh others."""
        n = len(boxes)
        if n <= 1:
            return boxes, scores, cls_ids
        boxes = np.asarray(boxes, dtype=np.float32)
        scores = np.asarray(scores, dtype=np.float32)
        cls_ids = np.asarray(cls_ids, dtype=np.int32)
        areas = (np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) *
                 np.maximum(0.0, boxes[:, 3] - boxes[:, 1]))
        margins = scores - self._conf_thres_array[cls_ids]
        order = np.lexsort((-areas, -margins))
        suppressed = np.zeros(n, dtype=bool)
        keep: list[int] = []
        for i in order:
            if suppressed[i]:
                continue
            keep.append(int(i))
            bi = boxes[i]
            xx1 = np.maximum(bi[0], boxes[:, 0])
            yy1 = np.maximum(bi[1], boxes[:, 1])
            xx2 = np.minimum(bi[2], boxes[:, 2])
            yy2 = np.minimum(bi[3], boxes[:, 3])
            inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
            a_i = max(1e-7, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
            iou = inter / (a_i + areas - inter + 1e-7)
            dup = iou > iou_thresh
            dup[i] = False
            suppressed |= dup
        keep_idx = np.array(keep, dtype=np.intp)
        return boxes[keep_idx], scores[keep_idx], cls_ids[keep_idx]

    def _apply_dual_label(self, boxes: np.ndarray, scores: np.ndarray,
                          cls_ids: np.ndarray
                          ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Duplicate low-confidence bus/truck boxes under the sibling class
        (hedges the bus<->truck confusion against per-class greedy matching).
        Must run AFTER cross-class dedup or the copy is immediately removed."""
        if self.dual_label_below <= 0.0 or len(boxes) == 0:
            return boxes, scores, cls_ids
        a = self.class_names.index(self.dual_label_pair[0])
        b = self.class_names.index(self.dual_label_pair[1])
        sib = np.full(len(self.class_names), -1, dtype=np.int32)
        sib[a], sib[b] = b, a
        mask = (sib[cls_ids] >= 0) & (scores < self.dual_label_below)
        if not mask.any():
            return boxes, scores, cls_ids
        boxes = np.concatenate([boxes, boxes[mask]], axis=0)
        scores = np.concatenate([scores, scores[mask]], axis=0)
        cls_ids = np.concatenate([cls_ids, sib[cls_ids[mask]]], axis=0)
        return boxes, scores, cls_ids

    def _per_view_pipeline(self, boxes: np.ndarray, scores: np.ndarray,
                           cls_ids: np.ndarray, orig_size: tuple[int, int]
                           ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """sane filter -> top-k cap -> per-class NMS -> max_det cap -> cross-class dedup."""
        boxes, scores, cls_ids = self._filter_sane_boxes(boxes, scores, cls_ids, orig_size)
        if len(boxes) == 0:
            return boxes, scores, cls_ids
        if len(scores) > self.pre_nms_topk:
            top = np.argpartition(-scores, self.pre_nms_topk)[: self.pre_nms_topk]
            boxes, scores, cls_ids = boxes[top], scores[top], cls_ids[top]
        if len(boxes) > 1:
            if self.use_soft_nms:
                keep, scores = self._per_class_soft_nms(boxes, scores, cls_ids)
                boxes, cls_ids = boxes[keep], cls_ids[keep]
            else:
                keep = self._per_class_hard_nms(boxes, scores, cls_ids)
                boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep]
        if len(scores) > self.max_det:
            top = np.argsort(-scores)[: self.max_det]
            boxes, scores, cls_ids = boxes[top], scores[top], cls_ids[top]
        if len(boxes) > 1:
            boxes, scores, cls_ids = self._cross_class_dedup_op(
                boxes, scores, cls_ids, self.cross_iou_thresh)
        return boxes, scores, cls_ids

    # ── Coordinate un-mapping (shared by both decode paths) ──────────────────

    def _unmap_and_finish(self, boxes: np.ndarray, scores: np.ndarray,
                          cls_ids: np.ndarray, meta: dict
                          ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Undo letterbox + square-pad, drop pad-center boxes, clip, post-process."""
        pad_w, pad_h = meta["pad"]
        ratio = meta["ratio"]
        orig_w, orig_h = meta["orig_size"]
        extra_left, extra_right = meta["extra_left"], meta["extra_right"]

        boxes[:, [0, 2]] -= pad_w
        boxes[:, [1, 3]] -= pad_h
        boxes /= ratio
        if extra_left:
            boxes[:, [0, 2]] -= extra_left
        if extra_left or extra_right:
            cx = (boxes[:, 0] + boxes[:, 2]) * 0.5
            inside = (cx >= 0) & (cx <= orig_w)
            boxes, scores, cls_ids = boxes[inside], scores[inside], cls_ids[inside]
            if len(boxes) == 0:
                return boxes, scores, cls_ids
        boxes = self._clip_boxes(boxes, (orig_w, orig_h))
        return self._per_view_pipeline(boxes, scores, cls_ids, (orig_w, orig_h))

    # ── Decoding ─────────────────────────────────────────────────────────────

    def _decode_final_dets(self, preds: np.ndarray, meta: dict
                           ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """End2end output rows [x1, y1, x2, y2, conf, cls_id] (NMS in graph)."""
        empty = (np.empty((0, 4), np.float32), np.empty(0, np.float32), np.empty(0, np.int32))
        if preds.ndim == 3 and preds.shape[0] == 1:
            preds = preds[0]
        if preds.ndim != 2 or preds.shape[1] < 6:
            raise ValueError(f"Unexpected final-det output shape: {preds.shape}")
        boxes = preds[:, :4].astype(np.float32)
        scores = preds[:, 4].astype(np.float32)
        cls_ids = self.cls_remap[preds[:, 5].astype(np.int32)]
        keep = self._conf_filter_mask(scores, cls_ids, meta["extra_left"])
        boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep]
        if len(boxes) == 0:
            return empty
        return self._unmap_and_finish(boxes, scores, cls_ids, meta)

    def _decode_raw_yolo(self, preds: np.ndarray, meta: dict
                         ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Raw YOLO output [1,C,N] or [1,N,C]; xywh + per-class scores."""
        empty = (np.empty((0, 4), np.float32), np.empty(0, np.float32), np.empty(0, np.int32))
        if preds.ndim != 3 or preds.shape[0] != 1:
            raise ValueError(f"Unexpected raw output shape: {preds.shape}")
        preds = preds[0]
        if preds.shape[0] <= 16 and preds.shape[1] > preds.shape[0]:
            preds = preds.T
        if preds.ndim != 2 or preds.shape[1] < 5:
            raise ValueError(f"Unexpected normalized raw shape: {preds.shape}")
        boxes_xywh = preds[:, :4].astype(np.float32)
        cls_part = preds[:, 4:].astype(np.float32)
        if cls_part.shape[1] == 1:
            scores = cls_part[:, 0]
            cls_ids = np.zeros(len(scores), dtype=np.int32)
        else:
            cls_ids = np.argmax(cls_part, axis=1).astype(np.int32)
            scores = cls_part[np.arange(len(cls_part)), cls_ids]
        cls_ids = self.cls_remap[cls_ids]
        keep = self._conf_filter_mask(scores, cls_ids, meta["extra_left"])
        boxes_xywh, scores, cls_ids = boxes_xywh[keep], scores[keep], cls_ids[keep]
        if len(boxes_xywh) == 0:
            return empty
        boxes = self._xywh_to_xyxy(boxes_xywh)
        return self._unmap_and_finish(boxes, scores, cls_ids, meta)

    def _decode(self, output: np.ndarray, meta: dict
                ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        if output.ndim == 2 and output.shape[1] >= 6:
            return self._decode_final_dets(output, meta)
        if output.ndim == 3 and output.shape[0] == 1 and output.shape[2] == 6:
            return self._decode_final_dets(output, meta)
        return self._decode_raw_yolo(output, meta)

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

    def _predict_single_arrays(self, image: np.ndarray
                               ) -> tuple[np.ndarray, np.ndarray, np.ndarray, dict]:
        if not isinstance(image, np.ndarray) or image.ndim != 3 or image.shape[2] != 3:
            raise ValueError(f"Expected HWC BGR image, got {type(image)} {getattr(image,'shape',None)}")
        if image.dtype != np.uint8:
            image = image.astype(np.uint8)
        inp, meta = self._preprocess(image)
        outputs = self.session.run(self.output_names, {self.input_name: inp})
        boxes, scores, cls_ids = self._decode(outputs[0], meta)
        return boxes, scores, cls_ids, meta

    @staticmethod
    def _to_boxes(boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray,
                  orig_size: tuple[int, int]) -> list[BoundingBox]:
        results: list[BoundingBox] = []
        orig_w, orig_h = orig_size
        for box, conf, cls_id in zip(boxes, scores, cls_ids):
            x1, y1, x2, y2 = box.tolist()
            if x2 <= x1 or y2 <= y1:
                continue
            results.append(BoundingBox(
                x1=max(0, min(orig_w, int(math.floor(x1)))),
                y1=max(0, min(orig_h, int(math.floor(y1)))),
                x2=max(0, min(orig_w, int(math.ceil(x2)))),
                y2=max(0, min(orig_h, int(math.ceil(y2)))),
                cls_id=int(cls_id),
                conf=float(max(0.0, min(1.0, conf))),
            ))
        return results

    def _infer_single(self, image: np.ndarray) -> list[BoundingBox]:
        boxes, scores, cls_ids, meta = self._predict_single_arrays(image)
        if len(boxes) == 0:
            return []
        if self.use_cluster_boost and len(boxes) > 1:
            scores = self._max_score_per_cluster(
                boxes, cls_ids, boxes, scores, cls_ids, self._avg_iou)
        boxes, scores, cls_ids = self._apply_dual_label(boxes, scores, cls_ids)
        return self._to_boxes(boxes, scores, cls_ids, meta["orig_size"])

    def _infer_tta(self, image: np.ndarray) -> list[BoundingBox]:
        """Horizontal-flip TTA: union of original + flipped, per-class NMS,
        same-class cluster-max boost, cross-class dedup."""
        b0, s0, c0, meta = self._predict_single_arrays(image)
        w = image.shape[1]
        bf, sf, cf, _ = self._predict_single_arrays(cv2.flip(image, 1))
        if len(bf):
            bf = bf.copy()
            bf[:, [0, 2]] = w - bf[:, [2, 0]]  # mirror x back to original coords
        coords = np.concatenate([b0, bf], axis=0) if len(bf) else b0
        scores = np.concatenate([s0, sf], axis=0) if len(sf) else s0
        cls_ids = np.concatenate([c0, cf], axis=0) if len(cf) else c0
        if len(coords) == 0:
            return []
        keep = self._per_class_hard_nms(coords, scores, cls_ids)
        if len(keep) == 0:
            return []
        if len(keep) > self.max_det:
            keep = keep[np.argsort(-scores[keep])[: self.max_det]]
        boosted = self._max_score_per_cluster(
            coords[keep], cls_ids[keep], coords, scores, cls_ids, self._avg_iou)
        kb, kc = coords[keep], cls_ids[keep]
        if len(kb) > 1:
            kb, boosted, kc = self._cross_class_dedup_op(kb, boosted, kc, self.cross_iou_thresh)
        kb, boosted, kc = self._apply_dual_label(kb, boosted, kc)
        return self._to_boxes(kb, boosted, kc, meta["orig_size"])

    def predict_batch(self, batch_images: list[ndarray], offset: int,
                      n_keypoints: int) -> list[TVFrameResult]:
        results: list[TVFrameResult] = []
        for idx, image in enumerate(batch_images):
            try:
                boxes = self._infer_tta(image) if self.use_tta else self._infer_single(image)
            except Exception as e:
                print(f"Inference failed for frame {offset + idx}: {e}")
                boxes = []
            keypoints = [(0, 0) for _ in range(max(0, int(n_keypoints)))]
            results.append(TVFrameResult(frame_id=offset + idx, boxes=boxes, keypoints=keypoints))
        return results