scorevision: push artifact

miner.py

@@ -33,19 +33,41 @@ class Miner:
     class_names = ["balaclava", "hoodie", "glove", "bat", "spray paint", "graffiti"]
     input_size = 1280
     iou_thres = 0.30
-    cross_iou_thresh = 0.70
+    cross_iou_thresh = 0.80
     max_aspect_ratio = 10.0
     max_det = 150
-    # tuned on consensus-GT, healthy distribution: [balaclava, hoodie, glove, bat, spray, graffiti]
     _conf_thres_array = np.array(
         [0.30, 0.70, 0.70, 0.50, 0.50, 0.40], dtype=np.float32
     )
-    # per-class rescue gap: if a class has no box above threshold, admit top-1
-    # if its score >= (threshold - bonus). Inspired by new-vmodel0's _bonus_array.
     _bonus_array = np.array(
         [0.20, 0.25, 0.20, 0.30, 0.20, 0.15], dtype=np.float32
     )
 
+    _glove_cls_id = 2
+    _glove_spurious_max_conf = 0.04
+    _glove_spurious_aspect = 2.5
+    _glove_crowd_min_count = 5
+    _glove_crowd_max_conf_trigger = 0.6
+    _glove_crowd_keep_min_conf = 0.27
+    _glove_split_low_max = 0.26
+    _glove_split_low_min_count = 4
+    _glove_split_high_min_count = 2
+    _glove_pair_low_max = 0.1
+    _glove_pair_high_min = 0.4
+    _glove_pair_max_center_x_dist = 0.028
+    _glove_reject_r_min = 100
+    _glove_reject_r_max = 125
+    _glove_reject_g_min = 88
+    _glove_reject_g_max = 115
+    _glove_reject_b_min = 68
+    _glove_reject_b_max = 88
+    _glove_reject_min_pixel_frac = 0.04
+    _glove_reject_p90_percentile = 90.0
+
+    _spray_paint_cls_id = 4
+    _spray_orphan_max_conf = 0.3
+    _spray_glove_min_overlap = 0.3
+
     def __init__(self, path_hf_repo: Path) -> None:
         model_path = path_hf_repo / "weights.onnx"
         print("ORT version:", ort.__version__)
@@ -234,6 +256,127 @@ class Miner:
         boxes = self._clip(boxes, orig_size)
         return boxes, scores, cls_ids
 
