from pathlib import Path import math import os import cv2 import numpy as np import onnxruntime as ort from numpy import ndarray from pydantic import BaseModel ELEMENT_CONFIG = { "element_id": "person", "classes": ["person"], "conf_floor": 0.30, "nms_iou": 0.4, "max_det": 75, "min_side": 3.0, "max_aspect": 10.0, "max_area_frac": 0.8, "ios": None, } CONF_POLICY = {"kind": "step", "steps": ((30, 0.40), (60, 0.50), (None, 0.60))} def dynamic_conf(n_raw: int) -> float: p = CONF_POLICY if p["kind"] == "lerp": if n_raw <= p["n_lo"]: return p["lo"] if n_raw >= p["n_hi"]: return p["hi"] return p["lo"] + (n_raw - p["n_lo"]) / (p["n_hi"] - p["n_lo"]) * (p["hi"] - p["lo"]) for limit, conf in p["steps"]: if limit is None or n_raw < limit: return conf return p["steps"][-1][1] 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: model_path = path_hf_repo / "weights.onnx" self.class_names = list(ELEMENT_CONFIG["classes"]) self.cfg = ELEMENT_CONFIG so = ort.SessionOptions() so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL so.intra_op_num_threads = min(2, len(os.sched_getaffinity(0))) so.inter_op_num_threads = 1 so.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL so.log_severity_level = 3 self.session = ort.InferenceSession(str(model_path), sess_options=so, providers=["CPUExecutionProvider"]) inp = self.session.get_inputs()[0] out = self.session.get_outputs()[0] if len(out.shape) != 3 or out.shape[1] != 5: raise ValueError(f"expected raw YOLO output [1,5,N], got {out.shape}") self.input_name = inp.name self.output_names = [out.name] self.input_height = int(inp.shape[2]) self.input_width = int(inp.shape[3]) self.last_n_raw = 0 self.last_conf = 0.0 self._warmup(2) def _warmup(self, iters: int) -> None: dummy = np.zeros((1024, 1024, 3), dtype=np.uint8) for _ in range(iters): self.predict_batch(batch_images=[dummy], offset=0, n_keypoints=0) @staticmethod def _letterbox(image: ndarray, new_shape: tuple[int, int], color=(114, 114, 114)): h, w = image.shape[:2] new_w, new_h = new_shape ratio = min(new_w / w, new_h / h) rw, rh = int(round(w * ratio)), int(round(h * ratio)) if (rw, rh) != (w, h): image = cv2.resize(image, (rw, rh), interpolation=cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR) dw, dh = (new_w - rw) / 2.0, (new_h - rh) / 2.0 left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) padded = cv2.copyMakeBorder(image, top, bottom, left, right, borderType=cv2.BORDER_CONSTANT, value=color) return padded, ratio, (dw, dh) def _preprocess(self, image: ndarray): padded, ratio, pad = self._letterbox(image, (self.input_width, self.input_height)) blob = cv2.dnn.blobFromImage(padded, scalefactor=1 / 255.0, swapRB=True) return np.ascontiguousarray(blob, dtype=np.float32), ratio, pad def _sane(self, boxes: np.ndarray, scores: np.ndarray, orig_w: int, orig_h: int): c = self.cfg bw, bh = boxes[:, 2] - boxes[:, 0], boxes[:, 3] - boxes[:, 1] ar = np.maximum(bw / np.maximum(bh, 1e-6), bh / np.maximum(bw, 1e-6)) keep = (bw >= c["min_side"]) & (bh >= c["min_side"]) & (bw * bh <= c["max_area_frac"] * orig_w * orig_h) & (ar <= c["max_aspect"]) return boxes[keep], scores[keep] @staticmethod def _nms(boxes: np.ndarray, scores: np.ndarray, iou_thr: float) -> np.ndarray: order = np.argsort(-scores, kind="stable") keep = [] while order.size > 0: i = order[0] keep.append(i) if order.size == 1: break rest = order[1:] xx1 = np.maximum(boxes[i, 0], boxes[rest, 0]) yy1 = np.maximum(boxes[i, 1], boxes[rest, 1]) xx2 = np.minimum(boxes[i, 2], boxes[rest, 2]) yy2 = np.minimum(boxes[i, 3], boxes[rest, 3]) inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1) area_i = (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1]) area_r = (boxes[rest, 2] - boxes[rest, 0]) * (boxes[rest, 3] - boxes[rest, 1]) iou = inter / (area_i + area_r - inter + 1e-7) order = rest[iou <= iou_thr] return np.array(keep, dtype=np.intp) @staticmethod def _ios(boxes: np.ndarray, scores: np.ndarray, thr): if thr is None or len(boxes) < 2: return boxes, scores xx1 = np.maximum(boxes[:, None, 0], boxes[None, :, 0]) yy1 = np.maximum(boxes[:, None, 1], boxes[None, :, 1]) xx2 = np.minimum(boxes[:, None, 2], boxes[None, :, 2]) yy2 = np.minimum(boxes[:, None, 3], boxes[None, :, 3]) inter = np.maximum(0.0, xx2 - xx1) * np.maximum(0.0, yy2 - yy1) area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) ios = inter / (area[None, :] + 1e-7) higher = np.tril(np.ones_like(ios, dtype=bool), -1).T drop = ((ios >= thr) & higher).any(axis=0) return boxes[~drop], scores[~drop] def _decode(self, raw: np.ndarray, ratio: float, pad: tuple[float, float], orig_w: int, orig_h: int) -> list[BoundingBox]: c = self.cfg p = raw[0].T keep = p[:, 4] >= c["conf_floor"] p = p[keep] if len(p) == 0: self.last_n_raw, self.last_conf = 0, c["conf_floor"] return [] boxes = np.empty((len(p), 4), dtype=np.float32) boxes[:, 0] = p[:, 0] - p[:, 2] / 2.0 boxes[:, 1] = p[:, 1] - p[:, 3] / 2.0 boxes[:, 2] = p[:, 0] + p[:, 2] / 2.0 boxes[:, 3] = p[:, 1] + p[:, 3] / 2.0 scores = p[:, 4].astype(np.float32) boxes[:, [0, 2]] -= pad[0] boxes[:, [1, 3]] -= pad[1] boxes /= ratio boxes[:, [0, 2]] = np.clip(boxes[:, [0, 2]], 0, orig_w - 1) boxes[:, [1, 3]] = np.clip(boxes[:, [1, 3]], 0, orig_h - 1) idx = self._nms(boxes, scores, c["nms_iou"]) boxes, scores = boxes[idx], scores[idx] boxes, scores = self._sane(boxes, scores, orig_w, orig_h) boxes, scores = self._ios(boxes, scores, c["ios"]) if len(boxes) == 0: self.last_n_raw, self.last_conf = 0, c["conf_floor"] return [] boxes, scores = boxes[: c["max_det"]], scores[: c["max_det"]] self.last_n_raw = len(scores) self.last_conf = dynamic_conf(self.last_n_raw) sel = scores >= self.last_conf return [ BoundingBox(x1=int(math.floor(b[0])), y1=int(math.floor(b[1])), x2=int(math.ceil(b[2])), y2=int(math.ceil(b[3])), cls_id=0, conf=float(s)) for b, s in zip(boxes[sel], scores[sel]) if b[2] > b[0] and b[3] > b[1] ] def _predict_single(self, image: np.ndarray) -> list[BoundingBox]: if image.dtype != np.uint8: image = image.astype(np.uint8) orig_h, orig_w = image.shape[:2] x, ratio, pad = self._preprocess(image) raw = self.session.run(self.output_names, {self.input_name: x})[0] return self._decode(raw, ratio, pad, orig_w, orig_h) def predict_batch(self, batch_images: list[ndarray], offset: int, n_keypoints: int) -> list[TVFrameResult]: results: list[TVFrameResult] = [] for i, image in enumerate(batch_images): try: boxes = self._predict_single(image) except Exception as e: print(f"Inference failed for frame {offset + i}: {e}") boxes = [] results.append(TVFrameResult(frame_id=offset + i, boxes=boxes, keypoints=[(0, 0) for _ in range(max(0, int(n_keypoints)))])) return results