update

README.md

@@ -0,0 +1,3 @@
+---
+license: mit
+---

chute_config.yml

@@ -2,20 +2,16 @@ Image:
   from_base: parachutes/python:3.12
   run_command:
     - pip install --upgrade setuptools wheel
-    - pip install 'numpy>=1.23' 'onnxruntime-gpu>=1.16' 'opencv-python>=4.7' 'pillow>=9.5' 'huggingface_hub>=0.19.4' 'pydantic>=2.0' 'pyyaml>=6.0' 'aiohttp>=3.9'
-    - pip install torch==2.8.0 torchvision==0.23.0 torchaudio==2.8.0 --index-url https://download.pytorch.org/whl/cu128
+    - pip install 'numpy>=1.23' 'onnxruntime-gpu[cuda,cudnn]>=1.16' 'opencv-python>=4.7' 'pillow>=9.5' 'huggingface_hub>=0.19.4' 'pydantic>=2.0' 'pyyaml>=6.0' 'aiohttp>=3.9'
+    - pip install torch torchvision
+  set_workdir: /app
 
 NodeSelector:
   gpu_count: 1
-  min_vram_gb_per_gpu: 16
-  max_hourly_price_per_gpu: 0.5
-
-  # Required for integrated SN44 chutes (TEE policy enforced as of 2026-04-27).
   include:
     - pro_6000
 
 Chute:
-  # Required for integrated SN44 chutes (TEE policy enforced as of 2026-04-27).
   tee: true
   timeout_seconds: 900
   concurrency: 4

class_names.txt

@@ -1,4 +1,4 @@
 petrol hose
 petrol pump
 price board
-roof canopy
+roof canopy

miner.py

@@ -1,3 +1,4 @@
+
 from pathlib import Path
 import math
 
@@ -22,32 +23,18 @@ class TVFrameResult(BaseModel):
     boxes: list[BoundingBox]
     keypoints: list[tuple[int, int]]
 
-SIZE = 1280
-TARGET_CLASS_NAMES = ["petrol hose", "petrol pump", "price board", "roof canopy"]
-
 
 class Miner:
+    """ONNX-backed petrol-tracking miner with canopy union-merge post-process."""
+
+    CANOPY_CLS = 3
+
     def __init__(self, path_hf_repo: Path) -> None:
-        model_path = path_hf_repo / "weights.onnx"
-        cn_path = model_path.with_name("class_names.txt")
-        self.class_names = TARGET_CLASS_NAMES.copy()
-        if cn_path.is_file():
-            lines = cn_path.read_text(encoding="utf-8").splitlines()
-            model_class_order = [
-                ln.strip()
-                for ln in lines
-                if ln.strip() and not ln.strip().startswith("#")
-            ]
-            if len(model_class_order) == len(self.class_names) and set(model_class_order) == set(self.class_names):
-                self.cls_remap = np.array(
-                    [self.class_names.index(n) for n in model_class_order], dtype=np.int32
-                )
-            else:
-                # If class_names.txt is missing/invalid for this target order, keep identity mapping.
-                self.cls_remap = np.arange(len(self.class_names), dtype=np.int32)
-        else:
-            # Fallback when no class_names.txt is present: assume ONNX class order == target order.
-            self.cls_remap = np.arange(len(self.class_names), dtype=np.int32)
+        model_path = path_hf_repo / "petrol.onnx"
+
+        # Class order as exported from the training pt: must match model.names
+        self.class_names = ["petrol hose", "petrol pump", "price board", "roof canopy"]
+
         print("ORT version:", ort.__version__)
 
         try:
@@ -88,21 +75,26 @@ class Miner:
         self.output_names = [output.name for output in self.session.get_outputs()]
         self.input_shape = self.session.get_inputs()[0].shape
 
-        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.input_height = self._safe_dim(self.input_shape[2], default=640)
+        self.input_width = self._safe_dim(self.input_shape[3], default=640)
+
+        # Thresholds
+        self.conf_thres = 0.38
+        self.iou_thres = 0.50
+        self.max_det = 300
 
-        self.conf_thres = 0.5
-        self.iou_thres = 0.45
-        self.max_det = 30
-        self.use_tta = True
+        # Canopy union-merge: same-class IoU above this triggers a union merge
+        # for class 3 only (roof canopy). Set to 0 to disable.
+        self.canopy_merge_iou = 0.30
 
