miner.py

"""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)]