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] keypoints: list[tuple[int, int]] class Miner: def __init__(self, path_hf_repo: Path) -> None: repo = Path(path_hf_repo) cfg = json.loads((repo / 'calib.json').read_text()) self.calib_x = np.asarray(cfg['calib_x'], dtype=np.float32) self.calib_y = np.asarray(cfg['calib_y'], dtype=np.float32) self.k = float(cfg.get('k', 15.0)) self.miss_prior = float(cfg.get('miss_prior', 0.0)) self.p_floor = float(cfg.get('p_floor', 0.05)) self.min_side = float(cfg.get('min_side', 2.0)) self.nms_iou = float(cfg.get('nms_iou', 0.6)) self.conf_min = float(cfg.get('conf_min', 0.02)) self.max_det = int(cfg.get('max_det', 300)) self.raw_xyxy = bool(cfg.get('raw_xyxy', True)) self.warp = tuple(cfg['warp']) if cfg.get('warp') else None so = ort.SessionOptions() so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL so.intra_op_num_threads = int(cfg.get('threads', 2)) so.inter_op_num_threads = 1 so.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL self.session = ort.InferenceSession(str(repo / 'weights.onnx'), sess_options=so, providers=['CPUExecutionProvider']) inp = self.session.get_inputs()[0] self.input_name = inp.name self.uint8_nhwc = 'uint8' in inp.type self.size = int(inp.shape[1] if self.uint8_nhwc else inp.shape[2]) dummy = np.zeros((1, self.size, self.size, 3), np.uint8) if self.uint8_nhwc else np.zeros((1, 3, self.size, self.size), np.float32) for _ in range(3): self.session.run(None, {self.input_name: dummy}) def __repr__(self) -> str: return f'PersonMiner(size={self.size}, k={self.k}, providers={self.session.get_providers()})' def _prep(self, img: ndarray): h, w = img.shape[:2] s = self.size if h == w: x = cv2.resize(img, (s, s), interpolation=cv2.INTER_LINEAR) r, left, top = (s / w, 0, 0) else: r = min(s / h, s / w) nh, nw = (int(round(h * r)), int(round(w * r))) x = np.full((s, s, 3), 114, np.uint8) top, left = ((s - nh) // 2, (s - nw) // 2) x[top:top + nh, left:left + nw] = cv2.resize(img, (nw, nh), interpolation=cv2.INTER_LINEAR) if self.uint8_nhwc: x = x[None] else: x = x[:, :, ::-1].transpose(2, 0, 1)[None].astype(np.float32) * (1.0 / 255.0) return (np.ascontiguousarray(x), r, left, top) def _warp_prep(self, img: ndarray): h, w = img.shape[:2] s = self.size hs, hd = (int(round(self.warp[0] * h)), int(round(self.warp[1] * s))) x = np.concatenate([cv2.resize(img[:hs], (s, hd), interpolation=cv2.INTER_LINEAR), cv2.resize(img[hs:], (s, s - hd), interpolation=cv2.INTER_LINEAR)], 0) return (np.ascontiguousarray(x[None]), hs, hd) def _unwarp(self, b: ndarray, w: int, h: int, hs: int, hd: int) -> ndarray: s = self.size b = b.copy() b[:, [0, 2]] *= w / s for j in (1, 3): y = b[:, j] b[:, j] = np.where(y < hd, y * hs / hd, hs + (y - hd) * (h - hs) / (s - hd)) return b def _nms_raw(self, y: ndarray, warp_args=None) -> ndarray: y = y.T sc = y[:, 4:].max(1) m = sc >= self.conf_min y, sc = (y[m], sc[m]) if sc.size == 0: return np.zeros((0, 6), np.float32) if self.raw_xyxy: b = y[:, :4].copy() else: b = np.stack([y[:, 0] - y[:, 2] / 2, y[:, 1] - y[:, 3] / 2, y[:, 0] + y[:, 2] / 2, y[:, 1] + y[:, 3] / 2], 1) if warp_args is not None: b = self._unwarp(b, *warp_args) xywh = np.concatenate([b[:, :2], b[:, 2:] - b[:, :2]], 1).astype(np.float64) k = np.asarray(cv2.dnn.NMSBoxes(xywh.tolist(), sc.astype(np.float32).tolist(), self.conf_min, self.nms_iou), dtype=int).reshape(-1)[:self.max_det] return np.concatenate([b[k], sc[k, None], np.zeros((k.size, 1), b.dtype)], 1) def _decide(self, det: ndarray, w: int, h: int, r: float, left: int, top: int) -> list[BoundingBox]: det = det[det[:, 4] > 0.0] if det.size == 0: return [] p = np.interp(det[:, 4], self.calib_x, self.calib_y).astype(np.float64) n_hat = float(p.sum()) + self.miss_prior thr = max(n_hat / (self.k + n_hat), self.p_floor) keep = p > thr det, p = (det[keep], p[keep]) order = np.argsort(-p, kind='stable') det, p = (det[order], p[order]) b = det[:, :4].copy() b[:, [0, 2]] = (b[:, [0, 2]] - left) / r b[:, [1, 3]] = (b[:, [1, 3]] - top) / r b[:, [0, 2]] = np.clip(b[:, [0, 2]], 0, w) b[:, [1, 3]] = np.clip(b[:, [1, 3]], 0, h) out = [] for (x1, y1, x2, y2), pi in zip(b, p): if x2 - x1 < self.min_side or y2 - y1 < self.min_side: continue out.append(BoundingBox(x1=int(round(x1)), y1=int(round(y1)), x2=int(round(x2)), y2=int(round(y2)), cls_id=0, conf=float(round(pi, 5)))) return out def predict_batch(self, batch_images: list[ndarray], offset: int, n_keypoints: int) -> list[TVFrameResult]: results = [] for i, img in enumerate(batch_images): h, w = img.shape[:2] if self.warp is not None and h == w and self.uint8_nhwc: x, hs, hd = self._warp_prep(img) det = self.session.run(None, {self.input_name: x})[0][0] det = self._nms_raw(det, (w, h, hs, hd)) boxes = self._decide(det, w, h, 1.0, 0, 0) else: x, r, left, top = self._prep(img) det = self.session.run(None, {self.input_name: x})[0][0] if det.shape[-1] != 6: det = self._nms_raw(det) boxes = self._decide(det, w, h, r, left, top) results.append(TVFrameResult(frame_id=offset + i, boxes=boxes, keypoints=[])) return results