-        print(f"✅ ONNX model loaded from: {model_path}")
+        print(f"✅ Petrol ONNX model loaded from: {model_path}")
         print(f"✅ ONNX providers: {self.session.get_providers()}")
         print(f"✅ ONNX input: name={self.input_name}, shape={self.input_shape}")
+        print(f"✅ Canopy merge IoU: {self.canopy_merge_iou}")
 
     def __repr__(self) -> str:
         return (
-            f"ONNXRuntime(session={type(self.session).__name__}, "
+            f"Petrol ONNXRuntime(session={type(self.session).__name__}, "
             f"providers={self.session.get_providers()})"
         )
 
@@ -116,13 +108,6 @@ class Miner:
         new_shape: tuple[int, int],
         color=(114, 114, 114),
     ) -> tuple[ndarray, float, tuple[float, float]]:
-        """
-        Resize with unchanged aspect ratio and pad to target shape.
-        Returns:
-            padded_image,
-            ratio,
-            (pad_w, pad_h)  # half-padding
-        """
         h, w = image.shape[:2]
         new_w, new_h = new_shape
 
@@ -158,14 +143,6 @@ class Miner:
     def _preprocess(
         self, image: ndarray
     ) -> tuple[np.ndarray, float, tuple[float, float], tuple[int, int]]:
-        """
-        Preprocess for fixed-size ONNX export:
-        - enhance image quality (CLAHE, denoise, sharpen)
-        - letterbox to model input size
-        - BGR -> RGB
-        - normalize to [0,1]
-        - HWC -> NCHW float32
-        """
         orig_h, orig_w = image.shape[:2]
 
         img, ratio, pad = self._letterbox(
@@ -196,230 +173,195 @@ class Miner:
         out[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
         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])
-            )
-            iou = inter / (area_i + areas_j - inter + 1e-7)
-            scores[i + 1:] *= np.exp(-(iou ** 2) / sigma)
+            area_i = max(0.0, (boxes[i, 2] - boxes[i, 0])) * max(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])
 
