| 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 |
|
|
