Upload miner.py with huggingface_hub

miner.py

@@ -24,20 +24,32 @@ class TVFrameResult(BaseModel):
 
 
 class Miner:
-    """ONNX Runtime miner. Hard global NMS + sanity filter + dedup + flip TTA, with per-class rescue bonus."""
+    """ONNX Runtime miner with per-class candidates, TTA fusion, and temporal rescue."""
 
     class_names = ["cup", "bottle", "can"]
-    model_class_names = ["bottle", "can", "cup"]
-    _model_to_competition_cls = np.array([1, 2, 0], dtype=np.int32)
+    model_class_names = ["cup", "bottle", "can"]
+    _model_to_competition_cls = np.array([0, 1, 2], dtype=np.int32)
     input_size = 1280
     iou_thres = 0.3
-    cross_iou_thresh = 0.7
+    cross_iou_thresh = 0.65
     min_side = 8.0
     min_box_area = 100.0
     max_aspect_ratio = 10.0
     max_det = 300
-    _conf_thres_array = np.array([0.35, 0.5, 0.5], dtype=np.float32)
-    _bonus_array = np.array([0.0, 0.0, 0.2], dtype=np.float32)
+    _conf_thres_array = np.array([0.60, 0.45, 0.50], dtype=np.float32)
+    _candidate_conf_thres_array = np.array([0.20, 0.30, 0.30], dtype=np.float32)
+    _tta_conf_thres_array = np.array([0.52, 0.37, 0.42], dtype=np.float32)
+    _temporal_conf_thres_array = np.array([0.54, 0.39, 0.44], dtype=np.float32)
+    _tta_confirmed_views = 1
+    temporal_iou_thresh = 0.25
+    track_iou_thresh = 0.35
+    track_keep_frames = 2
+    track_min_conf = np.array([0.50, 0.35, 0.40], dtype=np.float32)
+    sparse_candidate_count = 8
+    crowded_candidate_count = 28
+    crowded_area_ratio = 0.030
+    sparse_relax = 0.04
+    crowded_raise = 0.04
 
     def __init__(self, path_hf_repo: Path) -> None:
         model_path = path_hf_repo / "weights.onnx"
@@ -83,6 +95,8 @@ class Miner:
 
         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)
+        self._tracks: list[dict] = []
+        self._last_track_frame_id: int | None = None
 
         print(f"ONNX model loaded from: {model_path}")
         print(f"ONNX providers: {self.session.get_providers()}")
@@ -208,9 +222,18 @@ class Miner:
     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]:
+        keep_idx = self._cross_class_dedup_keep_indices(
+            boxes, scores, cls_ids, iou_thresh
+        )
+        return boxes[keep_idx], scores[keep_idx], cls_ids[keep_idx]
+
+    def _cross_class_dedup_keep_indices(self, boxes: np.ndarray,
+                                        scores: np.ndarray,
+                                        cls_ids: np.ndarray,
+                                        iou_thresh: float) -> np.ndarray:
         n = len(boxes)
         if n <= 1:
-            return boxes, scores, cls_ids
+            return np.arange(n, dtype=np.intp)
         boxes = np.asarray(boxes, dtype=np.float32)
         scores = np.asarray(scores, dtype=np.float32)
         cls_ids = np.asarray(cls_ids, dtype=np.int32)
@@ -235,8 +258,7 @@ class Miner:
             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]
