scorevision: push artifact

miner.py

@@ -1,19 +1,5 @@
-"""manak0/Detect-crime miner — 3-view TTA + Weighted Box Fusion.
-
-Same interface as king's miner.py (Miner.predict_batch returns list[TVFrameResult]).
-Differences vs king:
-  • 3 forward views: orig 1280, hflip 1280, zoomed 1408 (re-letterboxed)
-  • WBF fusion across views (better than score-tier consensus for recall)
-  • Per-class confidence thresholds (calibrated on SAM3-distilled val)
-  • Same per-class NMS at IoU=0.35 + same sane-box filters as king
-
-ONNX expected: input `images [1,3,H,W]`, output `[1, max_det, 6]` = (x1,y1,x2,y2,score,cls).
-"""
-from __future__ import annotations
-
-import math
 from pathlib import Path
-from typing import Any
+import math
 
 import cv2
 import numpy as np
@@ -22,34 +8,6 @@ from numpy import ndarray
 from pydantic import BaseModel
 
 
-CRIME_CLASSES = ["balaclava", "hoodie", "glove", "bat", "spray paint", "graffiti"]
-
-# Per-class confidence thresholds. Calibrated against the REAL validator
-# challenge images (133 images pulled from public R2 response shards,
-# scripts/05_calibrate_thresholds.py --val-dir dataset/real_chal/val).
-# v2 fine-tuned on web+real_chal mix; this calibration on real_chal val:
-# composite=0.571, map50=0.758, FP/img=1.09, fp_pillar=0.891.
-# Compare: king on same val=0.233, coolroman=0.271.
-PER_CLASS_CONF = np.array([
-    0.400,  # balaclava
-    0.500,  # hoodie
-    0.450,  # glove
-    0.325,  # bat
-    0.325,  # spray paint
-    0.500,  # graffiti
-], dtype=np.float32)
-
-# Per-class aspect-ratio constraint: keep box only if height/width >= ratio.
-PER_CLASS_MIN_H_OVER_W: dict[int, float] = {}
-
-# WBF parameters
-WBF_IOU_THR = 0.55
-WBF_SKIP_BOX_THR = 0.05  # below this, treat as no signal even pre-fusion
-
-# View weights (orig, hflip, zoomed-1408)
-VIEW_WEIGHTS = np.array([1.0, 1.0, 0.85], dtype=np.float32)
-
-
 class BoundingBox(BaseModel):
     x1: int
     y1: int
@@ -65,437 +23,492 @@ class TVFrameResult(BaseModel):
     keypoints: list[tuple[int, int]]
 
 
