File size: 4,569 Bytes
295eca3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | """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)]
|