Upload folder using huggingface_hub

miner.py

@@ -85,12 +85,31 @@ class Miner:
         self.input_width = self._safe_dim(self.input_shape[3], default=1280)
 
         # Tuned for validator scoring (pillars: 0.6*map50 + 0.4*false_positive).
-        # conf 0.40 was the best element-score point in the v1 eval sweep.
-        self.conf_thres = 0.32  # Higher = fewer FP, slightly lower recall
-        self.iou_thres = 0.5   # Lower = suppress duplicate detections (FP)
-        self.max_det = 200      # Cap detections per image
+        self.iou_thres = 0.5         # Per-class NMS IoU; lower = stricter dedup
+        self.cross_iou_thresh = 0.8  # Cross-class dedup IoU (suppress same physical object firing multiple classes)
+        self.max_det = 200
         self.use_tta = True
 
+        # Per-class confidence thresholds (ported pattern from fire001 miner).
+        # A single global conf cannot serve both tiny nozzles and large tracks.
+        # Indexed by class_names order: [broom, drainage gate, nozzle, track].
+        #   broom         (0.35) -- distinctive long handle, moderate
+        #   drainage gate (0.30) -- floor element often water-obscured
+        #   nozzle        (0.25) -- TINY GT objects (median ~290 px²), permissive
+        #   track         (0.35) -- large clear object when present, moderate
+        self._conf_thres_array = np.array(
+            [0.35, 0.015, 0.4, 0.3], dtype=np.float32
+        )
+        # Per-class rescue bonus: when a class has ZERO boxes passing the
+        # threshold in a frame, its top-1 candidate is admitted when its score
+        # is at least (per-class threshold - per-class bonus). Nozzles get the
+        # biggest rescue because spray + motion blur often shaves a few points
+        # off otherwise valid detections; track gets the smallest because it's
+        # rarely borderline.
+        self._bonus_array = np.array(
+            [0.10, 0.10, 0.05, 0.05], dtype=np.float32
+        )
+
         # Box sanity filter — kept loose: car-wash `nozzle` boxes are tiny
         # (GT median ~290 px², smallest ~32 px²). Fire's 14x14/min_side 8
         # would delete valid nozzles, so thresholds are dropped here.
@@ -370,35 +389,126 @@ class Miner:
 
     @staticmethod
     def _max_score_per_cluster(
-        coords: np.ndarray,
-        scores: np.ndarray,
-        keep_indices: np.ndarray,
+        post_boxes: np.ndarray,
+        post_cls: np.ndarray,
+        full_boxes: np.ndarray,
+        full_scores: np.ndarray,
+        full_cls: np.ndarray,
         iou_thresh: float,
     ) -> 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.
+
+        The previous version omitted the same-class constraint, which let a
+        confident broom raise the score of a coincident nozzle (or vice
+        versa) under TTA. That's a silent FP booster and is fixed here.
         """
-        For each kept box, return the max original score among itself and any
-        box that overlaps it with IOU >= iou_thresh (so TTA cluster keeps best conf).
+        n = len(post_boxes)
+        if n == 0:
+            return np.empty(0, dtype=np.float32)
+        full_areas = (np.maximum(0.0, full_boxes[:, 2] - full_boxes[:, 0]) *
+                      np.maximum(0.0, full_boxes[:, 3] - full_boxes[:, 1]))
+        out = np.empty(n, dtype=np.float32)
+        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])
+            out[i] = float(np.max(full_scores[cluster])) if np.any(cluster) else 0.0
+        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).
         """
-        n_keep = len(keep_indices)
-        if n_keep == 0:
-            return np.array([], dtype=np.float32)
-        out = np.empty(n_keep, dtype=np.float32)
-        coords = np.asarray(coords, dtype=np.float32)
+        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
+
+    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]:
+        """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. For
+        car-wash this kills the common failure where water spray makes the
+        model fire both `nozzle` and `track` on the same patch, or where a
+        broom handle overlaps a drainage-gate detection.
+        """
+        n = len(boxes)
+        if n <= 1:
+            return boxes, scores, cls_ids
+        boxes = np.asarray(boxes, dtype=np.float32)
         scores = np.asarray(scores, dtype=np.float32)