-# ===================== Weighted Box Fusion =====================
-def _wbf_one_class(
-    boxes: np.ndarray,         # [N,4] xyxy in image pixels
-    scores: np.ndarray,        # [N]
-    weights: np.ndarray,       # [N] view-weights per box
-    iou_thr: float,
-) -> tuple[np.ndarray, np.ndarray]:
-    """Single-class Weighted Box Fusion. Returns fused (boxes, scores)."""
-    if len(boxes) == 0:
-        return np.empty((0, 4), dtype=np.float32), np.empty((0,), dtype=np.float32)
-
-    order = np.argsort(scores)[::-1]
-    boxes = boxes[order]
-    scores = scores[order]
-    weights = weights[order]
-
-    clusters: list[list[int]] = []
-    fused_boxes: list[np.ndarray] = []
-
-    for i in range(len(boxes)):
-        bi = boxes[i]
-        best_j = -1
-        best_iou = 0.0
-        for j, fb in enumerate(fused_boxes):
-            xx1 = max(bi[0], fb[0]); yy1 = max(bi[1], fb[1])
-            xx2 = min(bi[2], fb[2]); yy2 = min(bi[3], fb[3])
-            iw = max(0.0, xx2 - xx1); ih = max(0.0, yy2 - yy1)
-            inter = iw * ih
-            area_i = max(0.0, bi[2] - bi[0]) * max(0.0, bi[3] - bi[1])
-            area_f = max(0.0, fb[2] - fb[0]) * max(0.0, fb[3] - fb[1])
-            union = area_i + area_f - inter
-            iou = inter / union if union > 0 else 0.0
-            if iou > best_iou:
-                best_iou = iou; best_j = j
-        if best_j >= 0 and best_iou >= iou_thr:
-            clusters[best_j].append(i)
-            members = clusters[best_j]
-            ws = weights[members] * scores[members]
-            wsum = ws.sum() + 1e-9
-            fused = (boxes[members] * ws[:, None]).sum(axis=0) / wsum
-            fused_boxes[best_j] = fused
-        else:
-            clusters.append([i])
-            fused_boxes.append(bi.copy())
-
-    out_boxes = np.stack(fused_boxes, axis=0).astype(np.float32)
-    out_scores = np.zeros(len(clusters), dtype=np.float32)
-    n_views = max(1.0, float(weights.max()) if len(weights) else 1.0)
-    for k, members in enumerate(clusters):
-        # WBF score: avg score weighted by view-weight, scaled by cluster coverage
-        ws = weights[members]
-        s = scores[members]
-        avg = float((s * ws).sum() / (ws.sum() + 1e-9))
-        coverage = min(1.0, len(members) / n_views)  # 1.0 if all views agreed
-        out_scores[k] = avg * (0.5 + 0.5 * coverage)
-    return out_boxes, out_scores
-
-
-def weighted_box_fusion(
-    boxes_per_view: list[np.ndarray],  # each [Ni,4]
-    scores_per_view: list[np.ndarray],  # each [Ni]
-    cls_per_view: list[np.ndarray],     # each [Ni]
-    view_weights: np.ndarray,
-    iou_thr: float = WBF_IOU_THR,
-) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-    all_b: list[np.ndarray] = []
-    all_s: list[np.ndarray] = []
-    all_w: list[np.ndarray] = []
-    all_c: list[np.ndarray] = []
-    for b, s, c, w in zip(boxes_per_view, scores_per_view, cls_per_view, view_weights):
-        if len(b) == 0:
-            continue
-        all_b.append(b); all_s.append(s); all_c.append(c)
-        all_w.append(np.full(len(b), float(w), dtype=np.float32))
-    if not all_b:
-        return (np.empty((0, 4), dtype=np.float32),
-                np.empty((0,), dtype=np.float32),
-                np.empty((0,), dtype=np.int32))
-    B = np.concatenate(all_b, axis=0)
-    S = np.concatenate(all_s, axis=0)
-    C = np.concatenate(all_c, axis=0)
-    W = np.concatenate(all_w, axis=0)
-
-    out_b: list[np.ndarray] = []
-    out_s: list[np.ndarray] = []
-    out_c: list[np.ndarray] = []
-    for c in np.unique(C):
-        mask = C == c
-        fb, fs = _wbf_one_class(B[mask], S[mask], W[mask], iou_thr)
-        if len(fb) == 0:
-            continue
-        out_b.append(fb)
-        out_s.append(fs)
-        out_c.append(np.full(len(fb), int(c), dtype=np.int32))
-    if not out_b:
-        return (np.empty((0, 4), dtype=np.float32),
-                np.empty((0,), dtype=np.float32),
-                np.empty((0,), dtype=np.int32))
-    return np.concatenate(out_b, axis=0), np.concatenate(out_s, axis=0), np.concatenate(out_c, axis=0)
-
-
-# ===================== Miner =====================
 class Miner:
+    """
+    ONNX Runtime miner for Detect-crime.
+
+    Recipe (see recipe.md for derivation):
+      * per-class confidence thresholds (king's lever)
+      * per-class IoU thresholds (closer pairs allowed for graffiti / spray paint)
+      * per-class soft-NMS with sigma=0.5
+      * horizontal-flip TTA merged with **weighted box fusion** (coord-averaged)
+      * per-class cluster confidence boost (fixes king's cross-class leak)
+      * no min-side / min-area filter (small TPs survive)
+      * no h/w aspect gate beyond a sane outlier cap (8x)
+    """
+
+    class_names = ["balaclava", "hoodie", "glove", "bat", "spray paint", "graffiti"]
+    input_size = 1280
+    max_det = 300
+    max_aspect_ratio = 8.0
+    soft_sigma = 0.5
+
+    _conf_thres_array = np.array(
+        # balaclava, hoodie, glove, bat, spray paint, graffiti
+        [0.50, 0.60, 0.30, 0.20, 0.45, 0.30],
+        dtype=np.float32,
+    )
+
+    _iou_thres_array = np.array(
+        # hoodie/balaclava: clean people, tight NMS
+        # glove/bat: small rare objects, mid
+        # spray/graffiti: legitimately overlapping marks, loose
+        [0.60, 0.60, 0.55, 0.55, 0.45, 0.45],
+        dtype=np.float32,
+    )
+
     def __init__(self, path_hf_repo: Path) -> None:
         model_path = path_hf_repo / "weights.onnx"
-        self.class_names = CRIME_CLASSES
         print("ORT version:", ort.__version__)
+
         try:
             ort.preload_dlls()
-        except Exception:
-            pass
+            print("preload_dlls success")
+        except Exception as e:
+            print(f"preload_dlls failed: {e}")
+
+        print("ORT providers available:", ort.get_available_providers())
+
+        sess_options = ort.SessionOptions()
+        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
 
-        sess_opt = ort.SessionOptions()
-        sess_opt.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
         try:
             self.session = ort.InferenceSession(
-                str(model_path), sess_options=sess_opt,
+                str(model_path),
+                sess_options=sess_options,
                 providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
             )
+            print("ORT session: CUDA provider")
         except Exception as e:
-            print(f"⚠️ CUDA session failed, CPU fallback: {e}")
+            print(f"CUDA session creation failed, falling back to CPU: {e}")
             self.session = ort.InferenceSession(
-                str(model_path), sess_options=sess_opt, providers=["CPUExecutionProvider"]
+                str(model_path),
+                sess_options=sess_options,
+                providers=["CPUExecutionProvider"],
             )
+
         print("ORT session providers:", self.session.get_providers())
 
+        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()]
-        in_meta = self.session.get_inputs()[0]
-        ishape = in_meta.shape
-        self.input_h = self._safe_dim(ishape[2], 1280)
-        self.input_w = self._safe_dim(ishape[3], 1280)
-        self.input_np_dtype = np.float16 if "float16" in (in_meta.type or "") else np.float32
-
-        # Inference recipe
-        self.use_tta = True
-        self.zoom_long_side = 1408     # second view at higher res, then letterbox to model input
-        self.iou_thres = 0.35
-        self.max_det = 200
-        self.per_class_conf = PER_CLASS_CONF.copy()
-
-        # box sanity filters (match king — they're well-tuned for crime)
-        self.min_box_area = 10 * 10
-        self.min_w = 6
-        self.min_h = 6
-        self.max_aspect_ratio = 8.0
-        self.max_box_area_ratio = 0.95
-
-        print(f"✅ ONNX loaded {model_path}  input={self.input_w}x{self.input_h}")
+        self.input_shape = self.session.get_inputs()[0].shape
+        # Detect FP16 vs FP32 input from the session metadata.
+        input_type = self.session.get_inputs()[0].type
+        self.input_dtype = np.float16 if "float16" in input_type else np.float32
+        self.input_height = self._safe_dim(self.input_shape[2], default=self.input_size)
+        self.input_width = self._safe_dim(self.input_shape[3], default=self.input_size)
+
+        print(f"ONNX model loaded: {model_path}")
+        print("per-class conf: " + ", ".join(
+            f"{n}={t:.2f}" for n, t in zip(self.class_names,
+                                           self._conf_thres_array.tolist())))
+        print("per-class iou : " + ", ".join(
+            f"{n}={t:.2f}" for n, t in zip(self.class_names,
+                                           self._iou_thres_array.tolist())))
+
+    def __repr__(self) -> str:
+        return (
+            f"ONNXRuntime(session={type(self.session).__name__}, "
+            f"providers={self.session.get_providers()})"
+        )
 
     @staticmethod
-    def _safe_dim(v, default):
-        return v if isinstance(v, int) and v > 0 else default
+    def _safe_dim(value, default: int) -> int:
+        return value if isinstance(value, int) and value > 0 else default
 
-    # ---- preprocessing ----
-    def _letterbox(self, image: ndarray, new_shape: tuple[int, int]):
+    # ---------------------------------------------------------------- preproc
+    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]
-        nw, nh = new_shape
-        ratio = min(nw / w, nh / h)
-        rw, rh = int(round(w * ratio)), int(round(h * ratio))
-        if (rw, rh) != (w, h):
+        new_w, new_h = new_shape
+        ratio = min(new_w / w, new_h / h)
+        resized_w = int(round(w * ratio))
+        resized_h = int(round(h * ratio))
+        if (resized_w, resized_h) != (w, h):
             interp = cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR
-            image = cv2.resize(image, (rw, rh), interpolation=interp)
-        dw, dh = (nw - rw) / 2.0, (nh - rh) / 2.0
-        l = int(round(dw - 0.1)); r = int(round(dw + 0.1))
-        t = int(round(dh - 0.1)); b = int(round(dh + 0.1))
-        padded = cv2.copyMakeBorder(image, t, b, l, r, borderType=cv2.BORDER_CONSTANT, value=(114, 114, 114))
+            image = cv2.resize(image, (resized_w, resized_h), interpolation=interp)
+        dw = (new_w - resized_w) / 2.0
+        dh = (new_h - resized_h) / 2.0
+        left = int(round(dw - 0.1))
+        right = int(round(dw + 0.1))
+        top = int(round(dh - 0.1))
+        bottom = int(round(dh + 0.1))
+        padded = cv2.copyMakeBorder(
+            image, top, bottom, left, right,
+            borderType=cv2.BORDER_CONSTANT, value=color,
+        )
         return padded, ratio, (dw, dh)
 
-    def _preproc(self, image: ndarray, target_long: int | None = None):
-        if target_long is not None:
-            h, w = image.shape[:2]
-            if max(h, w) != target_long:
-                s = target_long / max(h, w)
-                image = cv2.resize(image, (int(round(w * s)), int(round(h * s))), interpolation=cv2.INTER_AREA)
+    def _preprocess(
+        self, image: ndarray
+    ) -> tuple[np.ndarray, float, tuple[float, float], tuple[int, int]]:
         orig_h, orig_w = image.shape[:2]
-        img, ratio, pad = self._letterbox(image, (self.input_w, self.input_h))
+        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.input_np_dtype)
+        img = np.ascontiguousarray(img, dtype=self.input_dtype)
         return img, ratio, pad, (orig_w, orig_h)
 
