Add miner.py

miner.py

@@ -0,0 +1,303 @@
+import ast
+import math
+import os
+from pathlib import Path
+
+import cv2
+import numpy as np
+import onnxruntime as ort
+from numpy import ndarray
+from pydantic import BaseModel
+
+# manak0/Detect-Vehicle manifest order
+SUBNET_LABELS: tuple[str, ...] = ("car", "bus", "truck", "motorcycle", "van")
+SUBNET_INDEX: dict[str, int] = {name: i for i, name in enumerate(SUBNET_LABELS)}
+
+# Chosen against this checkpoint on 1024 challenge frames, not copied from another miner.
+# Truck is the only class with a real false-positive tail. Motorcycle precision is
+# already saturated at 0.25, so raising it throws away recall.
+CONF_THRES = np.array([0.25, 0.25, 0.35, 0.25, 0.25], dtype=np.float32)
+NMS_IOU = 0.50
+# Same physical vehicle emitted as two classes. Measured overlaps were ~0.98.
+CROSS_CLASS_IOU = 0.70
+
+
+def _normalize_label(name: str) -> str:
+    return "".join(ch for ch in name.lower() if ch.isalnum())
+
+
+class BoundingBox(BaseModel):
+    x1: int
+    y1: int
+    x2: int
+    y2: int
+    cls_id: int
+    conf: float
+
+
+class TVFrameResult(BaseModel):
+    frame_id: int
+    boxes: list[BoundingBox] | None = None
+    polygons: list[list] | None = None
+    keypoints: list[tuple[int, int]] | None = None
+
+
+class Miner:
+    """Detect-Vehicle miner. cls_id: 0=car, 1=bus, 2=truck, 3=motorcycle, 4=van.
+
+    Challenge frames are 1024x1024. The ONNX file is exported smaller so the
+    CPU compliance rerun stays under 100 ms; square frames are resized
+    uniformly and boxes are scaled back. The published weights have no van class.
+    """
+
+    def __init__(self, path_hf_repo: Path) -> None:
+        model_path = self._resolve_model_path(Path(path_hf_repo))
+        sess_options = ort.SessionOptions()
+        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+        sess_options.intra_op_num_threads = 2
+        sess_options.inter_op_num_threads = 1
+        sess_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
+
+        force_cpu = os.environ.get("MINER_FORCE_CPU", "") == "1"
+        use_cuda = (not force_cpu) and ("CUDAExecutionProvider" in ort.get_available_providers())
+        if use_cuda:
+            try:
+                ort.preload_dlls()
+            except Exception as exc:
+                print(f"preload_dlls skipped: {exc}")
+            providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+        else:
+            providers = ["CPUExecutionProvider"]
+        self.session = ort.InferenceSession(str(model_path), sess_options=sess_options, providers=providers)
+
+        order = self._read_model_class_order() or ["car", "bus", "truck", "motorcycle"]
+        remap = [SUBNET_INDEX.get(_normalize_label(name), -1) for name in order]
+        if all(idx < 0 for idx in remap):
+            raise ValueError(f"No vehicle classes in ONNX names: {order}")
+        self.cls_remap = np.asarray(remap, dtype=np.int32)
+
+        self.input_name = self.session.get_inputs()[0].name
+        shape = self.session.get_inputs()[0].shape
+        self.input_height = shape[2] if isinstance(shape[2], int) and shape[2] > 0 else 640
+        self.input_width = shape[3] if isinstance(shape[3], int) and shape[3] > 0 else 640
+        print(
+            f"ONNX {model_path.name} input={self.input_width}x{self.input_height} "
+            f"classes={order} providers={self.session.get_providers()}"
+        )
+        self._warmup()
+
+    def __repr__(self) -> str:
+        return (
+            f"VehicleONNX({self.input_width}x{self.input_height}, "
+            f"classes={self.cls_remap.tolist()}, providers={self.session.get_providers()})"
+        )
+
+    def _warmup(self) -> None:
+        dummy = np.zeros((self.input_height, self.input_width, 3), dtype=np.uint8)
+        self.predict_batch([dummy], offset=0, n_keypoints=0)
+
+    @staticmethod
+    def _resolve_model_path(repo: Path) -> Path:
+        for name in ("weights.onnx", "vehicle.onnx"):
+            path = repo / name
+            if path.is_file():
+                return path
+        found = sorted(repo.glob("*.onnx"))
+        if found:
+            return found[0]
+        raise FileNotFoundError(f"No ONNX weights in {repo}")
+
+    def _read_model_class_order(self) -> list[str] | None:
+        try:
+            raw = self.session.get_modelmeta().custom_metadata_map["names"]
+            names = ast.literal_eval(raw)
+        except Exception as exc:
+            print(f"cls order: metadata unreadable ({exc})")
+            return None
+        if isinstance(names, dict):
+            return [str(names[key]) for key in sorted(names, key=lambda item: int(item))]
+        if isinstance(names, (list, tuple)):
+            return [str(name) for name in names]
+        return None
+
+    def _preprocess(self, image: ndarray) -> tuple[np.ndarray, float, float, float]:
+        """Letterbox into the ONNX input. Square 1024 frames have no side padding."""
+        height, width = image.shape[:2]
+        if image.dtype != np.uint8:
+            image = image.astype(np.uint8)
+        new_w, new_h = self.input_width, self.input_height
+        ratio = min(new_w / width, new_h / height)
+        resized_w, resized_h = int(round(width * ratio)), int(round(height * ratio))
+        if (resized_w, resized_h) != (width, height):
+            image = cv2.resize(image, (resized_w, resized_h), interpolation=cv2.INTER_LINEAR)
+        pad_w = (new_w - resized_w) / 2.0
+        pad_h = (new_h - resized_h) / 2.0
+        left, right = int(round(pad_w - 0.1)), int(round(pad_w + 0.1))
+        top, bottom = int(round(pad_h - 0.1)), int(round(pad_h + 0.1))
+        padded = cv2.copyMakeBorder(image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114))
+        pad_h_px, pad_w_px = padded.shape[:2]
+        if pad_w_px != new_w or pad_h_px != new_h:
+            padded = cv2.copyMakeBorder(
+                padded,
+                0,
+                max(0, new_h - pad_h_px),
+                0,
+                max(0, new_w - pad_w_px),
+                cv2.BORDER_CONSTANT,
+                value=(114, 114, 114),
+            )
+            padded = padded[:new_h, :new_w]
+        blob = cv2.dnn.blobFromImage(padded, scalefactor=1.0 / 255.0, swapRB=True, crop=False)
+        return blob, ratio, pad_w, pad_h
+
+    @staticmethod
+    def _nms(boxes: np.ndarray, scores: np.ndarray, iou_thresh: float) -> np.ndarray:
+        if len(boxes) == 0:
+            return np.empty(0, dtype=np.intp)
+        x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
+        areas = np.maximum(0.0, x2 - x1) * np.maximum(0.0, y2 - y1)
+        order = np.argsort(-scores)
+        keep: list[int] = []
+        while order.size:
+            current = int(order[0])
+            keep.append(current)
+            if order.size == 1:
+                break
+            rest = order[1:]
+            xx1 = np.maximum(x1[current], x1[rest])
+            yy1 = np.maximum(y1[current], y1[rest])
+            xx2 = np.minimum(x2[current], x2[rest])
+            yy2 = np.minimum(y2[current], y2[rest])
+            inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1)
+            iou = inter / (areas[current] + areas[rest] - inter + 1e-7)
+            order = rest[iou <= iou_thresh]
+        return np.asarray(keep, dtype=np.intp)
+
+    def _decode(
+        self,
+        output: np.ndarray,
+        ratio: float,
+        pad_w: float,
+        pad_h: float,
+        width: int,
+        height: int,
+    ) -> list[BoundingBox]:
+        preds = output[0] if output.ndim == 3 else output
+        if preds.shape[0] < preds.shape[1]:
+            preds = preds.T
+        class_scores = preds[:, 4:].astype(np.float32)
+        raw_cls = np.argmax(class_scores, axis=1).astype(np.int32)
+        scores = class_scores[np.arange(len(class_scores)), raw_cls]
+        known = (raw_cls >= 0) & (raw_cls < len(self.cls_remap))
+        mapped = np.full(len(raw_cls), -1, dtype=np.int32)
+        mapped[known] = self.cls_remap[raw_cls[known]]
+        keep = (mapped >= 0) & (scores >= CONF_THRES[np.clip(mapped, 0, len(CONF_THRES) - 1)])
+        keep &= mapped >= 0
+        if not np.any(keep):
+            return []
+
+        xywh = preds[keep, :4].astype(np.float32)
+        scores = scores[keep]
+        cls_ids = mapped[keep]
+        boxes = np.empty_like(xywh)
+        boxes[:, 0] = xywh[:, 0] - xywh[:, 2] / 2.0
+        boxes[:, 1] = xywh[:, 1] - xywh[:, 3] / 2.0
+        boxes[:, 2] = xywh[:, 0] + xywh[:, 2] / 2.0
+        boxes[:, 3] = xywh[:, 1] + xywh[:, 3] / 2.0
+        boxes[:, [0, 2]] = (boxes[:, [0, 2]] - pad_w) / ratio
+        boxes[:, [1, 3]] = (boxes[:, [1, 3]] - pad_h) / ratio
+        np.clip(boxes[:, 0], 0, width - 1, out=boxes[:, 0])
+        np.clip(boxes[:, 1], 0, height - 1, out=boxes[:, 1])
+        np.clip(boxes[:, 2], 0, width - 1, out=boxes[:, 2])
+        np.clip(boxes[:, 3], 0, height - 1, out=boxes[:, 3])
+
+        kept: list[int] = []
+        for cls_id in np.unique(cls_ids):
+            idx = np.where(cls_ids == cls_id)[0]
+            chosen = self._nms(boxes[idx], scores[idx], NMS_IOU)
+            kept.extend(idx[chosen].tolist())
+        if not kept:
+            return []
+        order = np.asarray(kept, dtype=np.intp)
+        order = order[np.argsort(-scores[order])]
+        boxes, scores, cls_ids = boxes[order], scores[order], cls_ids[order]
+        boxes, scores, cls_ids = self._drop_cross_class_duplicates(boxes, scores, cls_ids)
+
+        results: list[BoundingBox] = []
+        for box, score, cls_id in zip(boxes, scores, cls_ids):
+            x1, y1, x2, y2 = box.tolist()
+            x1_i = max(0, min(width, int(math.floor(x1))))
+            y1_i = max(0, min(height, int(math.floor(y1))))
+            x2_i = max(0, min(width, int(math.ceil(x2))))
+            y2_i = max(0, min(height, int(math.ceil(y2))))
+            if x2_i <= x1_i or y2_i <= y1_i:
+                continue
+            results.append(
+                BoundingBox(
+                    x1=x1_i,
+                    y1=y1_i,
+                    x2=x2_i,
+                    y2=y2_i,
+                    cls_id=int(cls_id),
+                    conf=float(max(0.0, min(1.0, score))),
+                )
+            )
+        return results
+
+    @staticmethod
+    def _drop_cross_class_duplicates(
+        boxes: np.ndarray, scores: np.ndarray, cls_ids: np.ndarray
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        n = len(boxes)
+        if n < 2:
+            return boxes, scores, cls_ids
+        drop = np.zeros(n, dtype=bool)
+        for i in range(n):
+            if drop[i]:
+                continue
+            for j in range(i + 1, n):
+                if drop[j] or cls_ids[i] == cls_ids[j]:
+                    continue
+                xx1 = max(float(boxes[i, 0]), float(boxes[j, 0]))
+                yy1 = max(float(boxes[i, 1]), float(boxes[j, 1]))
+                xx2 = min(float(boxes[i, 2]), float(boxes[j, 2]))
+                yy2 = min(float(boxes[i, 3]), float(boxes[j, 3]))
+                inter = max(0.0, xx2 - xx1) * max(0.0, yy2 - yy1)
+                if inter <= 0.0:
+                    continue
+                area_i = max(0.0, float(boxes[i, 2] - boxes[i, 0])) * max(0.0, float(boxes[i, 3] - boxes[i, 1]))
+                area_j = max(0.0, float(boxes[j, 2] - boxes[j, 0])) * max(0.0, float(boxes[j, 3] - boxes[j, 1]))
+                iou = inter / (area_i + area_j - inter + 1e-7)
+                if iou >= CROSS_CLASS_IOU:
+                    drop[j if scores[j] < scores[i] else i] = True
+        keep = ~drop
+        return boxes[keep], scores[keep], cls_ids[keep]
+
+    def predict_batch(
+        self,
+        batch_images: list[ndarray],
+        offset: int,
+        n_keypoints: int,
+    ) -> list[TVFrameResult]:
+        n_keypoints = max(0, int(n_keypoints))
+        results: list[TVFrameResult] = []
+        for index, image in enumerate(batch_images):
+            try:
+                if not isinstance(image, np.ndarray) or image.ndim != 3 or image.shape[2] != 3:
+                    raise ValueError(f"Expected HWC BGR image, got {getattr(image, 'shape', None)}")
+                blob, ratio, pad_w, pad_h = self._preprocess(image)
+                height, width = image.shape[:2]
+                output = self.session.run(None, {self.input_name: blob})[0]
+                boxes = self._decode(output, ratio, pad_w, pad_h, width, height)
+            except Exception as exc:
+                print(f"Inference failed for frame {offset + index}: {exc}")
+                boxes = []
+            results.append(
+                TVFrameResult(
+                    frame_id=offset + index,
+                    boxes=boxes,
+                    polygons=[],
+                    keypoints=[(0, 0) for _ in range(n_keypoints)],
+                )
+            )
+        return results