scorevision: push artifact

miner.py

@@ -7,6 +7,7 @@ import onnxruntime as ort
 from numpy import ndarray
 from pydantic import BaseModel
 
+
 class BoundingBox(BaseModel):
     x1: int
     y1: int
@@ -24,59 +25,62 @@ class TVFrameResult(BaseModel):
 
 class Miner:
     """ONNX Runtime miner for road-sign detection (single class).
-
-    Post-processing: per-class confidence + rescue bonus, hard NMS,
-    cross-class dedup, sanity-box filter, optional TTA / tile merge.
+    Strategy (ported from offense / fire001 miner):
+      - per-class confidence threshold with per-class rescue bonus
+      - per-class hard NMS, then cross-class dedup (no-op for single class)
+      - horizontal-flip TTA with full-set cluster score boost
+    Plus: class remap, sanity-box filter tuned for small distant signs,
+    TTA toggle.
     """
 
     class_names = ["road_sign"]
+    # Order the model emits classes in -- remapped to `class_names` index.
     _model_class_order = ["road_sign"]
 
     iou_thres = 0.5
     cross_iou_thresh = 0.8
     max_det = 150
 
-    use_secondary_merge = True
-    secondary_conf = 0.70
-    merge_iou = 0.2
-    dual_head_dedup_iou = 0.35
-    remove_contained_boxes = True
+    # Per-class confidence threshold. Road signs in this dataset are
+    # frequently degraded / rear-facing / partly-obscured / distant, so we
+    # run noticeably below the fire/smoke baseline. The validator's
+    # false_positive pillar = max(0, 1 - ffpi/10): we can tolerate ~2 FP per
+    # image and still keep that pillar above 0.8.
     _conf_thres_array = np.array(
-    [0.28], dtype=np.float32
+        [0.28], dtype=np.float32
     )
+    # Per-class rescue bonus. If a class has ZERO boxes passing the threshold
+    # in a frame, its top-1 candidate is admitted when its score is at least
+    # (threshold - bonus). Bumped from 0.05 -> 0.10 so a single faint sign in
+    # an otherwise empty frame still produces a detection (map50 recall win,
+    # at most one extra FP per such frame).
     _bonus_array = np.array(
-    [0.18], dtype=np.float32
+        [0.18], dtype=np.float32
     )
 
+    # Box sanity filter: drop tiny / degenerate / image-spanning / extreme
+    # aspect ratio boxes.
+    #   min_box_area = 14x14  -> 14x14 is the smallest credible sign. The old
+    #                          value of 64 (8x8) silently discarded narrow
+    #                          distant signs like a 10x6 px overhead chevron.
+    #   min_side     = 3      -> matches min_box_area; anything thinner is
+    #                          almost certainly a pole or shadow false alarm.
+    #   max_aspect_ratio = 12.0
+    #                         -> overhead destination panels and lane-assignment
+    #                          signs are very wide (long, thin rectangles);
+    #                          8.0 was clipping legitimate detections.
     min_box_area = 8 * 8
     min_side = 3
     max_aspect_ratio = 12.0
 
+    # Tile-based TTA: when the source image is significantly larger than the
+    # model input, letterboxing throws away ~1.5x of effective resolution,
+    # which kills small-sign recall. Splitting into overlapping horizontal
+    # tiles preserves native resolution on each half. Triggered only when
+    # source width >= tile_trigger_ratio * model_input_width to avoid wasted
+    # compute on already-small images.
     tile_trigger_ratio = 1.4
     tile_overlap_ratio = 0.20
-    # use_fast_postprocess = True
-
-    @staticmethod
-    def _ort_provider_chain() -> list[str]:
-        """Prefer GPU/DML providers when installed; always end with CPU fallback."""
-        available = set(ort.get_available_providers())
-        preferred = (
-            "CPUExecutionProvider",
-            "CUDAExecutionProvider",
-            "DmlExecutionProvider",
-        )
-        return [p for p in preferred if p in available] or ["CPUExecutionProvider"]
-
-    @staticmethod
-    def _ort_session_options() -> ort.SessionOptions:
-        sess_options = ort.SessionOptions()
-        sess_options.enable_cpu_mem_arena = True
-        sess_options.enable_mem_pattern = True
-        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
-        sess_options.intra_op_num_threads = 2
-        sess_options.inter_op_num_threads = 1
-        sess_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
-        return sess_options
 
     def __init__(self, path_hf_repo: Path) -> None:
         model_path = path_hf_repo / "weights.onnx"
@@ -90,17 +94,24 @@ class Miner:
 
         print("ORT available providers BEFORE session:", ort.get_available_providers())
 
