scorevision: push artifact

miner.py

@@ -0,0 +1,97 @@
+"""SN44 public-track miner (our own). HF repo root must contain: miner.py, weights.onnx, sn44.json, chute_config.yml.
+
+- Images arrive as BGR uint8 (cv2.imdecode in the chute template).
+- Output cls_id = index into the element's manifest `objects` list (e.g. vehicle: car,bus,truck,motorcycle,van).
+  Our ONNX may have fewer classes (no 'van') -> `cls_map` in sn44.json maps onnx class -> manifest index.
+- CPU onnxruntime with 2 threads, deterministic (the conformity checker re-runs us and compares boxes by IoU).
+sn44.json: {"imgsz":1024, "conf": 0.30, "nms_iou": 0.6, "max_det": 300, "cls_map": [0], "threads": 2}
+"""
+import json
+from pathlib import Path
+
+import cv2
+import numpy as np
+import onnxruntime as ort
+from numpy import ndarray
+from pydantic import BaseModel
+
+
+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
+    keypoints: list[tuple[int, int]] | None = None
+
+
+class Miner:
+    def __init__(self, path_hf_repo: Path) -> None:
+        repo = Path(path_hf_repo)
+        cfg = json.loads((repo / "sn44.json").read_text())
+        self.size = int(cfg.get("imgsz", 1024))
+        self.conf = float(cfg.get("conf", 0.3))
+        self.conf_cls = {int(k): float(v) for k, v in (cfg.get("conf_per_class") or {}).items()}
+        self.iou = float(cfg.get("nms_iou", 0.6))
+        self.max_det = int(cfg.get("max_det", 300))
+        self.cls_map = list(cfg.get("cls_map", [0]))
+        so = ort.SessionOptions()
+        so.intra_op_num_threads = int(cfg.get("threads", 2)); so.inter_op_num_threads = 1
+        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+        self.sess = ort.InferenceSession(str(repo / "weights.onnx"), sess_options=so, providers=["CPUExecutionProvider"])
+        self.inp = self.sess.get_inputs()[0].name
+
+    def __repr__(self) -> str:
+        return f"SN44Miner(size={self.size}, conf={self.conf}, classes={len(self.cls_map)})"
+
+    def _letterbox(self, img):
+        h, w = img.shape[:2]
+        r = min(self.size / h, self.size / w)
+        nh, nw = int(round(h * r)), int(round(w * r))
+        top, left = (self.size - nh) // 2, (self.size - nw) // 2
+        canvas = np.full((self.size, self.size, 3), 114, np.uint8)
+        canvas[top:top + nh, left:left + nw] = cv2.resize(img, (nw, nh), interpolation=cv2.INTER_LINEAR)
+        x = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB).transpose(2, 0, 1)[None].astype(np.float32) / 255.0
+        return x, r, left, top
+
+    def _detect(self, img):
+        h, w = img.shape[:2]
+        x, r, left, top = self._letterbox(img)
+        y = self.sess.run(None, {self.inp: x})[0][0]  # ultralytics: (4+nc, N) cx,cy,w,h,scores
+        if y.shape[0] > y.shape[1]:
+            y = y.T
+        scores = y[4:]
+        cls = scores.argmax(0); conf = scores.max(0)
+        keep = conf >= min([self.conf] + list(self.conf_cls.values()))
+        if not keep.any():
+            return []
+        b = y[:4, keep].T; cls = cls[keep]; conf = conf[keep]
+        xyxy = np.stack([b[:, 0] - b[:, 2] / 2, b[:, 1] - b[:, 3] / 2, b[:, 0] + b[:, 2] / 2, b[:, 1] + b[:, 3] / 2], 1)
+        xyxy[:, [0, 2]] = ((xyxy[:, [0, 2]] - left) / r).clip(0, w)
+        xyxy[:, [1, 3]] = ((xyxy[:, [1, 3]] - top) / r).clip(0, h)
+        out = []
+        for c in np.unique(cls):  # class-wise NMS
+            m = cls == c
+            bb = xyxy[m]; cc = conf[m]
+            idx = cv2.dnn.NMSBoxes([[float(a), float(b_), float(c_ - a), float(d - b_)] for a, b_, c_, d in bb],
+                                   cc.astype(np.float32).tolist(), 0.0, self.iou)
+            if int(c) >= len(self.cls_map) or self.cls_map[int(c)] < 0:  # -1 = drop this onnx class
+                continue
+            for i in np.asarray(idx, dtype=int).reshape(-1):
+                if cc[i] < self.conf_cls.get(int(c), self.conf):
+                    continue
+                x1, y1, x2, y2 = bb[i]
+                out.append(BoundingBox(x1=int(round(x1)), y1=int(round(y1)), x2=int(round(x2)), y2=int(round(y2)),
+                                       cls_id=int(self.cls_map[int(c)]), conf=float(round(float(cc[i]), 5))))
+        out.sort(key=lambda b: -b.conf)
+        return out[: self.max_det]
+
+    def predict_batch(self, batch_images: list[ndarray], offset: int, n_keypoints: int) -> list[TVFrameResult]:
+        return [TVFrameResult(frame_id=offset + i, boxes=self._detect(img), keypoints=[(0, 0)] * n_keypoints)
+                for i, img in enumerate(batch_images)]