-        for i in range(n_keep):
-            idx = keep_indices[i]
-            bi = coords[idx]
-            xx1 = np.maximum(bi[0], coords[:, 0])
-            yy1 = np.maximum(bi[1], coords[:, 1])
-            xx2 = np.minimum(bi[2], coords[:, 2])
-            yy2 = np.minimum(bi[3], coords[:, 3])
+        cls_ids = np.asarray(cls_ids, dtype=np.int32)
+        areas = (np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) *
+                 np.maximum(0.0, boxes[:, 3] - boxes[:, 1]))
+        margins = scores - self._conf_thres_array[cls_ids]
+        order = np.lexsort((-areas, -margins))
+        suppressed = np.zeros(n, dtype=bool)
+        keep: list[int] = []
+        for i in order:
+            if suppressed[i]:
+                continue
+            keep.append(int(i))
+            bi = boxes[i]
+            xx1 = np.maximum(bi[0], boxes[:, 0])
+            yy1 = np.maximum(bi[1], boxes[:, 1])
+            xx2 = np.minimum(bi[2], boxes[:, 2])
+            yy2 = np.minimum(bi[3], boxes[:, 3])
             inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
-            area_i = (bi[2] - bi[0]) * (bi[3] - bi[1])
-            areas_j = (coords[:, 2] - coords[:, 0]) * (coords[:, 3] - coords[:, 1])
-            iou = inter / (area_i + areas_j - inter + 1e-7)
-            in_cluster = iou >= iou_thresh
-            out[i] = float(np.max(scores[in_cluster]))
-        return out
+            a_i = max(1e-7, float((bi[2] - bi[0]) * (bi[3] - bi[1])))
+            iou = inter / (a_i + areas - inter + 1e-7)
+            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]
+
+    def _per_view_pipeline(
+        self,
+        boxes: np.ndarray,
+        scores: np.ndarray,
+        cls_ids: np.ndarray,
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """Per-view post-processing: per-class NMS -> cap -> cross-class dedup."""
+        if len(boxes) > 1:
+            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]
+            boxes, scores, cls_ids = boxes[top], scores[top], cls_ids[top]
+        if len(boxes) > 1:
+            boxes, scores, cls_ids = self._cross_class_dedup_op(
+                boxes, scores, cls_ids, self.cross_iou_thresh
+            )
+        return boxes, scores, cls_ids
 
     def _decode_final_dets(
         self,
@@ -424,7 +534,8 @@ class Miner:
         cls_ids = preds[:, 5].astype(np.int32)
         cls_ids = self.cls_remap[cls_ids]
 
-        keep = scores >= self.conf_thres
+        # Per-class confidence filter with rescue (replaces scalar threshold)
+        keep = self._conf_filter_mask(scores, cls_ids)
         boxes = boxes[keep]
         scores = scores[keep]
         cls_ids = cls_ids[keep]
@@ -448,18 +559,21 @@ class Miner:
         if len(boxes) == 0:
             return []
 
-        # Per-class NMS to remove duplicates without suppressing across classes
-        if len(boxes) > 1:
-            if apply_optional_dedup:
-                keep_idx, scores = self._per_class_soft_nms(boxes, scores, cls_ids)
-                boxes = boxes[keep_idx]
-                cls_ids = cls_ids[keep_idx]
-            else:
-                keep_idx = self._per_class_hard_nms(boxes, scores, cls_ids, self.iou_thres)
-                keep_idx = keep_idx[: self.max_det]
-                boxes = boxes[keep_idx]
-                scores = scores[keep_idx]
-                cls_ids = cls_ids[keep_idx]
+        if apply_optional_dedup and len(boxes) > 1:
+            # Soft-NMS path preserved as a tunable option; default below.
+            keep_idx, scores = self._per_class_soft_nms(boxes, scores, cls_ids)
+            boxes = boxes[keep_idx]
+            cls_ids = cls_ids[keep_idx]
+            if len(scores) > self.max_det:
+                top = np.argsort(-scores)[: self.max_det]
+                boxes, scores, cls_ids = boxes[top], scores[top], cls_ids[top]
+            if len(boxes) > 1:
+                boxes, scores, cls_ids = self._cross_class_dedup_op(
+                    boxes, scores, cls_ids, self.cross_iou_thresh
+                )
+        else:
+            # Default: per-class hard NMS -> cap -> cross-class dedup
+            boxes, scores, cls_ids = self._per_view_pipeline(boxes, scores, cls_ids)
 
         results: list[BoundingBox] = []
         for box, conf, cls_id in zip(boxes, scores, cls_ids):
@@ -520,25 +634,20 @@ class Miner:
             scores = cls_part[np.arange(len(cls_part)), cls_ids]
         cls_ids = self.cls_remap[cls_ids]
 
-        keep = scores >= self.conf_thres
+        # Per-class confidence filter with rescue (replaces scalar threshold)
+        keep = self._conf_filter_mask(scores, cls_ids)
         boxes_xywh = boxes_xywh[keep]
         scores = scores[keep]
         cls_ids = cls_ids[keep]
-
         if len(boxes_xywh) == 0:
             return []
 
         boxes = self._xywh_to_xyxy(boxes_xywh)
 
-        keep_idx = self._per_class_hard_nms(boxes, scores, cls_ids, self.iou_thres)
-        keep_idx = keep_idx[: self.max_det]
-        boxes = boxes[keep_idx]
-        scores = scores[keep_idx]
-        cls_ids = cls_ids[keep_idx]
-
+        # Order matches fire001 / _decode_final_dets:
+        # unscale -> clip -> sanity filter -> per-view pipeline (NMS, cap, cross-class dedup).
         pad_w, pad_h = pad
         orig_w, orig_h = orig_size
-
         boxes[:, [0, 2]] -= pad_w
         boxes[:, [1, 3]] -= pad_h
         boxes /= ratio
@@ -550,6 +659,8 @@ class Miner:
         if len(boxes) == 0:
             return []
 
+        boxes, scores, cls_ids = self._per_view_pipeline(boxes, scores, cls_ids)
+
         results: list[BoundingBox] = []
         for box, conf, cls_id in zip(boxes, scores, cls_ids):
             x1, y1, x2, y2 = box.tolist()
@@ -620,10 +731,16 @@ class Miner:
         return self._postprocess(det_output, ratio, pad, orig_size)
 
     def _predict_tta(self, image: np.ndarray) -> list[BoundingBox]:
-        """
-        Horizontal-flip TTA: merge original + flipped via hard NMS.
-        Boost confidence for consensus detections (both views agree) to improve
-        mAP: validator sorts by confidence, so higher conf for TP helps PR curve.
+        """Horizontal-flip TTA.
+
+        Strategy (ported from fire001):
+          1. Predict on original and on flipped image.
+          2. Map flipped boxes back to original coordinates.
+          3. Per-class hard NMS on the union.
+          4. For each kept box, compute the max SAME-CLASS score across the
+             FULL union -- a high-confidence flipped detection raises a
+             borderline original one, but never one of a different class.
+          5. Cross-class dedup to suppress same-physical-object multi-class.
         """
         boxes_orig = self._predict_single(image)
 
@@ -652,24 +769,34 @@ class Miner:
         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:
+            top = np.argsort(-scores[hard_keep])[: self.max_det]
+            hard_keep = hard_keep[top]
 
-        hard_keep = hard_keep[: self.max_det]
-
-        # Boost confidence when both views agree (overlapping detections)
+        # Class-aware cluster-max score boost (fixes the silent cross-class
+        # leak in the previous _max_score_per_cluster).
         boosted = self._max_score_per_cluster(
-            coords, scores, hard_keep, self.iou_thres
+            coords[hard_keep], cls_ids[hard_keep],
+            coords, scores, cls_ids, self.iou_thres,
         )
 
+        kept_coords = coords[hard_keep]
+        kept_cls = cls_ids[hard_keep]
+        if len(kept_coords) > 1:
+            kept_coords, boosted, kept_cls = self._cross_class_dedup_op(
+                kept_coords, boosted, kept_cls, self.cross_iou_thresh
+            )
+
         return [
             BoundingBox(
-                x1=all_boxes[i].x1,
-                y1=all_boxes[i].y1,
-                x2=all_boxes[i].x2,
-                y2=all_boxes[i].y2,
-                cls_id=all_boxes[i].cls_id,
+                x1=int(math.floor(kept_coords[j, 0])),
+                y1=int(math.floor(kept_coords[j, 1])),
+                x2=int(math.ceil(kept_coords[j, 2])),
+                y2=int(math.ceil(kept_coords[j, 3])),
+                cls_id=int(kept_cls[j]),
                 conf=float(boosted[j]),
             )
-            for j, i in enumerate(hard_keep)
+            for j in range(len(kept_coords))
         ]
 
     def predict_batch(

weights.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c41ef35650af81eab3d4d57b8605dd93544df185af1171db61b6e091c5820e3b
-size 19408005
+oid sha256:fb83bf2b7e7948721246a6137da1e64389c9198d72400b62dc18a4c91552776b
+size 19408006