Upload folder using huggingface_hub

chute_config.yml

@@ -8,7 +8,7 @@ Image:
 NodeSelector:
   gpu_count: 1
   min_vram_gb_per_gpu: 16
-  max_hourly_price_per_gpu: 0.5
+  max_hourly_price_per_gpu: 1.0
 
   exclude:
     - "5090"

miner.py

@@ -6,8 +6,6 @@ import numpy as np
 import onnxruntime as ort
 from numpy import ndarray
 from pydantic import BaseModel
-import argparse
-import json
 
 
 class BoundingBox(BaseModel):
@@ -26,20 +24,45 @@ class TVFrameResult(BaseModel):
 
 SIZE = 1280
 
+# --- Class labels (edit here; cls_id 0..N-1 matches this order) ---
+CLASS_NAMES: tuple[str, ...] = (
+    "petrol hose",
+    "petrol pump",
+    "price board",
+    "roof canopy",
+)
+
+# If the ONNX model outputs class indices in a different order than CLASS_NAMES, set this
+# to the same names in that model order. None = identity.
+MODEL_CLASS_ORDER: tuple[str, ...] | None = None
+
+# --- Per-class confidence (edit here) ---
+# Same order as CLASS_NAMES. Use empty tuple () to use scalar defaults (conf_thres, conf_high).
+PER_CLASS_CONF_THRES = (0.25, 0.42, 0.32, 0.45)
+PER_CLASS_CONF_HIGH = (0.56, 0.62, 0.52, 0.6)
+
 
 class Miner:
     def __init__(self, path_hf_repo: Path) -> None:
         model_path = path_hf_repo / "weights.onnx"
-        cn_path = model_path.with_name("class_names.txt")
-        if cn_path.is_file():
-            lines = cn_path.read_text(encoding="utf-8").splitlines()
-            self.class_names = [
-                ln.strip()
-                for ln in lines
-                if ln.strip() and not ln.strip().startswith("#")
-            ]
+
+        self.class_names = list(CLASS_NAMES)
+        if MODEL_CLASS_ORDER is None:
+            self._train_cls_to_canonical = np.arange(
+                len(self.class_names), dtype=np.int32
+            )
         else:
-            self.class_names = ["petrol_hose", "petrol_pump", "price board", "roof canopy"]
+            if set(MODEL_CLASS_ORDER) != set(self.class_names) or len(
+                MODEL_CLASS_ORDER
+            ) != len(self.class_names):
+                raise ValueError(
+                    "MODEL_CLASS_ORDER must be a permutation of CLASS_NAMES "
+                    "(names in the order the ONNX model outputs cls indices)."
+                )
+            self._train_cls_to_canonical = np.array(
+                [self.class_names.index(n) for n in MODEL_CLASS_ORDER], dtype=np.int32
+            )
+
         print("ORT version:", ort.__version__)
 
         try:
@@ -83,14 +106,28 @@ class Miner:
         self.input_height = self._safe_dim(self.input_shape[2], default=SIZE)
         self.input_width = self._safe_dim(self.input_shape[3], default=SIZE)
 
-        self.conf_thres = 0.49
-        self.iou_thres = 0.28
-        self.max_det = 100
+        # --- VehicleDetection/vehicle3 scoring-oriented thresholds ---
+        self.conf_thres = 0.25  # low threshold for candidate generation
+        self.conf_high = 0.5  # high-conf boxes can survive without TTA confirmation
+        self.iou_thres = 0.50
+        self.tta_match_iou = 0.6  # TTA agreement IoU
+        self.max_det = 150
         self.use_tta = True