-    # ---- forward + decode (no per-class conf, no NMS — both happen in WBF/postproc) ----
-    def _forward_decode(self, image: ndarray, target_long: int | None) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-        if image.dtype != np.uint8:
-            image = image.astype(np.uint8)
-        tensor, ratio, pad, (ow, oh) = self._preproc(image, target_long)
-        outs = self.session.run(self.output_names, {self.input_name: tensor})
-        preds = outs[0]
-        if preds.ndim == 3 and preds.shape[0] == 1:
-            preds = preds[0]
-        if preds.ndim != 2 or preds.shape[1] < 6:
-            return (np.empty((0, 4), np.float32), np.empty((0,), np.float32), np.empty((0,), np.int32))
-
-        boxes = preds[:, :4].astype(np.float32)
-        scores = preds[:, 4].astype(np.float32)
-        cls_ids = preds[:, 5].astype(np.int32)
-
-        # de-letterbox back to ORIGINAL image coords (the un-resized one passed in to predict)
-        # NB: target_long step was a downsample of the input image; we de-letterbox to that
-        # downsampled size, then upscale to the true original.
-        keep = scores >= WBF_SKIP_BOX_THR
-        boxes = boxes[keep]; scores = scores[keep]; cls_ids = cls_ids[keep]
-        if len(boxes) == 0:
-            return boxes, scores, cls_ids
-
-        boxes[:, [0, 2]] -= pad[0]
-        boxes[:, [1, 3]] -= pad[1]
-        boxes /= ratio  # now in coords of the (possibly downsampled) image fed to letterbox
-
-        # If we had pre-resized for zoom view, scale boxes back to original image coords
-        true_h, true_w = image.shape[:2]
-        if target_long is not None and max(true_h, true_w) != target_long:
-            s = max(true_h, true_w) / target_long
-            boxes *= s
-
-        boxes[:, [0, 2]] = np.clip(boxes[:, [0, 2]], 0, true_w - 1)
-        boxes[:, [1, 3]] = np.clip(boxes[:, [1, 3]], 0, true_h - 1)
-        return boxes.astype(np.float32), scores, cls_ids
+    # ---------------------------------------------------------------- helpers
+    @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 _flip_boxes_h(boxes: np.ndarray, w: int) -> np.ndarray:
-        if len(boxes) == 0:
-            return boxes
-        out = boxes.copy()
-        x1 = out[:, 0].copy(); x2 = out[:, 2].copy()
-        out[:, 0] = w - x2
-        out[:, 2] = w - x1
+    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
 
-    # ---- box sanity ----
-    def _filter_sane(self, boxes, scores, cls_ids, true_size):
-        if len(boxes) == 0:
-            return boxes, scores, cls_ids
-        tw, th = true_size
-        area_img = float(tw * th)
-        keep = []
-        for i, (x1, y1, x2, y2) in enumerate(boxes):
-            bw = x2 - x1; bh = y2 - y1
-            if bw < self.min_w or bh < self.min_h:
-                continue
-            area = bw * bh
-            if area < self.min_box_area or area > self.max_box_area_ratio * area_img:
-                continue
-            ar = max(bw / max(bh, 1e-6), bh / max(bw, 1e-6))
-            if ar > self.max_aspect_ratio:
-                continue
-            keep.append(i)
-        if not keep:
-            return (np.empty((0, 4), np.float32), np.empty((0,), np.float32), np.empty((0,), np.int32))
-        keep = np.asarray(keep, dtype=np.intp)
-        return boxes[keep], scores[keep], cls_ids[keep]
-
-    # ---- per-class NMS ----
     @staticmethod