-        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(
+        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:
-        """
-        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:
+        if len(boxes) == 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]:
+        keep_all: list[int] = []
+        for c in np.unique(cls_ids):
+            idxs = np.nonzero(cls_ids == c)[0]
+            if len(idxs) == 0:
                 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)
+            local_keep = cls._hard_nms(boxes[idxs], scores[idxs], iou_thresh)
+            keep_all.extend(idxs[local_keep].tolist())
+        keep_all_arr = np.array(keep_all, dtype=np.intp)
+        order = np.argsort(scores[keep_all_arr])[::-1]
+        return keep_all_arr[order[:max_det]]
 
     @staticmethod
-    def _max_score_per_cluster(
-        coords: np.ndarray,
+    def _pairwise_iou(boxes: np.ndarray) -> np.ndarray:
+        """N×N IoU matrix for an [N,4] xyxy array."""
+        n = len(boxes)
+        if n == 0:
+            return np.zeros((0, 0), dtype=np.float32)
+        x1 = boxes[:, 0]; y1 = boxes[:, 1]
+        x2 = boxes[:, 2]; y2 = boxes[:, 3]
+        area = np.maximum(0.0, x2 - x1) * np.maximum(0.0, y2 - y1)
+
+        ix1 = np.maximum(x1[:, None], x1[None, :])
+        iy1 = np.maximum(y1[:, None], y1[None, :])
+        ix2 = np.minimum(x2[:, None], x2[None, :])
+        iy2 = np.minimum(y2[:, None], y2[None, :])
+        iw = np.maximum(0.0, ix2 - ix1)
+        ih = np.maximum(0.0, iy2 - iy1)
+        inter = iw * ih
+        union = area[:, None] + area[None, :] - inter
+        with np.errstate(divide="ignore", invalid="ignore"):
+            iou = np.where(union > 0, inter / union, 0.0)
+        np.fill_diagonal(iou, 0.0)
+        return iou.astype(np.float32)
+
+    def _union_merge_class(
+        self,
+        boxes: np.ndarray,
         scores: np.ndarray,
-        keep_indices: np.ndarray,
-        iou_thresh: float,
-    ) -> np.ndarray:
+        cls_ids: np.ndarray,
+        target_cls: int,
+        merge_iou: float,
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """Greedy union-merge for one class.
+
+        For boxes whose cls == target_cls, repeatedly fuse pairs whose IoU
+        exceeds `merge_iou`: replace them with the bounding-rectangle union
+        (max conf). Other classes are passed through unchanged.
         """
-        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
-
-    def _decode_final_dets(
+        if merge_iou <= 0 or len(boxes) == 0:
+            return boxes, scores, cls_ids
+
+        mask = cls_ids == target_cls
+        if mask.sum() < 2:
+            return boxes, scores, cls_ids
+
+        tgt_boxes = boxes[mask].astype(np.float32).copy()
+        tgt_scores = scores[mask].astype(np.float32).copy()
+
+        # Greedy merge: highest-conf box anchors each round; absorb all
+        # others above the IoU threshold; repeat until stable.
+        changed = True
+        while changed and len(tgt_boxes) > 1:
+            changed = False
+            order = np.argsort(tgt_scores)[::-1]
+            tgt_boxes = tgt_boxes[order]
+            tgt_scores = tgt_scores[order]
+
+            iou = self._pairwise_iou(tgt_boxes)
+            consumed = np.zeros(len(tgt_boxes), dtype=bool)
+            new_boxes: list[np.ndarray] = []
+            new_scores: list[float] = []
+            for i in range(len(tgt_boxes)):
+                if consumed[i]:
+                    continue
+                cur = tgt_boxes[i].copy()
+                cur_s = float(tgt_scores[i])
+                for j in range(i + 1, len(tgt_boxes)):
+                    if consumed[j]:
+                        continue
+                    if iou[i, j] > merge_iou:
+                        cur = np.array([
+                            min(cur[0], tgt_boxes[j, 0]),
+                            min(cur[1], tgt_boxes[j, 1]),
+                            max(cur[2], tgt_boxes[j, 2]),
+                            max(cur[3], tgt_boxes[j, 3]),
+                        ], dtype=np.float32)
+                        cur_s = max(cur_s, float(tgt_scores[j]))
+                        consumed[j] = True
+                        changed = True
+                new_boxes.append(cur)
+                new_scores.append(cur_s)
+            tgt_boxes = np.stack(new_boxes, axis=0)
+            tgt_scores = np.array(new_scores, dtype=np.float32)
+
+        # Stitch results back together with non-target classes
+        other_boxes = boxes[~mask]
+        other_scores = scores[~mask]
+        other_cls = cls_ids[~mask]
+
+        merged_cls = np.full(len(tgt_boxes), target_cls, dtype=cls_ids.dtype)
+        out_boxes = np.concatenate([other_boxes, tgt_boxes], axis=0)
+        out_scores = np.concatenate([other_scores, tgt_scores], axis=0)
+        out_cls = np.concatenate([other_cls, merged_cls], axis=0)
+        return out_boxes, out_scores, out_cls
+
+    def _decode_yolov8(
         self,
         preds: np.ndarray,
         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.
-        """
-        if preds.ndim == 3 and preds.shape[0] == 1:
-            preds = preds[0]
-
-        if preds.ndim != 2 or preds.shape[1] < 6:
-            raise ValueError(f"Unexpected ONNX final-det output shape: {preds.shape}")
-
-        boxes = preds[:, :4].astype(np.float32)
-        scores = preds[:, 4].astype(np.float32)
-        cls_ids = preds[:, 5].astype(np.int32)
-        cls_ids = self.cls_remap[np.clip(cls_ids, 0, len(self.cls_remap) - 1)]
-
-        keep = scores >= self.conf_thres
-        boxes = boxes[keep]
-        scores = scores[keep]
-        cls_ids = cls_ids[keep]
-
-        if len(boxes) == 0:
-            return []
-
-        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]
-
-        results: list[BoundingBox] = []
-        for box, conf, cls_id in zip(boxes, scores, cls_ids):
-            x1, y1, x2, y2 = box.tolist()
-
-            if x2 <= x1 or y2 <= y1:
-                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(cls_id),
-                    conf=float(conf),
-                )
-            )
-
-        return results
+        Decode a raw YOLOv8-style ONNX detection output.
 