-        self.tile_size = SIZE
-        self.overlap = 0.5
-        self.use_slicer = True
-        self.use_full_image_merge = True
+
+        n_cls = len(self.class_names)
+        self._conf_thres_per_class = self._per_class_vector(
+            n_cls, PER_CLASS_CONF_THRES, self.conf_thres, "PER_CLASS_CONF_THRES"
+        )
+        self._conf_high_per_class = self._per_class_vector(
+            n_cls, PER_CLASS_CONF_HIGH, self.conf_high, "PER_CLASS_CONF_HIGH"
+        )
+
+        # Box sanity (VehicleDetection/vehicle3)
+        self.min_box_area = 12 * 12
+        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()}")
@@ -106,6 +143,38 @@ class Miner:
     def _safe_dim(value, default: int) -> int:
         return value if isinstance(value, int) and value > 0 else default
 
+    @staticmethod
+    def _per_class_vector(
+        n_cls: int,
+        per_class: tuple[float, ...],
+        scalar: float,
+        name: str,
+    ) -> np.ndarray:
+        """Build length-`n_cls` vector from a tuple or broadcast `scalar` if tuple is empty."""
+        if not per_class:
+            return np.full(n_cls, scalar, dtype=np.float32)
+        if len(per_class) != n_cls:
+            raise ValueError(
+                f"{name}: expected {n_cls} values (same order as CLASS_NAMES), got {len(per_class)}"
+            )
+        return np.array(per_class, dtype=np.float32)
+
+    def _remap_train_cls_ids(self, cls_ids: np.ndarray) -> np.ndarray:
+        idx = np.clip(
+            cls_ids.astype(np.int64, copy=False),
+            0,
+            len(self._train_cls_to_canonical) - 1,
+        )
+        return self._train_cls_to_canonical[idx]
+
+    def _clip_cls_id(self, cls_id: int) -> int:
+        n = len(self.class_names)
+        if cls_id < 0:
+            return 0
+        if cls_id >= n:
+            return n - 1
+        return cls_id
+
     def _letterbox(
         self,
         image: ndarray,
@@ -192,182 +261,131 @@ class Miner:
         out[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
         return out
 
-    def _slice_image(
-        self, image: np.ndarray
-    ) -> list[tuple[np.ndarray, tuple[int, int], tuple[int, int]]]:
-        h, w = image.shape[:2]
-        t = self.tile_size
-        st = max(1, int(t * (1.0 - self.overlap)))
-
-        xs = []
-        x = 0
-        while True:
-            if x + t >= w:
-                xs.append(max(0, w - t))
-                break
-            xs.append(x)
-            x += st
-
-        ys = []
-        y = 0
-        while True:
-            if y + t >= h:
-                ys.append(max(0, h - t))
-                break
-            ys.append(y)
-            y += st
-
-        xs = list(dict.fromkeys(xs))
-        ys = list(dict.fromkeys(ys))
-
-        out = []
-        for y0 in ys:
-            for x0 in xs:
-                x1 = min(x0 + t, w)
-                y1 = min(y0 + t, h)
-                crop = image[y0:y1, x0:x1]
-                vh, vw = crop.shape[:2]
-                out.append((crop, (x0, y0), (vw, vh)))
-        return out
-
-    def _soft_nms(
-        self,
+    @staticmethod
+    def _hard_nms(
         boxes: np.ndarray,
         scores: np.ndarray,
-        sigma: float = 0.5,
-        score_thresh: float = 0.01,
-    ) -> tuple[np.ndarray, np.ndarray]:
-        """
-        Soft-NMS: Gaussian decay of overlapping scores instead of hard removal.
-        Returns (kept_original_indices, updated_scores).
-        """
-        N = len(boxes)
-        if N == 0:
-            return np.array([], dtype=np.intp), np.array([], dtype=np.float32)
-
-        boxes = boxes.astype(np.float32, copy=True)
-        scores = scores.astype(np.float32, copy=True)
-        order = np.arange(N)
+        iou_thresh: float,
+    ) -> np.ndarray:
+        if len(boxes) == 0:
+            return np.array([], dtype=np.intp)
 
-        for i in range(N):
-            max_pos = i + int(np.argmax(scores[i:]))
-            boxes[[i, max_pos]] = boxes[[max_pos, i]]
-            scores[[i, max_pos]] = scores[[max_pos, i]]
-            order[[i, max_pos]] = order[[max_pos, i]]
+        boxes = np.asarray(boxes, dtype=np.float32)
+        scores = np.asarray(scores, dtype=np.float32)
+        order = np.argsort(scores)[::-1]
+        keep = []
 
-            if i + 1 >= N:
+        while len(order) > 0:
+            i = order[0]
+            keep.append(i)
+            if len(order) == 1:
                 break
 
-            xx1 = np.maximum(boxes[i, 0], boxes[i + 1:, 0])
-            yy1 = np.maximum(boxes[i, 1], boxes[i + 1:, 1])
-            xx2 = np.minimum(boxes[i, 2], boxes[i + 1:, 2])
-            yy2 = np.minimum(boxes[i, 3], boxes[i + 1:, 3])
+            rest = order[1:]
+
+            xx1 = np.maximum(boxes[i, 0], boxes[rest, 0])
+            yy1 = np.maximum(boxes[i, 1], boxes[rest, 1])
+            xx2 = np.minimum(boxes[i, 2], boxes[rest, 2])
+            yy2 = np.minimum(boxes[i, 3], boxes[rest, 3])
+
             inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
 
-            area_i = max(0.0, float(
-                (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1])
-            ))
-            areas_j = (
-                np.maximum(0.0, boxes[i + 1:, 2] - boxes[i + 1:, 0])
-                * np.maximum(0.0, boxes[i + 1:, 3] - boxes[i + 1:, 1])
+            area_i = np.maximum(0.0, (boxes[i, 2] - boxes[i, 0])) * np.maximum(
+                0.0, (boxes[i, 3] - boxes[i, 1])
+            )
+            area_r = np.maximum(0.0, (boxes[rest, 2] - boxes[rest, 0])) * np.maximum(
+                0.0, (boxes[rest, 3] - boxes[rest, 1])
             )
-            iou = inter / (area_i + areas_j - inter + 1e-7)
-            scores[i + 1:] *= np.exp(-(iou ** 2) / sigma)
 
-        mask = scores > score_thresh
-        return order[mask], scores[mask]
+            iou = inter / (area_i + area_r - inter + 1e-7)
+            order = rest[iou <= iou_thresh]
 
-    @staticmethod
-    def _hard_nms(
-        boxes: np.ndarray,
-        scores: np.ndarray,
-        iou_thresh: float,
-    ) -> np.ndarray:
-        """
-        Standard NMS: keep one box per overlapping cluster (the one with highest score).
-        Returns indices of kept boxes (into the boxes/scores arrays).
-        """
-        N = len(boxes)
-        if N == 0:
-            return np.array([], dtype=np.intp)
-        boxes = np.asarray(boxes, dtype=np.float32)
-        scores = np.asarray(scores, dtype=np.float32)
-        order = np.argsort(scores)[::-1]
-        keep: list[int] = []
-        suppressed = np.zeros(N, dtype=bool)
-        for i in range(N):
-            idx = order[i]
-            if suppressed[idx]:
-                continue
-            keep.append(idx)
-            bi = boxes[idx]
-            for k in range(i + 1, N):
-                jdx = order[k]
-                if suppressed[jdx]:
-                    continue
-                bj = boxes[jdx]
-                xx1 = max(bi[0], bj[0])
-                yy1 = max(bi[1], bj[1])
-                xx2 = min(bi[2], bj[2])
-                yy2 = min(bi[3], bj[3])
-                inter = max(0.0, xx2 - xx1) * max(0.0, yy2 - yy1)
-                area_i = (bi[2] - bi[0]) * (bi[3] - bi[1])
-                area_j = (bj[2] - bj[0]) * (bj[3] - bj[1])
-                iou = inter / (area_i + area_j - inter + 1e-7)
-                if iou > iou_thresh:
-                    suppressed[jdx] = True
-        return np.array(keep)
-
-    def _hard_nms_by_class(
-        self,
+        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; then global top-`max_det` by score (VehicleDetection/vehicle3)."""
         if len(boxes) == 0:
             return np.array([], dtype=np.intp)
         keep_all: list[int] = []
         for c in np.unique(cls_ids):
-            m = cls_ids == c
-            inds = np.flatnonzero(m)
-            sub_keep = self._hard_nms(boxes[m], scores[m], iou_thresh)
-            keep_all.extend(int(inds[i]) for i in sub_keep)
-        keep_all = np.asarray(keep_all, dtype=np.intp)
-        order = np.argsort(scores[keep_all])[::-1][: self.max_det]
-        return keep_all[order]
+            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 _max_score_per_cluster(
-        coords: np.ndarray,
+    def _box_iou_one_to_many(box: np.ndarray, boxes: np.ndarray) -> np.ndarray:
+        xx1 = np.maximum(box[0], boxes[:, 0])
+        yy1 = np.maximum(box[1], boxes[:, 1])
+        xx2 = np.minimum(box[2], boxes[:, 2])
+        yy2 = np.minimum(box[3], boxes[:, 3])
+
+        inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
+
+        area_a = max(0.0, (box[2] - box[0]) * (box[3] - box[1]))
+        area_b = np.maximum(0.0, boxes[:, 2] - boxes[:, 0]) * np.maximum(
+            0.0, boxes[:, 3] - boxes[:, 1]
+        )
+
+        return inter / (area_a + area_b - inter + 1e-7)
+
+    def _filter_sane_boxes(
+        self,
+        boxes: np.ndarray,
         scores: np.ndarray,
-        keep_indices: np.ndarray,
-        iou_thresh: float,
-    ) -> np.ndarray:
-        """
-        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_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)
-        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])
-            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
+        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 = []
+        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),
+            )
+
+        keep = np.array(keep, dtype=np.intp)
+        return boxes[keep], scores[keep], cls_ids[keep]
 
     def _decode_final_dets(
         self,
@@ -375,12 +393,9 @@ class Miner:
         ratio: float,
         pad: tuple[float, float],
         orig_size: tuple[int, int],
-        apply_optional_dedup: bool = False,
     ) -> list[BoundingBox]:
         """
-        Primary path:
-        expected output rows like [x1, y1, x2, y2, conf, cls_id]
-        in letterboxed input coordinates.
+        Primary path: rows like [x1, y1, x2, y2, conf, cls_id] in letterboxed coords.
         """
         if preds.ndim == 3 and preds.shape[0] == 1:
             preds = preds[0]
