scorevision: push artifact

miner.py

@@ -0,0 +1,461 @@
+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]]
+
+
+class Miner:
+    """
+    YOLO (NMS-free) ONNX miner for road-sign detection. One forward pass per
+    frame — no TTA.
+
+    Class: single 'road sign'.
+
+    Pipeline per frame: preprocess -> ONNX -> decode -> conf threshold (+ rescue
+    bonus) -> un-letterbox -> sanity filter -> per-class NMS -> cross-class dedup
+    -> same-class cluster score boost -> results.
+
+    Speed: the detection pipeline runs exactly ONCE per frame (the old redundant
+    second NMS/dedup pass is gone), and the cluster boost is a single vectorized
+    IoU matrix, so cost stays flat as the number of detected objects grows.
+    """
+
+    class_names = ['road sign']
+    input_size = 1536
+    cross_iou_thresh = 0.8
+    max_det = 300
+    # NMS is O(n^2); if a frame yields a huge candidate list, keep the top-K by
+    # score before NMS. Set high enough never to touch real detections.
+    pre_nms_topk = 1000
+    #overlap_suppress_threshold = 0.85
+
+    # Per-class confidence threshold
+    _conf_thres_array = np.array([0.31], dtype=np.float32)
+
+    # Per-class IoU threshold for same-class NMS
+    _iou_thres_array = np.array([0.85], dtype=np.float32)
+
+    # Per-class rescue bonus
+    _bonus_array = np.array([0.2], dtype=np.float32)
+
+    # Per-class minimum box area (index 0 = road sign)
+    _min_box_area_array = np.array([9.0], dtype=np.float32)
+
+    def __init__(self, path_hf_repo: Path) -> None:
+        self.path_hf_repo = path_hf_repo
+
+        print("ORT version:", ort.__version__)
+
+        try:
+            ort.preload_dlls()
+            print("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
+
+        self.session = ort.InferenceSession(
+            str(path_hf_repo / "weights.onnx"),
+            sess_options=sess_options,
+            providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
+        )
+        print("Created ORT session with preferred CUDA provider list")
+        print("ORT session providers:", self.session.get_providers())
+
+        self.input_name = self.session.get_inputs()[0].name
+        input_shape = self.session.get_inputs()[0].shape
+
+        self.input_h = self._safe_dim(input_shape[2], default=self.input_size)
+        self.input_w = self._safe_dim(input_shape[3], default=self.input_size)
+
+        # Same-class cluster score boost (raises overlapping same-class survivors
+        # to their cluster max). Part of the current tuned behaviour; for a single
+        # class it is nearly a no-op. Set False to disable.
+        self.use_cluster_boost = True
+        self._avg_iou = float(np.mean(self._iou_thres_array))
+
+        self._warmup()
+
+    def _warmup(self, iters: int = 3) -> None:
+        try:
+            dummy = np.zeros((720, 1280, 3), dtype=np.uint8)
+            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"road sign Miner classes={len(self.class_names)}"
+
+    @staticmethod
+    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],
+        color: tuple[int, int, int] = (114, 114, 114),
+    ) -> tuple[ndarray, float, float, float]:
+        orig_h, orig_w = image.shape[:2]
+        target_w, target_h = new_shape
+
+        r = min(target_w / orig_w, target_h / orig_h)
+        new_unpad_w = int(round(orig_w * r))
+        new_unpad_h = int(round(orig_h * r))
+
+        resized = cv2.resize(image, (new_unpad_w, new_unpad_h), interpolation=cv2.INTER_LINEAR)
+
+        dw = target_w - new_unpad_w
+        dh = target_h - new_unpad_h
+        pad_w = dw / 2.0
+        pad_h = dh / 2.0
+
+        left = int(round(pad_w - 0.1))
+        right = int(round(pad_w + 0.1))
+        top = int(round(pad_h - 0.1))
+        bottom = int(round(pad_h + 0.1))
+
+        out = cv2.copyMakeBorder(
+            resized, top, bottom, left, right,
+            cv2.BORDER_CONSTANT, value=color,
+        )
+        return out, r, pad_w, pad_h
+
+    def _preprocess(self, image_bgr: ndarray) -> tuple[np.ndarray, dict]:
+        orig_h, orig_w = image_bgr.shape[:2]
+        rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
+        img, ratio, pad_w, pad_h = self._letterbox(rgb, (self.input_w, self.input_h))
+        x = img.astype(np.float32) / 255.0
+        x = np.transpose(x, (2, 0, 1))[None, ...]
+        x = np.ascontiguousarray(x)
+        return x, {
+            "orig_h": orig_h,
+            "orig_w": orig_w,
+            "ratio": ratio,
+            "pad_w": pad_w,
+            "pad_h": pad_h,
+        }
+
+    # ─── Vectorized box operations ───────────────────────────────
+
+    @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 _hard_nms(boxes: np.ndarray, scores: np.ndarray,
+                  iou_thresh: float) -> np.ndarray:
+        """Vectorized greedy NMS. Areas precomputed once. Returns indices to keep."""
+        n = len(boxes)
+        if n == 0:
+            return np.array([], dtype=np.intp)
+        x1, y1 = boxes[:, 0], boxes[:, 1]
+        x2, y2 = 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)
+
+    def _per_class_hard_nms(self, boxes: np.ndarray, scores: np.ndarray,
+                            cls_ids: np.ndarray) -> np.ndarray:
+        """Per-class NMS using per-class IoU thresholds."""
+        if len(boxes) == 0:
+            return np.array([], dtype=np.intp)
+        all_keep = []
+        for c in np.unique(cls_ids):
+            mask = cls_ids == c
+            indices = np.where(mask)[0]
+            cls_iou = float(self._iou_thres_array[c])  # per-class IoU threshold
+            keep = self._hard_nms(boxes[mask], scores[mask], cls_iou)
+            all_keep.extend(indices[keep].tolist())
+        all_keep.sort()
+        return np.array(all_keep, dtype=np.intp)
+
+    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]:
+        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 = []
+        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 _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]:
+        """Filter by per-class min area and max area ratio."""
+        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
+
+        class_min_area = self._min_box_area_array[cls_ids]
+
+        keep = (
+            (area >= class_min_area) &
+            (area <= 0.95 * image_area)
+        )
+        return boxes[keep], scores[keep], cls_ids[keep]
+
+    def _max_score_per_cluster(self, 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:
+        """For each kept box, confidence = max score in its SAME-CLASS IoU cluster.
+        Vectorized: single (n_post x n_full) IoU matrix, no per-box Python loop."""
+        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) -> np.ndarray:
+        """Per-class threshold with rescue bonus for missed classes."""
+        if len(scores) == 0:
+            return np.zeros(0, dtype=bool)
+        thr = self._conf_thres_array[cls_ids]
+        keep = scores >= thr
+        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] >= self._conf_thres_array[c] - b:
+                keep[top] = True
+        return keep
+
+    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]:
+        """Sanity filter -> (top-k cap) -> per-class NMS -> 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:
+            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
+
+    # ─── Decoding ─────────────────────────────────────────────────
+
+    def _decode_yolo_output(self, preds: np.ndarray, ratio: float,
+                           pad: tuple[float, float],
+                           orig_size: tuple[int, int]
+                           ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """Decode NMS-free output [1,N,6]=(x1,y1,x2,y2,conf,cls) -> conf filter ->
+        un-letterbox -> sanity+NMS+dedup. Returns arrays in ORIGINAL coords."""
+        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:
+            print(f"Warning: Unexpected output shape: {preds.shape}")
+            return empty
+
+        boxes = preds[:, :4].astype(np.float32)
+        scores = preds[:, 4].astype(np.float32)
+        cls_ids = preds[:, 5].astype(np.int32)
+
+        n_cls = len(self.class_names)
+        valid = (cls_ids >= 0) & (cls_ids < n_cls)
+        boxes, scores, cls_ids = boxes[valid], scores[valid], cls_ids[valid]
+        if len(boxes) == 0:
+            return empty
+
+        keep = self._conf_filter_mask(scores, cls_ids)
+        boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep]
+        if len(boxes) == 0:
+            return empty
+
+        pad_w, pad_h = pad
+        boxes[:, [0, 2]] -= pad_w
+        boxes[:, [1, 3]] -= pad_h
+        boxes /= ratio
+        boxes = self._clip_boxes(boxes, orig_size)
+
+        return self._per_view_pipeline(boxes, scores, cls_ids, orig_size)
+
+    @staticmethod
+    def _build_results(boxes: np.ndarray, scores: np.ndarray,
+                       cls_ids: np.ndarray,
+                       orig_size: tuple[int, int]) -> list[BoundingBox]:
+        results = []
+        orig_w, orig_h = orig_size
+        for box, conf, cls_id in zip(boxes, scores, cls_ids):
+            x1, y1, x2, y2 = box.tolist() if hasattr(box, "tolist") else box
+            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
+
+    # ─── Inference (single view, no TTA) ──────────────────────────
+
+    def _predict_single(self, image_bgr: np.ndarray
+                        ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """One forward pass -> decoded (boxes, scores, cls_ids) in original coords."""
+        if image_bgr is None or not isinstance(image_bgr, np.ndarray):
+            raise ValueError("Invalid image input")
+        if image_bgr.dtype != np.uint8:
+            image_bgr = image_bgr.astype(np.uint8)
+
+        inp, meta = self._preprocess(image_bgr)
+        outputs = self.session.run(None, {self.input_name: inp})
+
+        ratio = float(meta["ratio"])
+        pad = (float(meta["pad_w"]), float(meta["pad_h"]))
+        orig_size = (int(meta["orig_w"]), int(meta["orig_h"]))
+
+        return self._decode_yolo_output(outputs[0], ratio, pad, orig_size)
+
+    def _infer_single(self, image_bgr: ndarray) -> list[BoundingBox]:
+        orig_h, orig_w = image_bgr.shape[:2]
+        orig_size = (orig_w, orig_h)
+
+        boxes, scores, cls_ids = self._predict_single(image_bgr)
+        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)
+
+        return self._build_results(boxes, scores, cls_ids, orig_size)
+
+    # ─── Public API ───────────────────────────────────────────────
+
+    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_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