scorevision: push artifact

miner.py

@@ -149,9 +149,18 @@ class Miner:
     """
 
     def __init__(self, path_hf_repo: Path) -> None:
-        self.class_names = ["vehicle"]
-        # COCO class indices we keep:
-        self.vehicle_coco_ids = (2, 3, 5, 7)  # car, motorcycle, bus, truck
+        # IMPORTANT: element `manak0/Detect-detect-vehicle` declares:
+        #   objects = ["bus", "car", "truck", "motorcycle"]  (in this order)
+        # So the validator maps cls_id 0->"bus", 1->"car", 2->"truck", 3->"motorcycle".
+        # We must emit cls_id matching this order.
+        self.class_names = ["bus", "car", "truck", "motorcycle"]
+        # COCO -> element cls_id:
+        #   bus (5)        -> 0
+        #   car (2)        -> 1
+        #   truck (7)      -> 2
+        #   motorcycle (3) -> 3
+        self.coco_to_element = {5: 0, 2: 1, 7: 2, 3: 3}
+        self.vehicle_coco_ids = tuple(self.coco_to_element.keys())
 
         weights_name = os.environ.get("SN44_ONNX_WEIGHTS", _DEFAULT_WEIGHTS)
         weights_path = path_hf_repo / weights_name
@@ -306,16 +315,25 @@ class Miner:
         boxes_xywh = preds[:, :4].astype(np.float32)
         class_scores = preds[:, 4:].astype(np.float32)
 
-        # For each detection keep MAX confidence across the vehicle classes.
-        vehicle_idx = np.array(self.vehicle_coco_ids, dtype=np.intp)
-        vehicle_scores_all = class_scores[:, vehicle_idx]
-        vehicle_scores = vehicle_scores_all.max(axis=1)
+        # For each detection, determine which COCO vehicle class has the MAX score
+        # and retain that as the class id + score.
+        vehicle_coco_idx = np.array(self.vehicle_coco_ids, dtype=np.intp)
+        vehicle_class_scores = class_scores[:, vehicle_coco_idx]  # (N, 4)
+        best_in_vehicle = vehicle_class_scores.argmax(axis=1)     # (N,) index into vehicle_coco_idx
+        vehicle_scores = vehicle_class_scores.max(axis=1)         # (N,)
+        best_coco_ids = vehicle_coco_idx[best_in_vehicle]         # COCO ids per row
+
         mask = vehicle_scores >= self.conf_thres
         if not np.any(mask):
             return []
 
         boxes_xywh = boxes_xywh[mask]
         scores = vehicle_scores[mask]
+        # Map COCO id -> element cls_id (for emission)
+        element_cls_ids = np.array(
+            [self.coco_to_element[int(c)] for c in best_coco_ids[mask]],
+            dtype=np.int32,
+        )
 
         boxes = _xywh_to_xyxy(boxes_xywh)
 
@@ -325,7 +343,9 @@ class Miner:
         boxes /= ratio
         boxes = _clip_boxes(boxes, orig_size)
 
-        boxes, scores = self._filter_sane(boxes, scores, orig_size)
+        boxes, scores, element_cls_ids = self._filter_sane_with_cls(
+            boxes, scores, element_cls_ids, orig_size,
+        )
         if len(boxes) == 0:
             return []
 
@@ -333,9 +353,10 @@ class Miner:
         keep = keep[: self.max_det]
         boxes = boxes[keep]
         scores = scores[keep]
+        element_cls_ids = element_cls_ids[keep]
 
         out: list[BoundingBox] = []
-        for box, conf in zip(boxes, scores):
+        for box, conf, cls_id in zip(boxes, scores, element_cls_ids):
             if box[2] <= box[0] or box[3] <= box[1]:
                 continue
             out.append(
@@ -344,12 +365,50 @@ class Miner:
                     y1=int(math.floor(box[1])),
                     x2=int(math.ceil(box[2])),
                     y2=int(math.ceil(box[3])),
-                    cls_id=0,
+                    cls_id=int(cls_id),
                     conf=float(conf),
                 )
             )
         return out
 
+    def _filter_sane_with_cls(
+        self,
+        boxes: np.ndarray,
+        scores: np.ndarray,
+        cls_ids: np.ndarray,
+        orig_size: tuple[int, int],
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        if len(boxes) == 0:
+            return boxes, scores, cls_ids
+        orig_w, orig_h = orig_size
+        image_area = float(orig_w * orig_h)
+        keep: list[int] = []
+        for i, box in enumerate(boxes):
+            x1, y1, x2, y2 = box.tolist()
+            bw = x2 - x1
+            bh = y2 - y1
+            if bw <= 0 or bh <= 0:
+                continue
+            if bw < self.min_w or bh < self.min_h:
+                continue
+            area = bw * bh
+            if area < self.min_box_area:
+                continue
+            if area > self.max_box_area_ratio * image_area:
+                continue
+            ar = max(bw / max(bh, 1e-6), bh / max(bw, 1e-6))
+            if ar > self.max_aspect_ratio:
+                continue
+            keep.append(i)
+        if not keep:
+            return (
+                np.empty((0, 4), dtype=np.float32),
+                np.empty((0,), dtype=np.float32),
+                np.empty((0,), dtype=np.int32),
+            )
+        idx = np.array(keep, dtype=np.intp)
+        return boxes[idx], scores[idx], cls_ids[idx]
+
     @staticmethod
     def _coco_classes() -> list[str]:
         return [