+    @classmethod
+    def _is_spurious_glove(cls, cls_id, x1, y1, x2, y2, conf):
+        if cls_id != cls._glove_cls_id: return False
+        w, h = x2 - x1, y2 - y1
+        if w <= 0: return False
+        return h > cls._glove_spurious_aspect * w and conf < cls._glove_spurious_max_conf
+
+    @staticmethod
+    def _box_intersection_area(a, b):
+        ix1, iy1 = max(a.x1, b.x1), max(a.y1, b.y1)
+        ix2, iy2 = min(a.x2, b.x2), min(a.y2, b.y2)
+        return max(0.0, ix2 - ix1) * max(0.0, iy2 - iy1)
+
+    @classmethod
+    def _spray_overlaps_glove(cls, spray, glove):
+        inter = cls._box_intersection_area(spray, glove)
+        area = max(0.0, float(spray.x2 - spray.x1) * float(spray.y2 - spray.y1))
+        if area <= 0.0: return False
+        return inter / area >= cls._spray_glove_min_overlap
+
+    @classmethod
+    def _is_glove_rejection_color(cls, cls_id, image, x1, y1, x2, y2):
+        if cls_id != cls._glove_cls_id or image is None: return False
+        H, W = image.shape[:2]
+        xi1 = max(0, int(math.floor(x1)))
+        yi1 = max(0, int(math.floor(y1)))
+        xi2 = min(W, int(math.ceil(x2)))
+        yi2 = min(H, int(math.ceil(y2)))
+        if xi2 <= xi1 or yi2 <= yi1: return False
+        crop = image[yi1:yi2, xi1:xi2]
+        if crop.size == 0: return False
+        b = crop[:, :, 0].astype(np.float32)
+        g = crop[:, :, 1].astype(np.float32)
+        r = crop[:, :, 2].astype(np.float32)
+        in_range = ((r >= cls._glove_reject_r_min) & (r <= cls._glove_reject_r_max)
+                    & (g >= cls._glove_reject_g_min) & (g <= cls._glove_reject_g_max)
+                    & (b >= cls._glove_reject_b_min) & (b <= cls._glove_reject_b_max))
+        if float(in_range.mean()) >= cls._glove_reject_min_pixel_frac:
+            return True
+        p90_r = float(np.percentile(r, cls._glove_reject_p90_percentile))
+        p90_g = float(np.percentile(g, cls._glove_reject_p90_percentile))
+        p90_b = float(np.percentile(b, cls._glove_reject_p90_percentile))
+        return (cls._glove_reject_r_min <= p90_r <= cls._glove_reject_r_max
+                and cls._glove_reject_g_min <= p90_g <= cls._glove_reject_g_max
+                and cls._glove_reject_b_min <= p90_b <= cls._glove_reject_b_max)
+
+    @classmethod
+    def _filter_glove_by_color(cls, results, image):
+        if image is None: return results
+        return [b for b in results
+                if not cls._is_glove_rejection_color(b.cls_id, image, b.x1, b.y1, b.x2, b.y2)]
+
+    @staticmethod
+    def _glove_cx(b): return (b.x1 + b.x2) / 2.0
+
+    @classmethod
+    def _filter_paired_low_glove(cls, results, image_width):
+        if image_width <= 0: return results
+        gloves = [b for b in results if b.cls_id == cls._glove_cls_id]
+        if len(gloves) < 2: return results
+        remove = set()
+        for i, a in enumerate(gloves):
+            for b in gloves[i+1:]:
+                if a.conf < cls._glove_pair_low_max and b.conf > cls._glove_pair_high_min:
+                    low, high = a, b
+                elif b.conf < cls._glove_pair_low_max and a.conf > cls._glove_pair_high_min:
+                    low, high = b, a
+                else:
+                    continue
+                cx_dist = abs(cls._glove_cx(low) - cls._glove_cx(high)) / float(image_width)
+                if cx_dist < cls._glove_pair_max_center_x_dist:
+                    remove.add(id(low))
+        return [b for b in results if id(b) not in remove]
+
+    @classmethod
+    def _filter_weak_glove_with_strong_others(cls, results):
+        others = [b for b in results if b.cls_id != cls._glove_cls_id]
+        if not others: return results
+        if not all(b.conf > cls._glove_pair_high_min for b in others): return results
+        return [b for b in results
+                if b.cls_id != cls._glove_cls_id or b.conf >= cls._glove_pair_low_max]
+
+    @classmethod
+    def _filter_split_glove_confidence(cls, results):
+        gloves = [b for b in results if b.cls_id == cls._glove_cls_id]
+        low = [b for b in gloves if b.conf < cls._glove_split_low_max]
+        high = [b for b in gloves if b.conf > cls._glove_split_low_max]
+        if len(low) <= cls._glove_split_low_min_count: return results
+        if len(high) < cls._glove_split_high_min_count: return results
+        best_low = max(low, key=lambda b: b.conf)
+        others = [b for b in results if b.cls_id != cls._glove_cls_id]
+        return others + high + [best_low]
+
+    @classmethod
+    def _filter_crowded_gloves(cls, results):
+        gloves = [b for b in results if b.cls_id == cls._glove_cls_id]
+        if len(gloves) <= cls._glove_crowd_min_count: return results
+        if max(b.conf for b in gloves) <= cls._glove_crowd_max_conf_trigger: return results
+        return [b for b in results
+                if b.cls_id != cls._glove_cls_id or b.conf > cls._glove_crowd_keep_min_conf]
+
+    @classmethod
+    def _filter_orphan_low_conf_spray(cls, results):
+        gloves = [b for b in results if b.cls_id == cls._glove_cls_id]
+        kept = []
+        for box in results:
+            if box.cls_id != cls._spray_paint_cls_id or box.conf >= cls._spray_orphan_max_conf:
+                kept.append(box); continue
+            if any(cls._spray_overlaps_glove(box, g) for g in gloves):
+                kept.append(box)
+        return kept
+
+    @classmethod
+    def _apply_rules(cls, results, image_width, image=None):
+        results = cls._filter_glove_by_color(results, image)
+        results = cls._filter_paired_low_glove(results, image_width)
+        results = cls._filter_weak_glove_with_strong_others(results)
+        results = cls._filter_split_glove_confidence(results)
+        results = cls._filter_crowded_gloves(results)
+        return cls._filter_orphan_low_conf_spray(results)
+
     def _predict_single(self, image):
         x, ratio, pad, orig_size = self._preprocess(image)
         out = self.session.run(self.output_names, {self.input_name: x})[0]
@@ -264,15 +407,17 @@ class Miner:
         if len(scores) > self.max_det:
             top = np.argsort(-scores)[:self.max_det]
             boxes, scores, cls_ids = boxes[top], scores[top], cls_ids[top]
-        return [
-            BoundingBox(
-                x1=int(math.floor(b[0])), y1=int(math.floor(b[1])),
-                x2=int(math.ceil(b[2])), y2=int(math.ceil(b[3])),
+        results = []
+        for b, s, c in zip(boxes, scores, cls_ids):
+            if b[2] <= b[0] or b[3] <= b[1]: continue
+            x1, y1, x2, y2 = float(b[0]), float(b[1]), float(b[2]), float(b[3])
+            if self._is_spurious_glove(int(c), x1, y1, x2, y2, float(s)): continue
+            results.append(BoundingBox(
+                x1=int(math.floor(x1)), y1=int(math.floor(y1)),
+                x2=int(math.ceil(x2)), y2=int(math.ceil(y2)),
                 cls_id=int(c), conf=float(s),
-            )
-            for b, s, c in zip(boxes, scores, cls_ids)
-            if b[2] > b[0] and b[3] > b[1]
-        ]
+            ))
+        return self._apply_rules(results, image.shape[1], image=image)
 
     def predict_batch(self, batch_images, offset, n_keypoints):
         results = []