-    def _hard_nms(boxes, scores, iou_thr):
-        if len(boxes) == 0:
-            return np.array([], dtype=np.intp)
-        order = np.argsort(scores)[::-1]
-        keep = []
-        while len(order) > 0:
-            i = order[0]; keep.append(i)
-            if len(order) == 1:
+    def _iou_matrix(a: np.ndarray, b: np.ndarray) -> np.ndarray:
+        if len(a) == 0 or len(b) == 0:
+            return np.zeros((len(a), len(b)), dtype=np.float32)
+        ax1, ay1, ax2, ay2 = a[:, 0:1], a[:, 1:2], a[:, 2:3], a[:, 3:4]
+        bx1, by1, bx2, by2 = b[:, 0], b[:, 1], b[:, 2], b[:, 3]
+        ix1 = np.maximum(ax1, bx1)
+        iy1 = np.maximum(ay1, by1)
+        ix2 = np.minimum(ax2, bx2)
+        iy2 = np.minimum(ay2, by2)
+        inter = np.maximum(0.0, ix2 - ix1) * np.maximum(0.0, iy2 - iy1)
+        area_a = np.maximum(0.0, ax2 - ax1) * np.maximum(0.0, ay2 - ay1)
+        area_b = np.maximum(0.0, bx2 - bx1) * np.maximum(0.0, by2 - by1)
+        union = area_a + area_b - inter + 1e-7
+        return (inter / union).astype(np.float32)
+
+    # ---------------------------------------------------------------- NMS
+    def _soft_nms(
+        self, boxes: np.ndarray, scores: np.ndarray, sigma: float,
+        score_thresh: float = 0.001,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        n = len(boxes)
+        if n == 0:
+            return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
+        boxes = boxes.astype(np.float32, copy=True)
+        scores = scores.astype(np.float32, copy=True)
+        order = np.arange(n)
+        for i in range(n):
+            max_pos = i + int(np.argmax(scores[i:]))
+            boxes[[i, max_pos]] = boxes[[max_pos, i]]
+            scores[[i, max_pos]] = scores[[max_pos, i]]
+            order[[i, max_pos]] = order[[max_pos, i]]
+            if i + 1 >= n:
                 break
-            rest = order[1:]
-            xx1 = np.maximum(boxes[i, 0], boxes[rest, 0]); yy1 = np.maximum(boxes[i, 1], boxes[rest, 1])
-            xx2 = np.minimum(boxes[i, 2], boxes[rest, 2]); yy2 = np.minimum(boxes[i, 3], boxes[rest, 3])
-            inter = np.maximum(0, xx2 - xx1) * np.maximum(0, yy2 - yy1)
-            ai = max(0, boxes[i, 2] - boxes[i, 0]) * max(0, boxes[i, 3] - boxes[i, 1])
-            ar = np.maximum(0, boxes[rest, 2] - boxes[rest, 0]) * np.maximum(0, boxes[rest, 3] - boxes[rest, 1])
-            iou = inter / (ai + ar - inter + 1e-7)
-            order = rest[iou <= iou_thr]
-        return np.array(keep, dtype=np.intp)
-
-    @staticmethod
-    def _merge_close_graffiti(boxes, scores, cls_ids, gap_frac=0.25):
-        """Agglomerative proximity-merge for graffiti only.
-
-        Graffiti tags on the same wall often appear as multiple disjoint boxes with
-        near-zero pairwise IoU, so standard per-class NMS doesn't suppress them. King's
-        model produces one large box per wall panel; ours produces many small ones per
-        tag. We iteratively merge pairs of graffiti boxes whose nearest-edge gap is
-        smaller than `gap_frac` × the smaller box's max-side (already-overlapping
-        pairs trivially satisfy this, so they merge too).
-        """
+            xx1 = np.maximum(boxes[i, 0], boxes[i + 1:, 0])
+            yy1 = np.maximum(boxes[i, 1], boxes[i + 1:, 1])
+            xx2 = np.minimum(boxes[i, 2], boxes[i + 1:, 2])
+            yy2 = np.minimum(boxes[i, 3], boxes[i + 1:, 3])
+            inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
+            a_i = max(0.0, float(
+                (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1])
+            ))
+            a_j = (
+                np.maximum(0.0, boxes[i + 1:, 2] - boxes[i + 1:, 0]) *
+                np.maximum(0.0, boxes[i + 1:, 3] - boxes[i + 1:, 1])
+            )
+            iou = inter / (a_i + a_j - inter + 1e-7)
+            scores[i + 1:] *= np.exp(-(iou ** 2) / sigma)
+        mask = scores > score_thresh
+        return order[mask], scores[mask]
+
+    def _per_class_soft_nms(
+        self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray,
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
         if len(boxes) == 0:
             return boxes, scores, cls_ids
-        mask = cls_ids == 5
-        if mask.sum() < 2:
-            return boxes, scores, cls_ids
-        g_boxes = boxes[mask].tolist()
-        g_scores = scores[mask].tolist()
-        # leave non-graffiti alone
-        other_b = boxes[~mask]; other_s = scores[~mask]; other_c = cls_ids[~mask]
-
-        changed = True
-        while changed and len(g_boxes) > 1:
-            changed = False
-            for i in range(len(g_boxes)):
-                bi = g_boxes[i]
-                for j in range(i + 1, len(g_boxes)):
-                    bj = g_boxes[j]
-                    gap_x = max(bj[0] - bi[2], bi[0] - bj[2], 0.0)
-                    gap_y = max(bj[1] - bi[3], bi[1] - bj[3], 0.0)
-                    gap = max(gap_x, gap_y)
-                    ref = min(
-                        max(bi[2] - bi[0], bi[3] - bi[1]),
-                        max(bj[2] - bj[0], bj[3] - bj[1]),
-                    )
-                    if gap < gap_frac * ref:
-                        union = [
-                            min(bi[0], bj[0]), min(bi[1], bj[1]),
-                            max(bi[2], bj[2]), max(bi[3], bj[3]),
-                        ]
-                        new_score = max(g_scores[i], g_scores[j])
-                        g_boxes[j] = union
-                        g_scores[j] = new_score
-                        g_boxes.pop(i); g_scores.pop(i)
-                        changed = True
-                        break
-                if changed:
-                    break
-        if not g_boxes:
-            return other_b, other_s, other_c
-        merged_g = np.array(g_boxes, dtype=np.float32)
-        merged_s = np.array(g_scores, dtype=np.float32)
-        merged_c = np.full(len(g_boxes), 5, dtype=np.int32)
-        if len(other_b) == 0:
-            return merged_g, merged_s, merged_c
-        return (np.concatenate([other_b, merged_g], axis=0),
-                np.concatenate([other_s, merged_s], axis=0),
-                np.concatenate([other_c, merged_c], axis=0))
-
-    def _nms_per_class(self, boxes, scores, cls_ids):
+        out_b: list = []
+        out_s: list = []
+        out_c: list = []
+        for c in np.unique(cls_ids):
+            mask = cls_ids == c
+            sub_b = boxes[mask]
+            sub_s = scores[mask]
+            idx, decayed = self._soft_nms(sub_b, sub_s, self.soft_sigma)
+            if len(idx) == 0:
+                continue
+            out_b.append(sub_b[idx])
+            out_s.append(decayed)
+            out_c.append(np.full(len(idx), c, dtype=cls_ids.dtype))
+        if not out_b:
+            return (np.empty((0, 4), dtype=np.float32),
+                    np.empty((0,), dtype=np.float32),
+                    np.empty((0,), dtype=cls_ids.dtype))
+        return (np.concatenate(out_b, axis=0),
+                np.concatenate(out_s, axis=0),
+                np.concatenate(out_c, axis=0))
+
+    # ---------------------------------------------------------------- WBF
+    def _weighted_box_fusion(
+        self,
+        boxes: np.ndarray,
+        scores: np.ndarray,
+        cls_ids: np.ndarray,
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """Per-class confidence-weighted box fusion across orig+flip detections."""
         if len(boxes) == 0:
             return boxes, scores, cls_ids
-        keep_all = []
+        fused_b: list = []
+        fused_s: list = []
+        fused_c: list = []
         for c in np.unique(cls_ids):
-            idxs = np.nonzero(cls_ids == c)[0]
-            local = self._hard_nms(boxes[idxs], scores[idxs], self.iou_thres)
-            keep_all.extend(idxs[local].tolist())
-        keep_all = np.asarray(keep_all, dtype=np.intp)
-        order = np.argsort(scores[keep_all])[::-1][: self.max_det]
-        keep_all = keep_all[order]
-        return boxes[keep_all], scores[keep_all], cls_ids[keep_all]
-
-    # ---- main predict ----
-    def _predict_one(self, image: ndarray) -> list[BoundingBox]:
-        H, W = image.shape[:2]
-
-        # View 1: original at native resolution
-        b1, s1, c1 = self._forward_decode(image, target_long=None)
-
-        if not self.use_tta:
-            views_b, views_s, views_c = [b1], [s1], [c1]
-            view_weights = VIEW_WEIGHTS[:1]
-        else:
-            # View 2: hflip
-            flipped = cv2.flip(image, 1)
-            b2, s2, c2 = self._forward_decode(flipped, target_long=None)
-            b2 = self._flip_boxes_h(b2, W)
-
-            # View 3: zoomed (longer side resized to self.zoom_long_side, then letterbox to model input)
-            b3, s3, c3 = self._forward_decode(image, target_long=self.zoom_long_side)
-
-            views_b = [b1, b2, b3]
-            views_s = [s1, s2, s3]
-            views_c = [c1, c2, c3]
-            view_weights = VIEW_WEIGHTS
-
-        boxes, scores, cls_ids = weighted_box_fusion(
-            views_b, views_s, views_c, view_weights, iou_thr=WBF_IOU_THR
+            mask = cls_ids == c
+            sub_b = boxes[mask].astype(np.float32)
+            sub_s = scores[mask].astype(np.float32)
+            iou_thr = float(self._iou_thres_array[int(c)])
+            order = np.argsort(-sub_s)
+            sub_b = sub_b[order]
+            sub_s = sub_s[order]
+            used = np.zeros(len(sub_b), dtype=bool)
+            for i in range(len(sub_b)):
+                if used[i]:
+                    continue
+                used[i] = True
+                cluster_b = [sub_b[i]]
+                cluster_s = [sub_s[i]]
+                if i + 1 < len(sub_b):
+                    rest = sub_b[i + 1:]
+                    ious = self._iou_matrix(sub_b[i:i + 1], rest)[0]
+                    for j_offset, iou in enumerate(ious):
+                        j = i + 1 + j_offset
+                        if used[j]:
+                            continue
+                        if iou >= iou_thr:
+                            used[j] = True
+                            cluster_b.append(sub_b[j])
+                            cluster_s.append(sub_s[j])
+                ws = np.asarray(cluster_s, dtype=np.float32)
+                bs = np.asarray(cluster_b, dtype=np.float32)
+                w_sum = float(ws.sum())
+                if w_sum <= 0:
+                    continue
+                fused_box = (bs * ws[:, None]).sum(axis=0) / w_sum
+                # cluster confidence: max member, slight boost when >1 supporter
+                support = len(cluster_s)
+                fused_conf = float(ws.max())
+                if support > 1:
+                    fused_conf = min(1.0, fused_conf * (1.0 + 0.10 * (support - 1)))
+                fused_b.append(fused_box)
+                fused_s.append(fused_conf)
+                fused_c.append(int(c))
+        if not fused_b:
+            return (np.empty((0, 4), dtype=np.float32),
+                    np.empty((0,), dtype=np.float32),
+                    np.empty((0,), dtype=cls_ids.dtype))
+        return (
+            np.asarray(fused_b, dtype=np.float32),
+            np.asarray(fused_s, dtype=np.float32),
+            np.asarray(fused_c, dtype=cls_ids.dtype),
         )
-        if len(boxes) == 0:
-            return []
 
-        # Per-class confidence threshold
-        keep = scores >= self.per_class_conf[cls_ids]
-        boxes = boxes[keep]; scores = scores[keep]; cls_ids = cls_ids[keep]
+    # ---------------------------------------------------------------- sanity
+    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]:
         if len(boxes) == 0:
-            return []
+            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.where(
+            (bw > 0) & (bh > 0),
+            np.maximum(bw / np.maximum(bh, 1e-6), bh / np.maximum(bw, 1e-6)),
+            np.inf,
+        )
+        keep = (area > 0) & (area <= 0.95 * image_area) & (ar <= self.max_aspect_ratio)
+        return boxes[keep], scores[keep], cls_ids[keep]
 
-        # Sanity + NMS
-        boxes, scores, cls_ids = self._filter_sane(boxes, scores, cls_ids, (W, H))
+    # ---------------------------------------------------------------- decode
+    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]:
+        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(boxes) > 1:
+            boxes, scores, cls_ids = self._per_class_soft_nms(boxes, scores, cls_ids)
+        if len(scores) > self.max_det:
+            top = np.argsort(-scores)[: self.max_det]
+            boxes, scores, cls_ids = boxes[top], scores[top], cls_ids[top]
+        return boxes, scores, cls_ids
+
+    def _decode_final_dets(
+        self, preds: np.ndarray, ratio: float, pad: tuple[float, float],
+        orig_size: tuple[int, int],
+    ) -> list[BoundingBox]:
+        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 = preds[:, 5].astype(np.int32)
+        keep = scores >= self._conf_thres_array[cls_ids]
+        boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep]
         if len(boxes) == 0:
             return []