@@ -390,9 +405,10 @@ class Miner:
 
         boxes = preds[:, :4].astype(np.float32)
         scores = preds[:, 4].astype(np.float32)
-        cls_ids = preds[:, 5].astype(np.int32)
+        cls_ids = self._remap_train_cls_ids(preds[:, 5].astype(np.int32))
 
-        keep = scores >= self.conf_thres
+        ci = np.clip(cls_ids.astype(np.int64), 0, len(self._conf_thres_per_class) - 1)
+        keep = scores >= self._conf_thres_per_class[ci]
         boxes = boxes[keep]
         scores = scores[keep]
         cls_ids = cls_ids[keep]
@@ -403,36 +419,35 @@ class Miner:
         pad_w, pad_h = pad
         orig_w, orig_h = orig_size
 
-        # reverse letterbox
         boxes[:, [0, 2]] -= pad_w
         boxes[:, [1, 3]] -= pad_h
         boxes /= ratio
         boxes = self._clip_boxes(boxes, (orig_w, orig_h))
 
-        if apply_optional_dedup and len(boxes) > 1:
-            keep_idx, scores = self._soft_nms(boxes, scores)
-            boxes = boxes[keep_idx]
-            cls_ids = cls_ids[keep_idx]
+        boxes, scores, cls_ids = self._filter_sane_boxes(boxes, scores, cls_ids, orig_size)
+        if len(boxes) == 0:
+            return []
 
