Upload folder using huggingface_hub

miner.py

@@ -27,14 +27,17 @@ class Miner:
     def __init__(self,
         path_hf_repo: Path
     ) -> None:
-        model_path = path_hf_repo / "weights.onnx"
+        model_path = self._resolve_model_path(path_hf_repo)
         # car-wash element classes — cls_id order MUST match element `objects`
-        # (0=broom, 1=drainage gate, 2=nozzle, 3=track) and the YOLO training order.
+        # (0=broom, 1=drainage gate, 2=nozzle, 3=track). This is the canonical
+        # order every downstream consumer (validator, BoundingBox.cls_id) sees.
         self.class_names = ["broom", "drainage gate", "nozzle", "track"]
-        model_class_order = ["broom", "drainage gate", "nozzle", "track"]
-        self.cls_remap = np.array(
-            [self.class_names.index(n) for n in model_class_order], dtype=np.int32
-        )
+        # FALLBACK model-emit order: the authoritative order is read from the
+        # ONNX `names` metadata after the session is created (embedded by
+        # Ultralytics at export, ships inside weights.onnx), so a retrained
+        # model with a different class order is remapped correctly without
+        # code changes. This list is used only when metadata is missing.
+        self._model_class_order = ["broom", "drainage gate", "nozzle", "track"]
         print("ORT version:", ort.__version__)
 
         try:
@@ -65,6 +68,20 @@ class Miner:
 
         print("ORT session providers:", self.session.get_providers())
 
+        # Build cls_remap: for each model-emit index i,
+        #   cls_remap[i] = self.class_names.index(model_class_order[i])
+        # The model-side order comes from the ONNX metadata when available,
+        # else falls back to the static _model_class_order.
+        model_class_order = self._read_model_class_order()
+        if model_class_order is None:
+            model_class_order = list(self._model_class_order)
+            print(f"cls order: no usable ONNX metadata, FALLBACK {model_class_order}")
+        else:
+            print(f"cls order: from ONNX metadata {model_class_order}")
+        self.cls_remap = np.array(
+            [self.class_names.index(n) for n in model_class_order], dtype=np.int32
+        )
+
         for inp in self.session.get_inputs():
             print("INPUT:", inp.name, inp.shape, inp.type)
 
@@ -85,29 +102,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).
-        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)
+        # All values below are the measured optimum of a full grid sweep on
+        # the validator-style val split (tune_miner.py, 241 1024x1024 crops,
+        # composite 0.8002 -> 0.8103) -- re-run the sweep after any retrain.
+        self.iou_thres = 0.45        # Per-class NMS IoU; lower = stricter dedup
+        self.cross_iou_thresh = 0.9  # 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.
+#       conf thresholds:  broom=0.38  drainage gate=0.45  nozzle=0.30  track=0.60
+        # Per-class confidence thresholds.
         # 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
+        # broom/nozzle sit low: under the validator metric the mAP gained
+        # from the extra recall outweighs the FP-pillar cost (the previous
+        # 0.5/0.5 silently discarded many valid detections); track is the
+        # one class where false fires are common enough to need 0.38.
         self._conf_thres_array = np.array(
-            [0.5, 0.4, 0.5, 0.23], dtype=np.float32
+            [0.37, 0.23, 0.37, 0.45], 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.
+        # is at least (per-class threshold - per-class bonus).
+        # DISABLED (all zeros): the sweep showed rescue admits more false
+        # positives than true positives under the validator's FP pillar.
         self._bonus_array = np.array(
-            [0.15, 0.1, 0, 0.05], dtype=np.float32
+            [0.05, 0.05, 0.0, 0.15], dtype=np.float32
         )
 
         # Box sanity filter — kept loose: car-wash `nozzle` boxes are tiny
@@ -131,6 +150,55 @@ class Miner:
     def _safe_dim(value, default: int) -> int:
         return value if isinstance(value, int) and value > 0 else default
 
+    @staticmethod
+    def _resolve_model_path(repo: Path) -> Path:
+        """Locate the ONNX model in the repo dir.
+
+        Prefers weights.onnx (FP16/FP32 export), then weights_int8.onnx (the
+        training script's INT8-quantized export -- works as-is: quantization
+        preserves the Ultralytics metadata and QDQ models take regular fp32
+        input), then any other .onnx file. INT8 is the fallback when the FP16
+        export exceeds the 30 MB deployment limit (e.g. yolo26m).
+        """
+        for name in ("weights.onnx", "weights_int8.onnx"):
+            p = repo / name
+            if p.exists():
+                if name != "weights.onnx":
+                    print(f"model: weights.onnx not found, using {name}")
+                return p
+        candidates = sorted(repo.glob("*.onnx"))
+        if candidates:
+            print(f"model: using {candidates[0].name}")
+            return candidates[0]
+        return repo / "weights.onnx"  # let session creation raise the error
+
+    def _read_model_class_order(self) -> list[str] | None:
+        """Read the model's class order from Ultralytics ONNX metadata.
+
+        Returns the class names ordered by model-emit index, or None when
+        metadata is missing/unparsable or doesn't match `class_names` as a
+        set (in which case the static _model_class_order fallback is used).
+        """
+        try:
+            import ast
+
+            meta = self.session.get_modelmeta().custom_metadata_map
+            names = ast.literal_eval(meta["names"])  # e.g. {0: 'broom', ...}
+            if isinstance(names, dict):
+                order = [str(names[i]) for i in sorted(names)]
+            else:
+                order = [str(n) for n in names]
+        except Exception as e:
+            print(f"cls order: could not read ONNX names metadata ({e})")
+            return None
+        if sorted(order) != sorted(self.class_names):
+            print(
+                f"cls order: ONNX names {order} do not match expected classes "
+                f"{self.class_names}; ignoring metadata"
+            )
+            return None
+        return order
+
     def _letterbox(
         self,
         image: ndarray,

weights.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:9d5cc4c83c9b5a943d33a678a44839854645b639c8c4d1c68acf04906f8a0058
-size 19408004
+oid sha256:1f9c13ee4403ddbbccd9d9707a151b1474cbb33124a2ef2965fa57d1e821ed03
+size 19287011