-        sess_options = self._ort_session_options()
-        providers = self._ort_provider_chain()
-        print(
-            f"ORT session config: providers={providers} "
-            f"intra_op_threads={sess_options.intra_op_num_threads}"
-        )
-        self.session = ort.InferenceSession(
-            str(model_path),
-            sess_options=sess_options,
-            providers=providers,
-        )
+        sess_options = ort.SessionOptions()
+        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+        sess_options.intra_op_num_threads = 2
+        sess_options.inter_op_num_threads = 1
+        sess_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
+
+        try:
+            self.session = ort.InferenceSession(
+                str(model_path),
+                sess_options=sess_options,
+                providers=["CPUExecutionProvider"],
+            )
+        except Exception as e:
+            self.session = ort.InferenceSession(
+                str(model_path),
+                sess_options=sess_options,
+                providers=["CPUExecutionProvider"],
+            )
 
         print("ORT session providers:", self.session.get_providers())
 
@@ -121,10 +132,10 @@ class Miner:
             dtype=np.int32,
         )
 
-        for i, inp in enumerate(self.session.get_inputs()):
-            print(f"INPUT[{i}]: shape={inp.shape} type={inp.type}")
-        for i, out in enumerate(self.session.get_outputs()):
-            print(f"OUTPUT[{i}]: shape={out.shape} type={out.type}")
+        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 = [output.name for output in self.session.get_outputs()]
@@ -148,7 +159,7 @@ class Miner:
 
         print(f"✅ ONNX model loaded from: {model_path}")
         print(f"✅ ONNX providers: {self.session.get_providers()}")