-        results: list[BoundingBox] = []
-        for box, conf, cls_id in zip(boxes, scores, cls_ids):
-            x1, y1, x2, y2 = box.tolist()
+        keep_idx = self._nms_per_class(
+            boxes, scores, cls_ids, self.iou_thres, self.max_det
+        )
 
-            if x2 <= x1 or y2 <= y1:
-                continue
+        boxes = boxes[keep_idx]
+        scores = scores[keep_idx]
+        cls_ids = cls_ids[keep_idx]
 
-            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(cls_id),
-                    conf=float(conf),
-                )
+        return [
+            BoundingBox(
+                x1=int(math.floor(box[0])),
+                y1=int(math.floor(box[1])),
+                x2=int(math.ceil(box[2])),
+                y2=int(math.ceil(box[3])),
+                cls_id=int(cls_id),
+                conf=float(conf),
             )
-
-        return results
+            for box, conf, cls_id in zip(boxes, scores, cls_ids)
+            if box[2] > box[0] and box[3] > box[1]
+        ]
 
     def _decode_raw_yolo(
         self,
@@ -441,21 +456,13 @@ class Miner:
         pad: tuple[float, float],
         orig_size: tuple[int, int],
     ) -> list[BoundingBox]:
-        """
-        Fallback path for raw YOLO predictions.
-        Supports common layouts:
-        - [1, C, N]
-        - [1, N, C]
-        """
         if preds.ndim != 3:
             raise ValueError(f"Unexpected raw ONNX output shape: {preds.shape}")
