scorevision: push artifact

miner.py

@@ -24,11 +24,11 @@ class TVFrameResult(BaseModel):
 
 
 class Miner:
-    def __init__(self,
+    def __init__(self, 
         path_hf_repo: Path
     ) -> None:
         model_path = path_hf_repo / "weights.onnx"
-        self.class_names = ['bus', 'car', 'truck', 'motorcycle']
+        self.class_names = ["person"]
         print("ORT version:", ort.__version__)
 
         try:
@@ -74,26 +74,26 @@ class Miner:
 
         # ---------- Scoring-oriented thresholds ----------
         # Low threshold for candidate generation
-        self.conf_thres = 0.15
+        self.conf_thres = 0.68
 
         # High-confidence boxes can survive without TTA confirmation
-        self.conf_high = 0.50
+        self.conf_high = 0.30
 
         # NMS threshold
-        self.iou_thres = 0.66
+        self.iou_thres = 0.35
 
         # TTA confirmation IoU
-        self.tta_match_iou = 0.91
+        self.tta_match_iou = 0.68
 
         self.max_det = 150
         self.use_tta = True
 
         # Box sanity filters
-        self.min_box_area = 4 * 4
-        self.min_w = 2
-        self.min_h = 2
-        self.max_aspect_ratio = 12.0
-        self.max_box_area_ratio = 0.95
+        self.min_box_area = 14 * 14
+        self.min_w = 8
+        self.min_h = 8
+        self.max_aspect_ratio = 8.0
+        self.max_box_area_ratio = 0.8
 
         print(f"✅ ONNX model loaded from: {model_path}")
         print(f"✅ ONNX providers: {self.session.get_providers()}")
@@ -217,29 +217,6 @@ class Miner:
 
         return np.array(keep, dtype=np.intp)
 
-    @classmethod
-    def _nms_per_class(
-        cls,
-        boxes: np.ndarray,
-        scores: np.ndarray,
-        cls_ids: np.ndarray,
-        iou_thresh: float,
-        max_det: int,
-    ) -> np.ndarray:
-        """NMS within each class so overlapping car vs bus predictions are not merged away."""
-        if len(boxes) == 0:
-            return np.array([], dtype=np.intp)
-        keep_all: list[int] = []
-        for c in np.unique(cls_ids):
-            idxs = np.nonzero(cls_ids == c)[0]
-            if len(idxs) == 0:
-                continue
-            local_keep = cls._hard_nms(boxes[idxs], scores[idxs], iou_thresh)
-            keep_all.extend(idxs[local_keep].tolist())
-        keep_all = np.array(keep_all, dtype=np.intp)
-        order = np.argsort(scores[keep_all])[::-1]
-        return keep_all[order[:max_det]]
-
     @staticmethod
     def _box_iou_one_to_many(box: np.ndarray, boxes: np.ndarray) -> np.ndarray:
         xx1 = np.maximum(box[0], boxes[:, 0])
@@ -317,7 +294,11 @@ class Miner:
         scores = preds[:, 4].astype(np.float32)
         cls_ids = preds[:, 5].astype(np.int32)
 
-        # All trained vehicle classes pass: bus, car, truck, motorcycle.
+        # person only
+        keep = cls_ids == 0
+        boxes = boxes[keep]
+        scores = scores[keep]
+        cls_ids = cls_ids[keep]
 
         # candidate threshold
         keep = scores >= self.conf_thres
@@ -340,9 +321,8 @@ class Miner:
         if len(boxes) == 0:
             return []
 
-        keep_idx = self._nms_per_class(
-            boxes, scores, cls_ids, self.iou_thres, self.max_det
-        )
+        keep_idx = self._hard_nms(boxes, scores, self.iou_thres)
+        keep_idx = keep_idx[: self.max_det]
 
         boxes = boxes[keep_idx]
         scores = scores[keep_idx]
@@ -404,6 +384,11 @@ class Miner:
             cls_scores = class_probs[np.arange(len(class_probs)), cls_ids]
             scores = obj * cls_scores
 
+        keep = cls_ids == 0
+        boxes_xywh = boxes_xywh[keep]
+        scores = scores[keep]
+        cls_ids = cls_ids[keep]
+
         keep = scores >= self.conf_thres
         boxes_xywh = boxes_xywh[keep]
         scores = scores[keep]
@@ -426,9 +411,8 @@ class Miner:
         if len(boxes) == 0:
             return []
 
-        keep_idx = self._nms_per_class(
-            boxes, scores, cls_ids, self.iou_thres, self.max_det
-        )
+        keep_idx = self._hard_nms(boxes, scores, self.iou_thres)
+        keep_idx = keep_idx[: self.max_det]
 
         boxes = boxes[keep_idx]
         scores = scores[keep_idx]
@@ -505,15 +489,12 @@ class Miner:
 
         coords_o = np.array([[b.x1, b.y1, b.x2, b.y2] for b in boxes_orig], dtype=np.float32) if boxes_orig else np.empty((0, 4), dtype=np.float32)
         scores_o = np.array([b.conf for b in boxes_orig], dtype=np.float32) if boxes_orig else np.empty((0,), dtype=np.float32)
-        cls_o = np.array([b.cls_id for b in boxes_orig], dtype=np.int32) if boxes_orig else np.empty((0,), dtype=np.int32)
 
         coords_f = np.array([[b.x1, b.y1, b.x2, b.y2] for b in boxes_flip], dtype=np.float32) if boxes_flip else np.empty((0, 4), dtype=np.float32)
         scores_f = np.array([b.conf for b in boxes_flip], dtype=np.float32) if boxes_flip else np.empty((0,), dtype=np.float32)
-        cls_f = np.array([b.cls_id for b in boxes_flip], dtype=np.int32) if boxes_flip else np.empty((0,), dtype=np.int32)
 
         accepted_boxes = []
         accepted_scores = []
-        accepted_cls = []
 
         # Original view candidates
         for i in range(len(coords_o)):
@@ -521,7 +502,6 @@ class Miner:
             if score >= self.conf_high:
                 accepted_boxes.append(coords_o[i])
                 accepted_scores.append(score)
-                accepted_cls.append(int(cls_o[i]))
             elif len(coords_f) > 0:
                 ious = self._box_iou_one_to_many(coords_o[i], coords_f)
                 j = int(np.argmax(ious))
@@ -529,7 +509,6 @@ class Miner:
                     fused_score = max(score, scores_f[j])
                     accepted_boxes.append(coords_o[i])
                     accepted_scores.append(fused_score)
-                    accepted_cls.append(int(cls_o[i]))
 
         # Flipped-view high-confidence boxes that original missed
         for i in range(len(coords_f)):
@@ -540,23 +519,21 @@ class Miner:
             if len(coords_o) == 0:
                 accepted_boxes.append(coords_f[i])
                 accepted_scores.append(score)
-                accepted_cls.append(int(cls_f[i]))
                 continue
 
             ious = self._box_iou_one_to_many(coords_f[i], coords_o)
             if np.max(ious) < self.tta_match_iou:
                 accepted_boxes.append(coords_f[i])
                 accepted_scores.append(score)
-                accepted_cls.append(int(cls_f[i]))
 
         if not accepted_boxes:
             return []
 
         boxes = np.array(accepted_boxes, dtype=np.float32)
         scores = np.array(accepted_scores, dtype=np.float32)
-        cls_ids = np.array(accepted_cls, dtype=np.int32)
 
-        keep = self._nms_per_class(boxes, scores, cls_ids, self.iou_thres, self.max_det)
+        keep = self._hard_nms(boxes, scores, self.iou_thres)
+        keep = keep[: self.max_det]
 
         out = []
         for idx in keep:
@@ -567,7 +544,7 @@ class Miner:
                     y1=int(math.floor(y1)),
                     x2=int(math.ceil(x2)),
                     y2=int(math.ceil(y2)),
-                    cls_id=int(cls_ids[idx]),
+                    cls_id=0,
                     conf=float(scores[idx]),
                 )
             )
@@ -620,4 +597,4 @@ class Miner:
                 )
             )
 
-        return results
+        return results