+        return np.array(keep, dtype=np.intp)
 
     def _filter_sane_boxes(self, boxes: np.ndarray, scores: np.ndarray,
                            cls_ids: np.ndarray, orig_size: tuple[int, int]
@@ -286,27 +308,130 @@ class Miner:
             out[i] = float(np.max(full_scores[cluster])) if np.any(cluster) else 0.0
         return out
 
+    def _view_support_per_cluster(self, post_boxes: np.ndarray,
+                                  post_cls: np.ndarray,
+                                  full_boxes: np.ndarray,
+                                  full_cls: np.ndarray,
+                                  full_view_ids: np.ndarray,
+                                  iou_thresh: float) -> np.ndarray:
+        n = len(post_boxes)
+        if n == 0:
+            return np.empty(0, dtype=np.int32)
+        full_areas = (np.maximum(0.0, full_boxes[:, 2] - full_boxes[:, 0]) *
+                      np.maximum(0.0, full_boxes[:, 3] - full_boxes[:, 1]))
+        out = np.ones(n, dtype=np.int32)
+        for i in range(n):
+            bi = post_boxes[i]
+            xx1 = np.maximum(bi[0], full_boxes[:, 0])
+            yy1 = np.maximum(bi[1], full_boxes[:, 1])
+            xx2 = np.minimum(bi[2], full_boxes[:, 2])
+            yy2 = np.minimum(bi[3], full_boxes[:, 3])
+            inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
+            a_i = max(0.0, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
+            iou = inter / (a_i + full_areas - inter + 1e-7)
+            cluster = (iou >= iou_thresh) & (full_cls == post_cls[i])
+            if np.any(cluster):
+                out[i] = int(len(np.unique(full_view_ids[cluster])))
+        return out
+
     def _conf_filter_mask(self, scores: np.ndarray,
                           cls_ids: np.ndarray) -> np.ndarray:
-        """Boolean keep-mask: score >= per-class threshold, with a per-class
-        rescue — if a class has zero boxes passing, admit its top-1 candidate
-        when its score >= (per-class threshold - per-class bonus)."""
+        """Keep low-score candidates; final acceptance happens after evidence fusion."""
         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
+        return scores >= self._candidate_conf_thres_array[cls_ids]
+
+    def _scene_adjustment(self, boxes: list[BoundingBox],
+                          image_shape: tuple[int, int, int] | None) -> float:
+        if not boxes or image_shape is None:
+            return 0.0
+        h, w = image_shape[:2]
+        image_area = max(1.0, float(w * h))
+        total_box_area = sum(
+            max(0, box.x2 - box.x1) * max(0, box.y2 - box.y1)
+            for box in boxes
+        )
+        area_ratio = float(total_box_area) / image_area
+        if len(boxes) >= self.crowded_candidate_count or area_ratio >= self.crowded_area_ratio:
+            return self.crowded_raise
+        if len(boxes) <= self.sparse_candidate_count and area_ratio < self.crowded_area_ratio * 0.5:
+            return -self.sparse_relax
+        return 0.0
+
+    def _adaptive_thresholds(self, cls_ids: np.ndarray,
+                             boxes: list[BoundingBox],
+                             image_shape: tuple[int, int, int] | None
+                             ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        adjustment = self._scene_adjustment(boxes, image_shape)
+        auto = np.clip(
+            self._conf_thres_array[cls_ids] + adjustment,
+            self._candidate_conf_thres_array[cls_ids] + 0.05,
+            0.95,
+        )
+        tta = np.clip(
+            self._tta_conf_thres_array[cls_ids] + adjustment,
+            self._candidate_conf_thres_array[cls_ids],
+            auto,
+        )
+        temporal = np.clip(
+            self._temporal_conf_thres_array[cls_ids] + adjustment,
+            self._candidate_conf_thres_array[cls_ids],
+            auto,
+        )
+        return auto, tta, temporal
+
+    def _candidate_accept_mask(self, boxes: list[BoundingBox],
+                               view_support: np.ndarray,
+                               image_shape: tuple[int, int, int] | None
+                               ) -> np.ndarray:
+        _, scores, cls_ids = self._boxes_to_arrays(boxes)
+        if len(scores) == 0:
+            return np.zeros(0, dtype=bool)
+        auto_thres, tta_thres, _ = self._adaptive_thresholds(
+            cls_ids, boxes, image_shape
+        )
+        auto = scores >= auto_thres
+        tta_confirmed = (
+            (scores >= tta_thres) &
+            (view_support >= self._tta_confirmed_views)
+        )
+        return auto | tta_confirmed
+
+    @staticmethod
+    def _boxes_to_arrays(boxes: list[BoundingBox]
+                         ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        if not boxes:
+            return (
+                np.empty((0, 4), dtype=np.float32),
+                np.empty(0, dtype=np.float32),
+                np.empty(0, dtype=np.int32),
+            )
+        coords = np.array(
+            [[b.x1, b.y1, b.x2, b.y2] for b in boxes], dtype=np.float32
+        )
+        scores = np.array([b.conf for b in boxes], dtype=np.float32)
+        cls_ids = np.array([b.cls_id for b in boxes], dtype=np.int32)
+        return coords, scores, cls_ids
+
+    @staticmethod
+    def _image_shape(image: np.ndarray | None) -> tuple[int, int, int] | None:
+        if isinstance(image, np.ndarray) and image.ndim == 3:
+            return image.shape
+        return None
+
+    @staticmethod
+    def _single_box_iou(box: BoundingBox, boxes: np.ndarray) -> np.ndarray:
+        if len(boxes) == 0:
+            return np.empty(0, dtype=np.float32)
+        xx1 = np.maximum(float(box.x1), boxes[:, 0])
+        yy1 = np.maximum(float(box.y1), boxes[:, 1])
+        xx2 = np.minimum(float(box.x2), boxes[:, 2])
+        yy2 = np.minimum(float(box.y2), boxes[:, 3])
+        inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
+        a_i = max(0.0, float((box.x2 - box.x1) * (box.y2 - box.y1)))
+        areas = (np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) *
+                 np.maximum(0.0, boxes[:, 3] - boxes[:, 1]))
+        return inter / (a_i + areas - inter + 1e-7)
 
     def _per_view_pipeline(self, boxes: np.ndarray, scores: np.ndarray,
                            cls_ids: np.ndarray, orig_size: tuple[int, int]
@@ -317,7 +442,7 @@ class Miner:
         if len(boxes) == 0:
             return boxes, scores, cls_ids
         if len(boxes) > 1:
-            keep = self._hard_nms(boxes, scores, self.iou_thres)
+            keep = self._per_class_hard_nms(boxes, scores, cls_ids, self.iou_thres)
             boxes, scores, cls_ids = boxes[keep], scores[keep], cls_ids[keep]
         if len(scores) > self.max_det:
             top = np.argsort(-scores)[: self.max_det]
@@ -457,7 +582,8 @@ class Miner:
         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]:
+    def _predict_tta_candidates(self, image: np.ndarray
+                                ) -> tuple[list[BoundingBox], np.ndarray]:
         boxes_orig = self._predict_single(image)
         flipped = cv2.flip(image, 1)
         boxes_flip = self._predict_single(flipped)
@@ -471,17 +597,16 @@ class Miner:
         ]
         all_boxes = boxes_orig + boxes_flip
         if not all_boxes:
-            return []
+            return [], np.empty(0, dtype=np.int32)
 
-        coords = np.array(
-            [[b.x1, b.y1, b.x2, b.y2] for b in all_boxes], dtype=np.float32
+        coords, scores, cls_ids = self._boxes_to_arrays(all_boxes)
+        view_ids = np.array(
+            [0] * len(boxes_orig) + [1] * len(boxes_flip), dtype=np.int32
         )
-        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)
 
         hard_keep = self._per_class_hard_nms(coords, scores, cls_ids, self.iou_thres)
         if len(hard_keep) == 0:
-            return []
+            return [], np.empty(0, dtype=np.int32)
         if len(hard_keep) > self.max_det:
             top = np.argsort(-scores[hard_keep])[: self.max_det]
             hard_keep = hard_keep[top]
@@ -492,12 +617,19 @@ class Miner:
 
         kept_coords = coords[hard_keep]
         kept_cls = cls_ids[hard_keep]
+        view_support = self._view_support_per_cluster(
+            kept_coords, kept_cls, coords, cls_ids, view_ids, self.iou_thres,
+        )
         if len(kept_coords) > 1:
-            kept_coords, boosted, kept_cls = self._cross_class_dedup_op(
+            dedup_keep = self._cross_class_dedup_keep_indices(
                 kept_coords, boosted, kept_cls, self.cross_iou_thresh
             )
+            kept_coords = kept_coords[dedup_keep]
+            boosted = boosted[dedup_keep]
+            kept_cls = kept_cls[dedup_keep]
+            view_support = view_support[dedup_keep]
 
-        return [
+        boxes = [
             BoundingBox(
                 x1=int(math.floor(kept_coords[j, 0])),
                 y1=int(math.floor(kept_coords[j, 1])),
@@ -508,19 +640,159 @@ class Miner:
             )
             for j in range(len(kept_coords))
         ]
+        return boxes, view_support
+
+    def _filter_by_evidence(self, boxes: list[BoundingBox],
+                            view_support: np.ndarray,
+                            image_shape: tuple[int, int, int] | None
+                            ) -> list[BoundingBox]:
+        keep = self._candidate_accept_mask(boxes, view_support, image_shape)
+        return [box for box, ok in zip(boxes, keep) if bool(ok)]
+
+    def _predict_tta(self, image: np.ndarray) -> list[BoundingBox]:
+        boxes, view_support = self._predict_tta_candidates(image)
+        return self._filter_by_evidence(boxes, view_support, image.shape)
+
+    def _has_temporal_support(self, frame_idx: int, box: BoundingBox,
+                              candidate_boxes: list[list[BoundingBox]],
+                              initial_keep: list[np.ndarray]) -> bool:
+        neighbor_indices = [
+            idx for idx in (frame_idx - 1, frame_idx + 1)
+            if 0 <= idx < len(candidate_boxes)
+        ]
+        for idx in neighbor_indices:
+            coords, _, cls_ids = self._boxes_to_arrays(candidate_boxes[idx])
+            same_cls = cls_ids == box.cls_id
+            if not np.any(same_cls):
+                continue
+            accepted = same_cls & initial_keep[idx]
+            if np.any(accepted):
+                if np.max(self._single_box_iou(box, coords[accepted])) >= self.temporal_iou_thresh:
+                    return True
+
+        two_sided_candidate_support = []
+        for idx in (frame_idx - 1, frame_idx + 1):
+            if not 0 <= idx < len(candidate_boxes):
+                two_sided_candidate_support.append(False)
+                continue
+            coords, scores, cls_ids = self._boxes_to_arrays(candidate_boxes[idx])
+            same_cls = cls_ids == box.cls_id
+            if not np.any(same_cls):
+                two_sided_candidate_support.append(False)
+                continue
+            score_ok = scores >= self._temporal_conf_thres_array[cls_ids]
+            neighbor_ok = same_cls & score_ok
+            supported = (
+                np.any(neighbor_ok) and
+                np.max(self._single_box_iou(box, coords[neighbor_ok])) >= self.temporal_iou_thresh
+            )
+            two_sided_candidate_support.append(bool(supported))
+        return all(two_sided_candidate_support)
+
+    def _reset_tracks_if_needed(self, frame_id: int) -> None:
+        if self._last_track_frame_id is None:
+            self._last_track_frame_id = frame_id - 1
+            return
+        if frame_id <= self._last_track_frame_id:
+            self._tracks = []
+        self._last_track_frame_id = frame_id
+
+    def _track_supported(self, box: BoundingBox, frame_id: int) -> bool:
+        best_iou = 0.0
+        for track in self._tracks:
+            if int(track["cls_id"]) != box.cls_id:
+                continue
+            age = frame_id - int(track["frame_id"])
+            if age < 1 or age > self.track_keep_frames:
+                continue
+            iou = self._single_box_iou(box, track["coords"])[0]
+            best_iou = max(best_iou, float(iou))
+        return best_iou >= self.track_iou_thresh
+
+    def _update_tracks(self, boxes: list[BoundingBox], frame_id: int) -> None:
+        fresh_tracks = []
+        for track in self._tracks:
+            if frame_id - int(track["frame_id"]) <= self.track_keep_frames:
+                fresh_tracks.append(track)
+        for box in boxes:
+            coords = np.array(
+                [[box.x1, box.y1, box.x2, box.y2]], dtype=np.float32
+            )
+            updated = False
+            for track in fresh_tracks:
+                if int(track["cls_id"]) != box.cls_id:
+                    continue
+                iou = self._single_box_iou(box, track["coords"])[0]
+                if iou >= self.track_iou_thresh:
+                    track["coords"] = coords
+                    track["frame_id"] = frame_id
+                    track["conf"] = box.conf
+                    updated = True
+                    break
+            if not updated:
+                fresh_tracks.append(
+                    {
+                        "coords": coords,
+                        "cls_id": box.cls_id,
+                        "conf": box.conf,
+                        "frame_id": frame_id,
+                    }
+                )
+        self._tracks = fresh_tracks
 
     def predict_batch(self, batch_images: list[ndarray], offset: int,
                       n_keypoints: int) -> list[TVFrameResult]:
+        candidate_boxes: list[list[BoundingBox]] = []
+        view_supports: list[np.ndarray] = []
+        image_shapes = [self._image_shape(image) for image in batch_images]
         results: list[TVFrameResult] = []
         for frame_number_in_batch, image in enumerate(batch_images):
             try:
-                boxes = self._predict_tta(image)
+                boxes, view_support = self._predict_tta_candidates(image)
             except Exception as e:
                 print(f"Inference failed for frame {offset + frame_number_in_batch}: {e}")
                 boxes = []
+                view_support = np.empty(0, dtype=np.int32)
+            candidate_boxes.append(boxes)
+            view_supports.append(view_support)
+
+        initial_keep: list[np.ndarray] = []
+        for boxes, view_support, image_shape in zip(
+            candidate_boxes, view_supports, image_shapes
+        ):
+            initial_keep.append(
+                self._candidate_accept_mask(boxes, view_support, image_shape)
+            )
+
+        for frame_number_in_batch, boxes in enumerate(candidate_boxes):
+            frame_id = offset + frame_number_in_batch
+            self._reset_tracks_if_needed(frame_id)
+            keep = initial_keep[frame_number_in_batch].copy()
+            _, scores, cls_ids = self._boxes_to_arrays(boxes)
+            _, _, temporal_thres = self._adaptive_thresholds(
+                cls_ids, boxes, image_shapes[frame_number_in_batch]
+            )
+            temporal_ready = scores >= temporal_thres
+            track_ready = scores >= self.track_min_conf[cls_ids]
+            for i, box in enumerate(boxes):
+                if keep[i]:
+                    continue
+                has_neighbor_support = (
+                    bool(temporal_ready[i]) and
+                    self._has_temporal_support(
+                        frame_number_in_batch, box, candidate_boxes, initial_keep
+                    )
+                )
+                has_track_support = (
+                    bool(track_ready[i]) and self._track_supported(box, frame_id)
+                )
+                if has_neighbor_support or has_track_support:
+                    keep[i] = True
+            boxes = [box for box, ok in zip(boxes, keep) if bool(ok)]
+            self._update_tracks(boxes, frame_id)
             results.append(
                 TVFrameResult(
-                    frame_id=offset + frame_number_in_batch,
+                    frame_id=frame_id,
                     boxes=boxes,
                     keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))],
                 )