-
         if preds.shape[0] != 1:
             raise ValueError(f"Unexpected batch dimension in raw output: {preds.shape}")
 
         preds = preds[0]
 
-        # Normalize to [N, C]
         if preds.shape[0] <= 16 and preds.shape[1] > preds.shape[0]:
             preds = preds.T
 
@@ -463,16 +470,27 @@ class Miner:
             raise ValueError(f"Unexpected normalized raw output shape: {preds.shape}")
 
         boxes_xywh = preds[:, :4].astype(np.float32)
-        cls_part = preds[:, 4:].astype(np.float32)
+        tail = preds[:, 4:].astype(np.float32)
 
-        if cls_part.shape[1] == 1:
-            scores = cls_part[:, 0]
+        if tail.shape[1] == 1:
+            scores = tail[:, 0]
+            cls_ids = np.zeros(len(scores), dtype=np.int32)
+        elif tail.shape[1] == 2:
+            obj = tail[:, 0]
+            cls_prob = tail[:, 1]
+            scores = obj * cls_prob
             cls_ids = np.zeros(len(scores), dtype=np.int32)
         else:
-            cls_ids = np.argmax(cls_part, axis=1).astype(np.int32)
-            scores = cls_part[np.arange(len(cls_part)), cls_ids]
+            obj = tail[:, 0]
+            class_probs = tail[:, 1:]
+            cls_ids = np.argmax(class_probs, axis=1).astype(np.int32)
+            cls_scores = class_probs[np.arange(len(class_probs)), cls_ids]
+            scores = obj * cls_scores
 
-        keep = scores >= self.conf_thres
+        cls_ids = self._remap_train_cls_ids(cls_ids)
+
+        ci = np.clip(cls_ids.astype(np.int64), 0, len(self._conf_thres_per_class) - 1)
+        keep = scores >= self._conf_thres_per_class[ci]
         boxes_xywh = boxes_xywh[keep]
         scores = scores[keep]
         cls_ids = cls_ids[keep]
@@ -481,12 +499,6 @@ class Miner:
             return []
 
         boxes = self._xywh_to_xyxy(boxes_xywh)