-        boxes, scores, cls_ids = self._nms_per_class(boxes, scores, cls_ids)
-        # Graffiti-only proximity merge: collapse tiny adjacent tags into wall-panel boxes,
-        # which matches the validator's preferred granularity (one box per panel).
-        boxes, scores, cls_ids = self._merge_close_graffiti(boxes, scores, cls_ids, gap_frac=0.25)
-
-        # Per-class aspect-ratio filter (currently: hoodie must be tall ≥1.2× wide,
-        # so we keep only detections that include the hooded-head region).
-        if len(boxes) > 0 and PER_CLASS_MIN_H_OVER_W:
-            keep = np.ones(len(boxes), dtype=bool)
-            for i in range(len(boxes)):
-                c = int(cls_ids[i])
-                if c in PER_CLASS_MIN_H_OVER_W:
-                    bw = float(boxes[i, 2] - boxes[i, 0])
-                    bh = float(boxes[i, 3] - boxes[i, 1])
-                    if bw <= 0 or (bh / bw) < PER_CLASS_MIN_H_OVER_W[c]:
-                        keep[i] = False
-            boxes = boxes[keep]; scores = scores[keep]; cls_ids = cls_ids[keep]
-
-        return [
+        pad_w, pad_h = pad
+        boxes[:, [0, 2]] -= pad_w
+        boxes[:, [1, 3]] -= pad_h
+        boxes /= ratio
+        boxes = self._clip_boxes(boxes, orig_size)
+        boxes, scores, cls_ids = self._per_view_pipeline(boxes, scores, cls_ids, orig_size)
+        return self._build_results(boxes, scores, cls_ids)
+
+    def _decode_raw_yolo(
+        self, preds: np.ndarray, ratio: float, pad: tuple[float, float],
+        orig_size: tuple[int, int],
+    ) -> list[BoundingBox]:
+        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 raw output 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]
+        keep = scores >= self._conf_thres_array[cls_ids]
+        boxes_xywh, scores, cls_ids = boxes_xywh[keep], scores[keep], cls_ids[keep]
+        if len(boxes_xywh) == 0:
+            return []
+        boxes = self._xywh_to_xyxy(boxes_xywh)
+        pad_w, pad_h = pad
+        boxes[:, [0, 2]] -= pad_w
+        boxes[:, [1, 3]] -= pad_h
+        boxes /= ratio
+        boxes = self._clip_boxes(boxes, orig_size)
+        boxes, scores, cls_ids = self._per_view_pipeline(boxes, scores, cls_ids, orig_size)
+        return self._build_results(boxes, scores, cls_ids)
+
+    @staticmethod
+    def _build_results(
+        boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray,
+    ) -> list[BoundingBox]:
+        results: list[BoundingBox] = []
+        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=int(math.floor(x1)),
+                    y1=int(math.floor(y1)),
+                    x2=int(math.ceil(x2)),
+                    y2=int(math.ceil(y2)),
+                    cls_id=int(cls_id),
+                    conf=float(conf),
+                )
+            )
+        return results
+
+    def _postprocess(
+        self, output: np.ndarray, ratio: float, pad: tuple[float, float],
+        orig_size: tuple[int, int],
+    ) -> list[BoundingBox]:
+        if output.ndim == 2 and output.shape[1] >= 6:
+            return self._decode_final_dets(output, ratio, pad, orig_size)
+        if output.ndim == 3 and output.shape[0] == 1 and output.shape[2] == 6:
+            return self._decode_final_dets(output, ratio, pad, orig_size)
+        return self._decode_raw_yolo(output, ratio, pad, orig_size)
+
+    # ---------------------------------------------------------------- inference
+    def _predict_single(self, image: np.ndarray) -> list[BoundingBox]:
+        if image is None:
+            raise ValueError("Input image is None")
+        if not isinstance(image, np.ndarray):
+            raise TypeError(f"Input is not numpy array: {type(image)}")
+        if image.ndim != 3:
+            raise ValueError(f"Expected HWC image, got shape={image.shape}")
+        if image.shape[2] != 3:
+            raise ValueError(f"Expected 3 channels, got shape={image.shape}")
+        if image.dtype != np.uint8:
+            image = image.astype(np.uint8)
+
+        input_tensor, ratio, pad, orig_size = self._preprocess(image)
+        expected = (1, 3, self.input_height, self.input_width)
+        if input_tensor.shape != expected:
+            raise ValueError(
+                f"Bad input tensor shape={input_tensor.shape}, expected={expected}"
+            )
+        outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
+        return self._postprocess(outputs[0], ratio, pad, orig_size)
+
+    def _predict_tta(self, image: np.ndarray) -> list[BoundingBox]:
+        boxes_orig = self._predict_single(image)
+        flipped = cv2.flip(image, 1)
+        boxes_flip = self._predict_single(flipped)
+        w = image.shape[1]
+        boxes_flip = [
             BoundingBox(
-                x1=int(math.floor(b[0])), y1=int(math.floor(b[1])),
-                x2=int(math.ceil(b[2])), y2=int(math.ceil(b[3])),
-                cls_id=int(c), conf=float(s),
+                x1=w - b.x2, y1=b.y1, x2=w - b.x1, y2=b.y2,
+                cls_id=b.cls_id, conf=b.conf,
             )
-            for b, s, c in zip(boxes, scores, cls_ids)
-            if b[2] > b[0] and b[3] > b[1]
+            for b in boxes_flip
         ]