-    def _decode_raw_yolo(
-        self,
-        preds: np.ndarray,
-        ratio: float,
-        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]
+        Expected shape: [1, 4 + nc, num_boxes] (no objectness channel).
+        Some exporters emit [1, num_boxes, 4 + nc]; both are handled.
         """
-        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}")
+        if preds.ndim != 3 or preds.shape[0] != 1:
+            raise ValueError(f"Unexpected ONNX output shape: {preds.shape}")
 
         preds = preds[0]
 
-        # Normalize to [N, C]
-        if preds.shape[0] <= 16 and preds.shape[1] > preds.shape[0]:
+        # Normalize to [N, C] where C = 4 + nc
+        nc = len(self.class_names)
+        expected_c = 4 + nc
+        if preds.shape[0] == expected_c:
             preds = preds.T
+        elif preds.shape[1] != expected_c:
+            # Fall back: treat smaller dim as channels
+            if preds.shape[0] < preds.shape[1]:
+                preds = preds.T
 
         if preds.ndim != 2 or preds.shape[1] < 5:
-            raise ValueError(f"Unexpected normalized raw output shape: {preds.shape}")
+            raise ValueError(f"Unexpected normalized output shape: {preds.shape}")
 
         boxes_xywh = preds[:, :4].astype(np.float32)
-        cls_part = preds[:, 4:].astype(np.float32)
+        class_probs = preds[:, 4:].astype(np.float32)
 
-        if cls_part.shape[1] == 1:
-            scores = cls_part[:, 0]
-            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]
-        cls_ids = self.cls_remap[np.clip(cls_ids, 0, len(self.cls_remap) - 1)]
+        cls_ids = np.argmax(class_probs, axis=1).astype(np.int32)
+        scores = class_probs[np.arange(len(class_probs)), cls_ids]
 
         keep = scores >= self.conf_thres
         boxes_xywh = boxes_xywh[keep]
@@ -430,12 +372,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
@@ -445,47 +381,33 @@ 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()
-
-            if x2 <= x1 or y2 <= y1:
-                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(cls_id),
-                    conf=float(conf),
-                )
-            )
-
-        return results
+        keep_idx = self._nms_per_class(
+            boxes, scores, cls_ids, self.iou_thres, self.max_det
+        )
 
-    def _postprocess(
-        self,
-        output: np.ndarray,
-        ratio: float,
-        pad: tuple[float, float],
-        orig_size: tuple[int, int],
-    ) -> list[BoundingBox]:
-        """
-        Prefer final detections first.
-        Fallback to raw decode only if needed.
-        """
-        # final detections: [N,6]
-        if output.ndim == 2 and output.shape[1] >= 6:
-            return self._decode_final_dets(output, ratio, pad, orig_size)
+        boxes = boxes[keep_idx]
+        scores = scores[keep_idx]
+        cls_ids = cls_ids[keep_idx]
 
-        # final detections: [1,N,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)
+        # Class-3 union-merge: rejoin half-canopy splits into one box.
+        boxes, scores, cls_ids = self._union_merge_class(
+            boxes, scores, cls_ids,
+            target_cls=self.CANOPY_CLS,
+            merge_iou=self.canopy_merge_iou,
+        )
 
-        # fallback raw decode
-        return self._decode_raw_yolo(output, ratio, pad, orig_size)
+        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 _predict_single(self, image: np.ndarray) -> list[BoundingBox]:
         if image is None:
@@ -512,51 +434,7 @@ class Miner:
 
         outputs = self.session.run(self.output_names, {self.input_name: input_tensor})
         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:
-            return []
-
-        coords = np.array(
-            [[b.x1, b.y1, b.x2, b.y2] for b in all_boxes], dtype=np.float32
-        )
-        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 []
-
-        # _hard_nms already orders kept indices by descending score.
-        hard_keep = hard_keep[: self.max_det]
-
-        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
-        ]
+        return self._decode_yolov8(det_output, ratio, pad, orig_size)
 
     def predict_batch(
         self,
@@ -564,81 +442,113 @@ class Miner:
         offset: int,
         n_keypoints: int,
     ) -> list[TVFrameResult]:
+        """
+        Miner prediction for a batch of images using ONNX Runtime.
+
+        The petrol detector is a plain object-detection model (no pose),
+        so keypoints are returned as `n_keypoints` padding entries of (0, 0)
+        to keep the TVFrameResult schema stable across challenge types.
+        """
         results: list[TVFrameResult] = []
+        n_kp = max(0, int(n_keypoints))
 
         for frame_number_in_batch, image in enumerate(batch_images):
+            frame_idx = offset + frame_number_in_batch
             try:
-                if self.use_tta:
-                    boxes = self._predict_tta(image)
-                else:
-                    boxes = self._predict_single(image)
+                boxes = self._predict_single(image)
             except Exception as e:
-                print(f"⚠️ Inference failed for frame {offset + frame_number_in_batch}: {e}")
+                print(f"⚠️ Inference failed for frame {frame_idx}: {e}")
                 boxes = []
-            # for box in boxes:
-            #     if box.cls_id == 2:
-            #         box.cls_id = 3
-            #     elif box.cls_id == 3:
-            #         box.cls_id = 2
-                    
-            
-            
+
             results.append(
                 TVFrameResult(
-                    frame_id=offset + frame_number_in_batch,
+                    frame_id=frame_idx,
                     boxes=boxes,
-                    keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))],
+                    keypoints=[(0, 0) for _ in range(n_kp)],
                 )
             )
 
+        print("✅ Petrol ONNX predictions complete")
         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)]
+
+def main() -> None:
+    """Example runner — same CLI as miner.py for direct A/B comparison."""
+    import sys
+
+    repo_path = Path(__file__).parent
+    print(f"Loading miner_v2 from: {repo_path}")
+    miner = Miner(path_hf_repo=repo_path)
+    print(repr(miner))
+
+    batch_images: list[np.ndarray] = []
+
+    if len(sys.argv) > 1:
+        for image_path in sys.argv[1:]:
+            image = cv2.imread(image_path)
+            if image is None:
+                raise ValueError(f"Cannot read image: {image_path}")
+            batch_images.append(image)
+        print(f"Loaded {len(batch_images)} image(s)")
+    else:
+        batch_images = [np.zeros((640, 640, 3), dtype=np.uint8)]
+        print("No image provided — running on a single blank dummy frame")
+
+    results = miner.predict_batch(
+        batch_images=batch_images,
+        offset=0,
+        n_keypoints=32,
+    )
+
+    output_dir = repo_path / "predictions_v2"
+    output_dir.mkdir(exist_ok=True)
+
+    class_names = {i: n for i, n in enumerate(miner.class_names)}
+
+    def color_for_class(cls_id: int) -> tuple[int, int, int]:
+        hue = (cls_id * 47) % 180
+        hsv = np.uint8([[[hue, 220, 255]]])
+        bgr = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)[0, 0]
+        return int(bgr[0]), int(bgr[1]), int(bgr[2])
+
+    for image, r in zip(batch_images, results):
+        print(
+            f"frame={r.frame_id} "
+            f"boxes={len(r.boxes)} "
+            f"keypoints={len(r.keypoints)}"
+        )
+
+        vis = image.copy()
+        for box in r.boxes:
+            name = class_names.get(box.cls_id, str(box.cls_id))
+            color = color_for_class(box.cls_id)
+            print(
+                f"  box cls={box.cls_id}({name}) conf={box.conf:.2f} "
+                f"[{box.x1},{box.y1},{box.x2},{box.y2}]"
+            )
+            cv2.rectangle(vis, (box.x1, box.y1), (box.x2, box.y2), color, 2)
+            label = f"{name} {box.conf:.2f}"
+            (tw, th), baseline = cv2.getTextSize(
+                label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1
+            )
+            top = max(box.y1 - th - baseline, 0)
             cv2.rectangle(
-                vis,
-                (box.x1, box.y1),
-                (box.x2, box.y2),
-                color,
-                2,
+                vis, (box.x1, top), (box.x1 + tw, top + th + baseline), color, -1
             )
-            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,
+                vis, label, (box.x1, top + th),
+                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 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}")
+        for x, y in r.keypoints:
+            if x == 0 and y == 0:
+                continue
+            cv2.circle(vis, (x, y), 3, (0, 0, 255), -1)
+
+        out_path = output_dir / f"frame_{r.frame_id:04d}.jpg"
+        cv2.imwrite(str(out_path), vis)
+        print(f"  saved: {out_path}")
+
+
+if __name__ == "__main__":
+    main()

weights.onnx → petrol.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:bcb98181cc317a8a83ffc13a390cc9b1ba37ef27c6cb55c05eade472ebdc0df2
-size 10264218
+oid sha256:bb8ff9dbe935b06f64e6049b0604c2c871386b633b36ef9d320d0e02e5f35c36
+size 22664875