-        keep_idx, scores = self._soft_nms(boxes, scores)
-        keep_idx = keep_idx[: self.max_det]
-        scores = scores[: self.max_det]
-
-        boxes = boxes[keep_idx]
-        cls_ids = cls_ids[keep_idx]
 
         pad_w, pad_h = pad
         orig_w, orig_h = orig_size
@@ -496,25 +508,30 @@ class Miner:
         boxes /= ratio
         boxes = self._clip_boxes(boxes, (orig_w, orig_h))
 
-        results: list[BoundingBox] = []
-        for box, conf, cls_id in zip(boxes, scores, cls_ids):
-            x1, y1, x2, y2 = box.tolist()
+        boxes, scores, cls_ids = self._filter_sane_boxes(boxes, scores, cls_ids, orig_size)
+        if len(boxes) == 0:
+            return []
 
-            if x2 <= x1 or y2 <= y1:
-                continue
+        keep_idx = self._nms_per_class(
+            boxes, scores, cls_ids, self.iou_thres, self.max_det
+        )
 
-            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(cls_id),
-                    conf=float(conf),
-                )
-            )
+        boxes = boxes[keep_idx]
+        scores = scores[keep_idx]
+        cls_ids = cls_ids[keep_idx]
 
-        return results
+        return [
+            BoundingBox(
+                x1=int(math.floor(box[0])),
+                y1=int(math.floor(box[1])),
+                x2=int(math.ceil(box[2])),
+                y2=int(math.ceil(box[3])),
+                cls_id=int(cls_id),
+                conf=float(conf),
+            )
+            for box, conf, cls_id in zip(boxes, scores, cls_ids)
+            if box[2] > box[0] and box[3] > box[1]
+        ]
 
     def _postprocess(
         self,
@@ -531,8 +548,8 @@ class Miner:
         if output.ndim == 2 and output.shape[1] >= 6:
             return self._decode_final_dets(output, ratio, pad, orig_size)
 
-        # final detections: [1,N,6]
-        if output.ndim == 3 and output.shape[0] == 1 and output.shape[2] == 6:
+        # final detections: [1,N,C] with C>=6
+        if output.ndim == 3 and output.shape[0] == 1 and output.shape[2] >= 6:
             return self._decode_final_dets(output, ratio, pad, orig_size)
 
         # fallback raw decode
@@ -565,130 +582,137 @@ class Miner:
         det_output = outputs[0]
         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."""
-        boxes_orig = self._predict_single(image)
-
-        flipped = cv2.flip(image, 1)
-        boxes_flip = self._predict_single(flipped)
-
-        w = image.shape[1]
-        boxes_flip = [
-            BoundingBox(
-                x1=w - b.x2, y1=b.y1, x2=w - b.x1, y2=b.y2,
-                cls_id=b.cls_id, conf=b.conf,
-            )
-            for b in boxes_flip
-        ]
-
-        all_boxes = boxes_orig + boxes_flip
-        if len(all_boxes) == 0:
+    def _merge_tta_consensus(
+        self,
+        boxes_orig: list[BoundingBox],
+        boxes_flip: list[BoundingBox],
+    ) -> list[BoundingBox]:
+        """
+        VehicleDetection/vehicle3 strategy:
+        - keep any original-view box with conf >= conf_high
+        - keep lower-conf original boxes only if confirmed in flipped view (IoU)
+        - add flipped high-conf boxes the original view missed
+        - final per-class NMS
+        """
+        if not boxes_orig and not boxes_flip:
             return []
 
-        coords = np.array(
-            [[b.x1, b.y1, b.x2, b.y2] for b in all_boxes], dtype=np.float32
+        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)
         )
-        scores = np.array([b.conf for b in all_boxes], dtype=np.float32)
 
-        hard_keep = self._hard_nms(coords, scores, self.iou_thres)
-        if len(hard_keep) == 0:
-            return []
+        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)
+        )
 
-        # _hard_nms already orders kept indices by descending score.
-        hard_keep = hard_keep[: self.max_det]
+        accepted_boxes = []
+        accepted_scores = []
+        accepted_cls = []
+
+        for i in range(len(coords_o)):
+            score = scores_o[i]
+            c = self._clip_cls_id(int(cls_o[i]))
+            ch = float(self._conf_high_per_class[c])
+            if score >= ch:
+                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))
+                if ious[j] >= self.tta_match_iou:
+                    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]))
+
+        for i in range(len(coords_f)):
+            score = scores_f[i]
+            c = self._clip_cls_id(int(cls_f[i]))
+            if score < float(self._conf_high_per_class[c]):
+                continue
 
-        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,
-                conf=float(scores[i]),
-            )
-            for i in hard_keep
-        ]
+            if len(coords_o) == 0:
+                accepted_boxes.append(coords_f[i])
+                accepted_scores.append(score)
+                accepted_cls.append(int(cls_f[i]))
+                continue
 
-    def predict_image(self, image: np.ndarray) -> list[BoundingBox]:
-        if image is None:
-            raise ValueError("Input image is None")
-        if not isinstance(image, np.ndarray):
-            raise TypeError(f"Input is not numpy array: {type(image)}")
-        if image.ndim != 3 or image.shape[2] != 3:
-            raise ValueError(f"Expected HWC image with 3 channels, got shape={image.shape}")
-
-        H, W = image.shape[:2]
-        all_boxes: list[list[float]] = []
-        all_scores: list[float] = []
-        all_cls: list[int] = []
-
-        if self.use_slicer:
-            tiles = self._slice_image(image)
-            for tile_img, (ox, oy), (vw, vh) in tiles:
-                try:
-                    dets = self._predict_tta(tile_img) if self.use_tta else self._predict_single(tile_img)
-                except Exception as e:
-                    print(f"⚠️ Tile inference failed at ({ox}, {oy}): {e}")
-                    continue
-
-                left_edge = ox == 0
-                top_edge = oy == 0
-                right_edge = (ox + vw) >= W
-                bottom_edge = (oy + vh) >= H
-
-                for b in dets:
-                    bw = b.x2 - b.x1
-                    bh = b.y2 - b.y1
-                    m = max(8, int(min(bw, bh) * 0.2))
-                    if not left_edge and b.x1 < m:
-                        continue
-                    if not top_edge and b.y1 < m:
-                        continue
-                    if not right_edge and b.x2 > (vw - m):
-                        continue
-                    if not bottom_edge and b.y2 > (vh - m):
-                        continue
-
-                    x1 = max(0, min(W - 1, int(b.x1 + ox)))
-                    y1 = max(0, min(H - 1, int(b.y1 + oy)))
-                    x2 = max(0, min(W - 1, int(b.x2 + ox)))
-                    y2 = max(0, min(H - 1, int(b.y2 + oy)))
-                    if x2 > x1 and y2 > y1:
-                        all_boxes.append([x1, y1, x2, y2])
-                        all_scores.append(float(b.conf))
-                        all_cls.append(int(b.cls_id))
-
-        if self.use_full_image_merge or not self.use_slicer:
-            full_dets = self._predict_tta(image) if self.use_tta else self._predict_single(image)
-            for b in full_dets:
-                if b.x2 > b.x1 and b.y2 > b.y1:
-                    all_boxes.append([b.x1, b.y1, b.x2, b.y2])
-                    all_scores.append(float(b.conf))
-                    all_cls.append(int(b.cls_id))
-
-        if not all_boxes:
+            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.asarray(all_boxes, dtype=np.float32)
-        scores = np.asarray(all_scores, dtype=np.float32)
-        cls_ids = np.asarray(all_cls, dtype=np.int32)
-        keep = self._hard_nms_by_class(boxes, scores, cls_ids, self.iou_thres)
+        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)
 
-        out: list[BoundingBox] = []
-        for i in keep:
-            b = boxes[i]
+        out = []
+        for idx in keep:
+            x1, y1, x2, y2 = boxes[idx].tolist()
             out.append(
                 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])),
-                    cls_id=int(cls_ids[i]),
-                    conf=float(scores[i]),
+                    x1=int(math.floor(x1)),
+                    y1=int(math.floor(y1)),
+                    x2=int(math.ceil(x2)),
+                    y2=int(math.ceil(y2)),
+                    cls_id=int(cls_ids[idx]),
+                    conf=float(scores[idx]),
                 )
             )
         return out
 
+    def _predict_tta(self, image: np.ndarray) -> list[BoundingBox]:
+        boxes_orig = self._predict_single(image)
+
+        flipped = cv2.flip(image, 1)
+        boxes_flip_raw = self._predict_single(flipped)
+
+        w = image.shape[1]
+        boxes_flip = [
+            BoundingBox(
+                x1=w - b.x2,
+                y1=b.y1,
+                x2=w - b.x1,
+                y2=b.y2,
+                cls_id=b.cls_id,
+                conf=b.conf,
+            )
+            for b in boxes_flip_raw
+        ]
+
+        return self._merge_tta_consensus(boxes_orig, boxes_flip)
+
     def predict_batch(
         self,
         batch_images: list[ndarray],
@@ -699,7 +723,10 @@ class Miner:
 
         for frame_number_in_batch, image in enumerate(batch_images):
             try:
-                boxes = self.predict_image(image)
+                if self.use_tta:
+                    boxes = self._predict_tta(image)
+                else:
+                    boxes = self._predict_single(image)
             except Exception as e:
                 print(f"⚠️ Inference failed for frame {offset + frame_number_in_batch}: {e}")
                 boxes = []
@@ -713,3 +740,53 @@ class Miner:
             )
 
         return results
+    
+
+if __name__ == "__main__":
+    # Simple manual test: load weights.onnx, run on 1.png, and draw bboxes
+    repo_dir = Path(__file__).parent
+    miner = Miner(repo_dir)
+
+    image_path = repo_dir / "car1.png"
+    if not image_path.exists():
+        raise FileNotFoundError(f"Test image not found: {image_path}")
+
+    image = cv2.imread(str(image_path), cv2.IMREAD_COLOR)
+    if image is None:
+        raise RuntimeError(f"Failed to read image: {image_path}")
+
+    results = miner.predict_batch([image], offset=0, n_keypoints=0)
+    # Draw bounding boxes on a copy of the image
+    vis = image.copy()
+    colors = [(0, 255, 0), (0, 0, 255), (255, 0, 0)]
+    for frame in results:
+        print(f"Frame {frame.frame_id}:")
+        for i, box in enumerate(frame.boxes):
+            color = colors[i % len(colors)]
+            cv2.rectangle(
+                vis,
+                (box.x1, box.y1),
+                (box.x2, box.y2),
+                color,
+                2,
+            )
+            label = f"{box.cls_id }_{miner.class_names[box.cls_id] if box.cls_id < len(miner.class_names) else box.cls_id}:{box.conf:.2f}"
+            cv2.putText(
+                vis,
+                label,
+                (box.x1, max(0, box.y1 - 5)),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                box.conf,
+                color,
+                1,
+                cv2.LINE_AA,
+            )
+            print(
+                f"  cls={box.cls_id} conf={box.conf:.3f} "
+                f"box=({box.x1},{box.y1},{box.x2},{box.y2})"
+            )
+        print(len(frame.boxes))
+
+    out_path = repo_dir / f"1_out_iou{miner.iou_thres:.2f}.png"
+    cv2.imwrite(str(out_path), vis)
+    print(f"Saved visualization to: {out_path}")