+        all_boxes = boxes_orig + boxes_flip
+        if not all_boxes:
+            return []
+        coords = np.array(
+            [[b.x1, b.y1, b.x2, b.y2] for b in all_boxes], dtype=np.float32
+        )
+        scores = np.array([b.conf for b in all_boxes], dtype=np.float32)
+        cls_ids = np.array([b.cls_id for b in all_boxes], dtype=np.int32)
+
+        # Per-class weighted box fusion (replaces hard-NMS + cluster-boost)
+        boxes_f, scores_f, cls_f = self._weighted_box_fusion(coords, scores, cls_ids)
+        if len(boxes_f) == 0:
+            return []
+        if len(scores_f) > self.max_det:
+            top = np.argsort(-scores_f)[: self.max_det]
+            boxes_f, scores_f, cls_f = boxes_f[top], scores_f[top], cls_f[top]
+        return self._build_results(boxes_f, scores_f, cls_f)
 
     def predict_batch(
-        self,
-        batch_images: list[ndarray],
-        offset: int,
-        n_keypoints: int,
+        self, batch_images: list[ndarray], offset: int, n_keypoints: int,
     ) -> list[TVFrameResult]:
         results: list[TVFrameResult] = []
-        for i, image in enumerate(batch_images):
+        for frame_number_in_batch, image in enumerate(batch_images):
             try:
-                boxes = self._predict_one(image)
+                boxes = self._predict_tta(image)
             except Exception as e:
-                print(f"⚠️ inference failed frame {offset + i}: {e}")
+                print(
+                    f"Inference failed for frame "
+                    f"{offset + frame_number_in_batch}: {e}"
+                )
                 boxes = []
-            results.append(TVFrameResult(
-                frame_id=offset + i,
-                boxes=boxes,
-                keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))],
-            ))
+            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