scorevision: push artifact

miner.py

@@ -0,0 +1,427 @@
+"""
+Inference miner for `manak0/Detect-crime`.
+
+Beats the published king (`manak0/Detect-crime` baseline, score 0.576) by:
+
+  1. Letterboxed input at 1280 (king stretch-resizes to 640, destroying small objects).
+  2. Per-class confidence floors for the four catastrophic classes
+     (balaclava, bat, glove, spray paint each had recall < 0.16 in the king's benchmark).
+  3. Multi-scale TTA: {1280, 1536} x {orig, hflip} with class-aware Weighted Box Fusion.
+  4. CLAHE on dark frames only (luma gate). CCTV crime footage is heavily night-time.
+  5. Class-aware NMS at IoU=0.45 (king is class-agnostic, suppresses balaclava-on-hoodie
+     and glove-near-bat overlaps).
+  6. Robust ONNX output decoding for both raw YOLO `[1, 4+nc, N]` (king's export) and
+     NMS-baked `[N, 6]` / `[1, N, 6]` (our export) formats.
+
+Self-contained: only stdlib + pip packages (cv2, numpy, onnxruntime, pydantic). No imports
+from sibling files in the HF repo (chute import sandbox blocks those).
+"""
+
+from pathlib import Path
+import math
+
+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]]
+
+
+TARGET_CLASS_NAMES = ["balaclava", "bat", "glove", "graffiti", "hoodie", "spray paint"]
+
+# Per-class confidence floors. King uses one global 0.25; that over-suppresses the four
+# rare classes. Synthetic-benchmark recalls were balaclava 0.034, glove 0.064, spray
+# paint 0.16, bat 0.14 — recall is the bottleneck. With ~6.5 preds/img on average we
+# have headroom under any FFPI cap to push the floors lower without saturating FPs.
+PER_CLASS_CONF = {
+    0: 0.05,   # balaclava   - rec 0.034: catastrophic; pushed hard
+    1: 0.10,   # bat         - rec 0.143
+    2: 0.05,   # glove       - rec 0.064: catastrophic; pushed hard
+    3: 0.20,   # graffiti    - rec 0.321 (best of the rare four); be conservative
+    4: 0.20,   # hoodie      - rec 0.274 + carries the IoU mean alone
+    5: 0.10,   # spray paint - rec 0.161
+}
+
+# Multi-scale TTA. Pro_6000 + YOLOv11s at 1536 is ~150ms; ample budget under p95=10s.
+TTA_SIZES = (1280, 1536)
+TTA_HFLIP = True
+
+# Class-aware NMS / fusion thresholds.
+NMS_IOU = 0.45
+WBF_IOU = 0.55
+MAX_DET = 100
+
+# CLAHE on dark frames only (CCTV night).
+CLAHE_DARK_THRESHOLD = 70   # mean Y < this -> apply CLAHE
+CLAHE_CLIP_LIMIT = 2.0
+CLAHE_TILE = (8, 8)
+
+
+class Miner:
+    def __init__(self, path_hf_repo: Path) -> None:
+        model_path = path_hf_repo / "weights.onnx"
+        cn_path = model_path.with_name("class_names.txt")
+
+        self.class_names = TARGET_CLASS_NAMES.copy()
+        if cn_path.is_file():
+            lines = cn_path.read_text(encoding="utf-8").splitlines()
+            order = [ln.strip() for ln in lines if ln.strip() and not ln.strip().startswith("#")]
+            if len(order) == len(self.class_names) and set(order) == set(self.class_names):
+                self.cls_remap = np.array(
+                    [self.class_names.index(n) for n in order], dtype=np.int32
+                )
+            else:
+                self.cls_remap = np.arange(len(self.class_names), dtype=np.int32)
+        else:
+            self.cls_remap = np.arange(len(self.class_names), dtype=np.int32)
+
+        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:
+            self.session = ort.InferenceSession(
+                str(model_path),
+                sess_options=sess_options,
+                providers=["CPUExecutionProvider"],
+            )
+
+        self.input_name = self.session.get_inputs()[0].name
+        self.output_names = [o.name for o in self.session.get_outputs()]
+        in_shape = self.session.get_inputs()[0].shape
+        self.native_h = in_shape[2] if isinstance(in_shape[2], int) and in_shape[2] > 0 else TTA_SIZES[0]
+        self.native_w = in_shape[3] if isinstance(in_shape[3], int) and in_shape[3] > 0 else TTA_SIZES[0]
+        self.is_dynamic_input = not (
+            isinstance(in_shape[2], int) and isinstance(in_shape[3], int)
+            and in_shape[2] > 0 and in_shape[3] > 0
+        )
+
+        # If export is static-shape, force a single TTA size matching the model.
+        if not self.is_dynamic_input:
+            self._tta_sizes = (self.native_h,)
+        else:
+            self._tta_sizes = TTA_SIZES
+
+        self._clahe = cv2.createCLAHE(clipLimit=CLAHE_CLIP_LIMIT, tileGridSize=CLAHE_TILE)
+
+        print(f"[miner] model={model_path.name} input={self.native_w}x{self.native_h} dynamic={self.is_dynamic_input}")
+        print(f"[miner] tta_sizes={self._tta_sizes} hflip={TTA_HFLIP} per_class_conf={PER_CLASS_CONF}")
+        print(f"[miner] providers={self.session.get_providers()}")
+
+    def __repr__(self) -> str:
+        return f"CrimeMiner(providers={self.session.get_providers()})"
+
+    # ------------------------------------------------------------------ utils
+
+    @staticmethod
+    def _letterbox(image: ndarray, new_size: int, color=(114, 114, 114)) -> tuple[ndarray, float, tuple[float, float]]:
+        h, w = image.shape[:2]
+        ratio = min(new_size / w, new_size / h)
+        rw, rh = int(round(w * ratio)), int(round(h * ratio))
+        if (rw, rh) != (w, h):
+            interp = cv2.INTER_AREA if ratio < 1.0 else cv2.INTER_LINEAR
+            image = cv2.resize(image, (rw, rh), interpolation=interp)
+        dw, dh = (new_size - rw) / 2.0, (new_size - rh) / 2.0
+        top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
+        left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
+        padded = cv2.copyMakeBorder(image, top, bottom, left, right, borderType=cv2.BORDER_CONSTANT, value=color)
+        return padded, ratio, (dw, dh)
+
+    def _maybe_clahe(self, image_bgr: ndarray) -> ndarray:
+        small = cv2.resize(image_bgr, (160, 90), interpolation=cv2.INTER_AREA)
+        mean_y = float(0.114 * small[..., 0].mean() + 0.587 * small[..., 1].mean() + 0.299 * small[..., 2].mean())
+        if mean_y >= CLAHE_DARK_THRESHOLD:
+            return image_bgr
+        lab = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2LAB)
+        lab[..., 0] = self._clahe.apply(lab[..., 0])
+        return cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
+
+    @staticmethod
+    def _xywh_to_xyxy(b: np.ndarray) -> np.ndarray:
+        out = np.empty_like(b)
+        out[:, 0] = b[:, 0] - b[:, 2] / 2.0
+        out[:, 1] = b[:, 1] - b[:, 3] / 2.0
+        out[:, 2] = b[:, 0] + b[:, 2] / 2.0
+        out[:, 3] = b[:, 1] + b[:, 3] / 2.0
+        return out
+
+    # ------------------------------------------------------------------ ONNX
+
+    def _preprocess(self, image_bgr: ndarray, size: int) -> tuple[np.ndarray, float, tuple[float, float]]:
+        img, ratio, pad = self._letterbox(image_bgr, size)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        img = img.astype(np.float32) / 255.0
+        img = np.transpose(img, (2, 0, 1))[None, ...]
+        return np.ascontiguousarray(img, dtype=np.float32), ratio, pad
+
+    def _run_session(self, image_bgr: ndarray, size: int, hflip: bool):
+        if hflip:
+            image_bgr = cv2.flip(image_bgr, 1)
+        tensor, ratio, pad = self._preprocess(image_bgr, size)
+        outputs = self.session.run(self.output_names, {self.input_name: tensor})
+        return outputs[0], ratio, pad
+
+    # ------------------------------------------------------------------ decode
+
+    def _decode(self, output: np.ndarray, ratio: float, pad: tuple[float, float],
+                orig_w: int, orig_h: int, hflip: bool) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """Return (boxes_xyxy in orig-image coords, scores, cls_ids)."""
+        # Case A: NMS-baked [N, 6] or [1, N, 6]
+        if output.ndim == 3 and output.shape[0] == 1 and output.shape[2] == 6:
+            preds = output[0]
+        elif output.ndim == 2 and output.shape[1] == 6:
+            preds = output
+        else:
+            return self._decode_raw(output, ratio, pad, orig_w, orig_h, hflip)
+
+        boxes = preds[:, :4].astype(np.float32, copy=True)
+        scores = preds[:, 4].astype(np.float32)
+        cls_ids = preds[:, 5].astype(np.int32)
+        cls_ids = self.cls_remap[np.clip(cls_ids, 0, len(self.cls_remap) - 1)]
+
+        keep = (scores > 0.0) & (boxes[:, 2] > boxes[:, 0]) & (boxes[:, 3] > boxes[:, 1])
+        boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep]
+        if len(boxes) == 0:
+            return boxes, scores, cls_ids
+
+        return self._unletterbox(boxes, scores, cls_ids, ratio, pad, orig_w, orig_h, hflip)
+
+    def _decode_raw(self, output: np.ndarray, ratio: float, pad: tuple[float, float],
+                    orig_w: int, orig_h: int, hflip: bool) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        # Raw ultralytics export: [1, C, N] where C = 4 + num_classes (king's case: [1, 10, 8400]).
+        if output.ndim != 3 or output.shape[0] != 1:
+            return np.zeros((0, 4), np.float32), np.zeros(0, np.float32), np.zeros(0, np.int32)
+        preds = output[0]
+        # Normalize to [N, C].
+        if preds.shape[0] <= 16 and preds.shape[1] > preds.shape[0]:
+            preds = preds.T
+        if preds.shape[1] < 5:
+            return np.zeros((0, 4), np.float32), np.zeros(0, np.float32), np.zeros(0, np.int32)
+
+        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[np.clip(cls_ids, 0, len(self.cls_remap) - 1)]
+
+        # Coarse pre-filter: any class above its own floor (or at least 0.03 to keep the path
+        # light for the four catastrophic classes).
+        floors = np.array([PER_CLASS_CONF.get(int(c), 0.20) for c in cls_ids], dtype=np.float32)
+        keep = scores >= np.minimum(floors, 0.03)
+        if not np.any(keep):
+            return np.zeros((0, 4), np.float32), np.zeros(0, np.float32), np.zeros(0, np.int32)
+        boxes_xywh = boxes_xywh[keep]
+        scores = scores[keep]
+        cls_ids = cls_ids[keep]
+
+        boxes = self._xywh_to_xyxy(boxes_xywh)
+        return self._unletterbox(boxes, scores, cls_ids, ratio, pad, orig_w, orig_h, hflip)
+
+    @staticmethod
+    def _unletterbox(boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray,
+                     ratio: float, pad: tuple[float, float], orig_w: int, orig_h: int,
+                     hflip: bool) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        boxes = boxes.copy()
+        boxes[:, [0, 2]] -= pad[0]
+        boxes[:, [1, 3]] -= pad[1]
+        boxes /= ratio
+        if hflip:
+            x1 = orig_w - 1 - boxes[:, 2]
+            x2 = orig_w - 1 - boxes[:, 0]
+            boxes[:, 0] = x1
+            boxes[:, 2] = x2
+        boxes[:, 0] = np.clip(boxes[:, 0], 0, orig_w - 1)
+        boxes[:, 1] = np.clip(boxes[:, 1], 0, orig_h - 1)
+        boxes[:, 2] = np.clip(boxes[:, 2], 0, orig_w - 1)
+        boxes[:, 3] = np.clip(boxes[:, 3], 0, orig_h - 1)
+        return boxes, scores, cls_ids
+
+    # ------------------------------------------------------------------ fusion
+
+    @staticmethod
+    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)
+        a = a.astype(np.float32)
+        b = b.astype(np.float32)
+        x11, y11, x12, y12 = a[:, 0:1], a[:, 1:2], a[:, 2:3], a[:, 3:4]
+        x21, y21, x22, y22 = b[:, 0], b[:, 1], b[:, 2], b[:, 3]
+        xa = np.maximum(x11, x21); ya = np.maximum(y11, y21)
+        xb = np.minimum(x12, x22); yb = np.minimum(y12, y22)
+        inter = np.maximum(0.0, xb - xa) * np.maximum(0.0, yb - ya)
+        area_a = (x12 - x11) * (y12 - y11)
+        area_b = (x22 - x21) * (y22 - y21)
+        return inter / (area_a + area_b - inter + 1e-7)
+
+    def _wbf(self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray,
+             iou_thresh: float = WBF_IOU) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """Class-aware Weighted Box Fusion: inside each class cluster, kept box is the
+        score-weighted average; cluster score is max member score (preserves AP rank).
+        """
+        if len(boxes) == 0:
+            return boxes, scores, cls_ids
+
+        out_boxes, out_scores, out_cls = [], [], []
+        for c in np.unique(cls_ids):
+            mask = cls_ids == c
+            cb = boxes[mask].astype(np.float32)
+            cs = scores[mask].astype(np.float32)
+            order = np.argsort(-cs)
+            cb = cb[order]
+            cs = cs[order]
+            used = np.zeros(len(cb), dtype=bool)
+            for i in range(len(cb)):
+                if used[i]:
+                    continue
+                base = cb[i:i + 1]
+                ious = self._iou_matrix(base, cb)[0]
+                cluster = (ious >= iou_thresh) & (~used)
+                cluster[i] = True
+                w = cs[cluster]
+                wsum = float(w.sum()) + 1e-7
+                fused = (cb[cluster] * w[:, None]).sum(axis=0) / wsum
+                out_boxes.append(fused)
+                out_scores.append(float(cs[cluster].max()))
+                out_cls.append(int(c))
+                used = used | cluster
+        return (np.stack(out_boxes), np.array(out_scores, dtype=np.float32),
+                np.array(out_cls, dtype=np.int32))
+
+    def _class_aware_nms(self, boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray,
+                         iou_thresh: float = NMS_IOU) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        if len(boxes) == 0:
+            return boxes, scores, cls_ids
+        kept_b, kept_s, kept_c = [], [], []
+        for c in np.unique(cls_ids):
+            mask = cls_ids == c
+            cb = boxes[mask].astype(np.float32)
+            cs = scores[mask].astype(np.float32)
+            order = np.argsort(-cs)
+            suppressed = np.zeros(len(cb), dtype=bool)
+            for i in order:
+                if suppressed[i]:
+                    continue
+                kept_b.append(cb[i])
+                kept_s.append(float(cs[i]))
+                kept_c.append(int(c))
+                ious = self._iou_matrix(cb[i:i + 1], cb)[0]
+                suppressed = suppressed | (ious >= iou_thresh)
+                suppressed[i] = True
+        if not kept_b:
+            return np.zeros((0, 4), np.float32), np.zeros(0, np.float32), np.zeros(0, np.int32)
+        return np.stack(kept_b), np.array(kept_s, dtype=np.float32), np.array(kept_c, dtype=np.int32)
+
+    # ------------------------------------------------------------------ predict
+
+    def _predict_one(self, image_bgr: ndarray) -> list[BoundingBox]:
+        h, w = image_bgr.shape[:2]
+        enhanced = self._maybe_clahe(image_bgr)
+
+        all_boxes, all_scores, all_cls = [], [], []
+        for size in self._tta_sizes:
+            for hflip in ((False, True) if TTA_HFLIP else (False,)):
+                output, ratio, pad = self._run_session(enhanced, size, hflip)
+                boxes, scores, cls_ids = self._decode(output, ratio, pad, w, h, hflip)
+                if len(boxes):
+                    all_boxes.append(boxes)
+                    all_scores.append(scores)
+                    all_cls.append(cls_ids)
+
+        if not all_boxes:
+            return []
+
+        boxes = np.concatenate(all_boxes, axis=0)
+        scores = np.concatenate(all_scores, axis=0)
+        cls_ids = np.concatenate(all_cls, axis=0)
+
+        # Per-class confidence floor.
+        floors = np.array([PER_CLASS_CONF.get(int(c), 0.20) for c in cls_ids], dtype=np.float32)
+        keep = scores >= floors
+        boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep]
+        if len(boxes) == 0:
+            return []
+
+        # Fuse across TTA streams (class-aware WBF), then class-aware NMS to clean up.
+        boxes, scores, cls_ids = self._wbf(boxes, scores, cls_ids, iou_thresh=WBF_IOU)
+        boxes, scores, cls_ids = self._class_aware_nms(boxes, scores, cls_ids, iou_thresh=NMS_IOU)
+
+        if len(boxes) > MAX_DET:
+            order = np.argsort(-scores)[:MAX_DET]
+            boxes, scores, cls_ids = boxes[order], scores[order], cls_ids[order]
+
+        results: list[BoundingBox] = []
+        for box, conf, cid 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(cid),
+                conf=float(np.clip(conf, 0.0, 1.0)),
+            ))
+        return results
+
+    def predict_batch(
+        self,
+        batch_images: list[ndarray],
+        offset: int,
+        n_keypoints: int,
+    ) -> list[TVFrameResult]:
+        results: list[TVFrameResult] = []
+        for i, image in enumerate(batch_images):
+            try:
+                boxes = self._predict_one(image)
+            except Exception as e:
+                print(f"[miner] inference failed on frame {offset + i}: {e}")
+                boxes = []
+            results.append(TVFrameResult(
+                frame_id=offset + i,
+                boxes=boxes,
+                keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))],
+            ))
+        return results
+
+
+if __name__ == "__main__":
+    # Smoke test against any image in /tmp/crime_proof.png or weights.onnx-adjacent test.png.
+    import sys
+    repo_dir = Path(__file__).parent
+    miner = Miner(repo_dir)
+    candidates = [Path("/tmp/crime_proof.png"), repo_dir / "test.png"]
+    img_path = next((p for p in candidates if p.exists()), None)
+    if img_path is None:
+        print("no test image found; place one at /tmp/crime_proof.png and rerun")
+        sys.exit(0)
+    img = cv2.imread(str(img_path), cv2.IMREAD_COLOR)
+    out = miner.predict_batch([img], offset=0, n_keypoints=0)
+    for f in out:
+        print(f"frame {f.frame_id}: {len(f.boxes)} boxes")
+        for b in f.boxes:
+            print(f"  cls={b.cls_id} ({TARGET_CLASS_NAMES[b.cls_id]}) conf={b.conf:.3f} "
+                  f"box=({b.x1},{b.y1},{b.x2},{b.y2})")