-        print(f"✅ ONNX input shape={self.input_shape}")
+        print(f"✅ ONNX input: name={self.input_name}, shape={self.input_shape}")
         print(f"✅ ONNX input size: {self.input_width}x{self.input_height}, "
               f"use_tta={self.use_tta}, use_tile_tta={self.use_tile_tta}")
         print("per-class conf: " + ", ".join(
@@ -170,7 +181,6 @@ class Miner:
 
     def _read_model_class_order(self) -> "list[str] | None":
         """Read the model's class order from Ultralytics ONNX metadata.
-
         Returns the class names ordered by model-emit index, or None when the
         metadata is missing/unparsable or doesn't match `class_names` as a set
         (in which case the static _model_class_order fallback is used)."""
@@ -257,42 +267,6 @@ class Miner:
         boxes[:, 3] = np.clip(boxes[:, 3], 0, h - 1)
         return boxes
 
-    @staticmethod
-    def _box_area(b: BoundingBox) -> int:
-        return max(0, b.x2 - b.x1) * max(0, b.y2 - b.y1)
-
-    @staticmethod
-    def _box_fully_contains(outer: BoundingBox, inner: BoundingBox) -> bool:
-        """True when outer strictly contains inner (outer area must be larger)."""
-        if (outer.x1 > inner.x1 or outer.y1 > inner.y1
-                or outer.x2 < inner.x2 or outer.y2 < inner.y2):
-            return False
-        outer_area = max(0, outer.x2 - outer.x1) * max(0, outer.y2 - outer.y1)
-        inner_area = max(0, inner.x2 - inner.x1) * max(0, inner.y2 - inner.y1)
-        return outer_area > inner_area
-
-    def _remove_contained_boxes(self, boxes: list[BoundingBox]) -> list[BoundingBox]:
-        """Drop smaller boxes fully contained inside a larger kept box."""
-        if not boxes or not self.remove_contained_boxes:
-            return boxes
-
-        sorted_boxes = sorted(boxes, key=self._box_area, reverse=True)
-        kept: list[BoundingBox] = []
-        for b in sorted_boxes:
-            if any(self._box_fully_contains(k, b) for k in kept):
-                continue
-            kept.append(b)
-        return kept
-
-    def _finalize_boxes(
-        self,
-        boxes: list[BoundingBox],
-        image_size: tuple[int, int],
-    ) -> list[BoundingBox]:
-        """Expand boxes, then remove smaller boxes contained in larger ones."""
-
-        return self._remove_contained_boxes(boxes)
-
     @staticmethod
     def _xywh_to_xyxy(boxes: np.ndarray) -> np.ndarray:
         out = np.empty_like(boxes)
@@ -415,7 +389,6 @@ class Miner:
         iou_thresh: float,
     ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
         """Remove near-duplicate boxes across classes.
-
         Order candidates by (score - per_class_threshold) margin, then by area;
         keep the highest, suppress every other box with IoU > iou_thresh.
         With a single road_sign class this is effectively a no-op, but the
@@ -463,7 +436,6 @@ class Miner:
     ) -> np.ndarray:
         """For each kept (post-NMS) box, return the max score over the FULL
         candidate set among same-class boxes with IoU >= iou_thresh.
-
         Used after horizontal-flip TTA: a high-confidence flipped detection
         can raise the score of the corresponding original detection.
         """
@@ -601,10 +573,6 @@ class Miner:
         pad: tuple[float, float],
         orig_size: tuple[int, int],
     ) -> list[BoundingBox]:
-
-        # if self.use_fast_postprocess:
-        #     return self._decode_final_dets_fast(preds, ratio, pad, orig_size)
-
         """Final-detection output path: rows shaped [x1, y1, x2, y2, conf, cls_id]."""
         if preds.ndim == 3 and preds.shape[0] == 1:
             preds = preds[0]
@@ -615,7 +583,7 @@ class Miner:
         scores = preds[:, 4].astype(np.float32)
         cls_ids = preds[:, 5].astype(np.int32)
         cls_ids = self.cls_remap[cls_ids]
-        
+
         keep = self._conf_filter_mask(scores, cls_ids)
         boxes = boxes[keep]
         scores = scores[keep]
@@ -634,6 +602,7 @@ class Miner:
         )
         if len(boxes) == 0:
             return []
+
         boxes, scores, cls_ids = self._per_view_pipeline(boxes, scores, cls_ids)
         return self._build_results(boxes, scores, cls_ids)
 
@@ -699,109 +668,6 @@ class Miner:
             return self._decode_final_dets(output, ratio, pad, orig_size)
         return self._decode_raw_yolo(output, ratio, pad, orig_size)
 
-    @staticmethod
-    def _iou_box(a: BoundingBox, b: BoundingBox) -> float:
-        x1 = max(a.x1, b.x1)
-        y1 = max(a.y1, b.y1)
-        x2 = min(a.x2, b.x2)
-        y2 = min(a.y2, b.y2)
-        inter = max(0, x2 - x1) * max(0, y2 - y1)
-        if inter <= 0:
-            return 0.0
-        area_a = max(0, a.x2 - a.x1) * max(0, a.y2 - a.y1)
-        area_b = max(0, b.x2 - b.x1) * max(0, b.y2 - b.y1)
-        union = area_a + area_b - inter
-        return float(inter / union) if union > 0 else 0.0
-
-    def _selective_merge(
-        self,
-        primary: list[BoundingBox],
-        secondary_boxes: list[BoundingBox],
-    ) -> list[BoundingBox]:
-        """Keep high-conf secondary boxes that do not overlap primary."""
-        if not secondary_boxes:
-            return []
-        merged: list[BoundingBox] = []
-        for a in secondary_boxes:
-            if a.conf < self.secondary_conf:
-                continue
-            if primary:
-                max_iou = max(self._iou_box(a, p) for p in primary)
-                if max_iou >= self.merge_iou:
-                    continue
-            merged.append(a)
-        return merged
-
-    def _harmonize_overlapping_heads(
-        self,
-        primary: list[BoundingBox],
-        secondary: list[BoundingBox],
-    ) -> list[BoundingBox]:
-        """Where primary and secondary overlap, keep primary geometry and partner score."""
-        if not primary or not secondary:
-            return primary
-
-        out: list[BoundingBox] = []
-        for p in primary:
-            partner_conf = p.conf
-            best_iou = 0.0
-            for s in secondary:
-                if s.cls_id != p.cls_id:
-                    continue
-                iou = self._iou_box(p, s)
-                if iou >= self.merge_iou and iou > best_iou:
-                    best_iou = iou
-                    partner_conf = s.conf if s.conf >= 0.90 else p.conf
-            out.append(
-                BoundingBox(
-                    x1=p.x1,
-                    y1=p.y1,
-                    x2=p.x2,
-                    y2=p.y2,
-                    cls_id=p.cls_id,
-                    conf=float(partner_conf),
-                )
-            )
-        return out
-
-    def _merge_dual_head_boxes(
-        self,
-        primary_boxes: list[BoundingBox],
-        secondary_boxes: list[BoundingBox],
-        orig_size: tuple[int, int],
-    ) -> list[BoundingBox]:
-        """Preserve primary boxes; append alex fill after extra-only dedupe.
-
-        `merge_iou` only gates cross-head exclusion in `_selective_merge`.
-        `dual_head_dedup_iou` only dedupes alex extras (never re-NMS primary).
-        Primary geometry/conf are left unchanged (no harmonize).
-        """
-        extra = self._selective_merge(primary_boxes, secondary_boxes)
-        if not extra:
-            return primary_boxes
-
-        deduped_extra = self._merge_views(
-            [extra], orig_size, iou_thresh=self.dual_head_dedup_iou
-        )
-        return primary_boxes + deduped_extra
-
-    def _merge_model_outputs(
-        self,
-        primary_out: np.ndarray,
-        secondary_out: np.ndarray,
-        ratio: float,
-        pad: tuple[float, float],
-        orig_size: tuple[int, int],
-    ) -> list[BoundingBox]:
-        """Primary detections plus optional non-overlapping secondary boxes."""
-        primary_boxes = self._postprocess(primary_out, ratio, pad, orig_size)
-        if not self.use_secondary_merge:
-            return primary_boxes
-        secondary_boxes = self._postprocess(secondary_out, ratio, pad, orig_size)
-        return self._merge_dual_head_boxes(
-            primary_boxes, secondary_boxes, orig_size
-        )
-
     def _predict_single(self, image: np.ndarray) -> list[BoundingBox]:
         if image is None:
             raise ValueError("Input image is None")
@@ -824,20 +690,10 @@ class Miner:
             )
 
         outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
-        if len(outputs) < 2:
-            return self._postprocess(outputs[0], ratio, pad, orig_size)
-
-        out0, out1 = outputs[0], outputs[1]
-        # Drop batch dim when present ([1, N, 6] -> [N, 6]).
-        if isinstance(out0, np.ndarray) and out0.ndim == 3 and out0.shape[0] == 1:
-            out0 = out0[0]
-        if isinstance(out1, np.ndarray) and out1.ndim == 3 and out1.shape[0] == 1:
-            out1 = out1[0]
-        return self._merge_model_outputs(out0, out1, ratio, pad, orig_size)
+        return self._postprocess(outputs[0], ratio, pad, orig_size)
 
     def _predict_tta(self, image: np.ndarray) -> list[BoundingBox]:
         """Horizontal-flip TTA.
-
         Strategy:
           1. Predict on original and on flipped image.
           2. Map flipped boxes back to original coordinates.
@@ -901,13 +757,11 @@ class Miner:
 
     def _predict_tiles(self, image: np.ndarray) -> list[BoundingBox]:
         """Tile-based TTA for high-resolution images.
-
         Splits the source image into two overlapping horizontal tiles, runs
         single-pass inference on each at native scale, and translates boxes
         back to the global frame. Useful when source width >> model input
         width because letterboxing otherwise discards effective resolution
         that small / distant signs depend on.
-
         Returns an empty list if the image isn't wide enough to benefit; the
         caller falls back to the regular pipeline in that case.
         """
@@ -943,16 +797,13 @@ class Miner:
         self,
         view_boxes: list[list[BoundingBox]],
         image_size: tuple[int, int],
-        iou_thresh: float | None = None,
     ) -> list[BoundingBox]:
         """Merge boxes from multiple views (single / hflip / tiles).
-
         Same logic as `_predict_tta`'s tail: per-class hard NMS to dedupe,
         then for each kept box take the max same-class score across the full
         candidate union — a high-confidence detection in any view boosts
         borderline matches in others.
         """
-        nms_iou = self.iou_thres if iou_thresh is None else float(iou_thresh)
         all_boxes: list[BoundingBox] = []
         for vb in view_boxes:
             all_boxes.extend(vb)
@@ -967,7 +818,7 @@ class Miner:
 
         coords = self._clip_boxes(coords, image_size)
 
-        hard_keep = self._per_class_hard_nms(coords, scores, cls_ids, nms_iou)
+        hard_keep = self._per_class_hard_nms(coords, scores, cls_ids, self.iou_thres)
         if len(hard_keep) == 0:
             return []
         if len(hard_keep) > self.max_det:
@@ -976,7 +827,7 @@ class Miner:
 
         boosted = self._max_score_per_cluster(
             coords[hard_keep], cls_ids[hard_keep],
-            coords, scores, cls_ids, nms_iou,
+            coords, scores, cls_ids, self.iou_thres,
         )
 
         kept_coords = coords[hard_keep]
@@ -1000,21 +851,18 @@ class Miner:
 
     def _predict_full(self, image: np.ndarray) -> list[BoundingBox]:
         """Top-level per-frame prediction with all enabled augmentations.
-
         - `use_tta=True`: original + horizontal flip
         - `use_tile_tta=True` AND image wide enough: two overlapping tiles
         All views are merged via per-class NMS + cluster-max score boost.
         """
-        h, w = image.shape[:2]
-        image_size = (w, h)
-
         if not self.use_tta and not self.use_tile_tta:
-            return self._finalize_boxes(self._predict_single(image), image_size)
+            return self._predict_single(image)
 
         views: list[list[BoundingBox]] = []
         if self.use_tta:
             views.append(self._predict_single(image))
             flipped = cv2.flip(image, 1)
+            w = image.shape[1]
             flipped_dets = self._predict_single(flipped)
             views.append([
                 BoundingBox(
@@ -1031,7 +879,8 @@ class Miner:
             if tile_boxes:
                 views.append(tile_boxes)
 
-        return self._finalize_boxes(self._merge_views(views, image_size), image_size)
+        h, w = image.shape[:2]
+        return self._merge_views(views, (w, h))
 
     def predict_batch(
         self,