from __future__ import annotations import math from pathlib import Path import cv2 import numpy as np import onnxruntime as ort from numpy import ndarray from pydantic import BaseModel PERSON_CLS_ID = 0 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: BOX_FORMAT = "xyxy" CONF_THRESHOLD = 0.56 NMS_IOU = 0.45 MAX_DET = 50 # Scores are rounded to this many decimals before anything is compared or # sorted. The compliance checker reruns this repo on its own machine and # holds the boxes against the ones the chute returned, so an ordering that a # difference in the sixth digit can flip turns arithmetic noise into a # missing or an extra box. Four decimals is far coarser than any such # difference and far finer than the gaps that decide real detections. SCORE_DECIMALS = 4 PRE_NMS_TOPK = 1000 MIN_SIDE = 3.0 MAX_ASPECT_RATIO = 10.0 MAX_BOX_AREA_RATIO = 0.95 def __init__(self, path_hf_repo: Path) -> None: cv2.setNumThreads(1) options = ort.SessionOptions() options.intra_op_num_threads = 2 options.inter_op_num_threads = 1 options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL model_path = Path(path_hf_repo) / "weights.onnx" self.session = ort.InferenceSession( str(model_path), sess_options=options, providers=["CPUExecutionProvider"], ) model_input = self.session.get_inputs()[0] self.input_name = model_input.name self.input_height = int(model_input.shape[2]) self.input_width = int(model_input.shape[3]) self._warmup() def _warmup(self, iters: int = 3) -> None: dummy = np.zeros((720, 1280, 3), dtype=np.uint8) for _ in range(max(0, iters)): self._predict_single(dummy) def _letterbox(self, image: ndarray) -> tuple[ndarray, float, tuple[float, float]]: h0, w0 = image.shape[:2] ratio = min(self.input_width / w0, self.input_height / h0) new_w, new_h = int(round(w0 * ratio)), int(round(h0 * ratio)) if (new_w, new_h) != (w0, h0): image = cv2.resize( image, (new_w, new_h), interpolation=cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR, ) pad_w = (self.input_width - new_w) / 2.0 pad_h = (self.input_height - new_h) / 2.0 padded = cv2.copyMakeBorder( image, int(round(pad_h - 0.1)), int(round(pad_h + 0.1)), int(round(pad_w - 0.1)), int(round(pad_w + 0.1)), cv2.BORDER_CONSTANT, value=(114, 114, 114), ) return padded, ratio, (pad_w, pad_h) def _preprocess(self, image: ndarray) -> tuple[ndarray, float, tuple[float, float]]: padded, ratio, pad = self._letterbox(image) blob = cv2.dnn.blobFromImage(padded, scalefactor=1 / 255.0, swapRB=True) return blob, ratio, pad def _candidates(self, preds: ndarray) -> tuple[ndarray, ndarray]: anchors = preds[0].T boxes = anchors[:, :4].astype(np.float32) if self.BOX_FORMAT == "xywh": half_w, half_h = boxes[:, 2] / 2.0, boxes[:, 3] / 2.0 boxes = np.stack([boxes[:, 0] - half_w, boxes[:, 1] - half_h, boxes[:, 0] + half_w, boxes[:, 1] + half_h], axis=1) return boxes, anchors[:, 4].astype(np.float32) def _rank(self, boxes: ndarray, scores: ndarray) -> ndarray: """Best first, by a key two machines cannot disagree on. The quantised score decides, and where it ties the box's own top-left corner does, so the result never depends on the order the candidates arrived in or on a difference in the last float digits. """ q = np.round(scores.astype(np.float64), self.SCORE_DECIMALS) return np.lexsort((np.round(boxes[:, 1], 3), np.round(boxes[:, 0], 3), -q)) def _screen(self, boxes: ndarray, scores: ndarray) -> tuple[ndarray, ndarray]: keep = np.round(scores.astype(np.float64), self.SCORE_DECIMALS) >= self.CONF_THRESHOLD boxes, scores = boxes[keep], scores[keep] if boxes.shape[0] > self.PRE_NMS_TOPK: order = self._rank(boxes, scores)[: self.PRE_NMS_TOPK] boxes, scores = boxes[order], scores[order] return boxes, scores def _nms(self, boxes: ndarray, scores: ndarray, threshold: float) -> ndarray: x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3] areas = np.maximum(x2 - x1, 0.0) * np.maximum(y2 - y1, 0.0) order = self._rank(boxes, scores) keep: list[int] = [] while order.size > 0: best = int(order[0]) keep.append(best) if order.size == 1: break rest = order[1:] xx1 = np.maximum(x1[best], x1[rest]) yy1 = np.maximum(y1[best], y1[rest]) xx2 = np.minimum(x2[best], x2[rest]) yy2 = np.minimum(y2[best], y2[rest]) inter = np.clip(xx2 - xx1, 0.0, None) * np.clip(yy2 - yy1, 0.0, None) iou = inter / np.maximum(areas[best] + areas[rest] - inter, 1e-9) order = rest[iou < threshold] return np.array(keep, dtype=np.int64) def _suppress( self, boxes: ndarray, scores: ndarray, orig_size: tuple[int, int] ) -> tuple[ndarray, ndarray]: w, h = orig_size widths = boxes[:, 2] - boxes[:, 0] heights = boxes[:, 3] - boxes[:, 1] aspect = np.maximum(widths / np.maximum(heights, 1e-6), heights / np.maximum(widths, 1e-6)) sane = (widths >= self.MIN_SIDE) & (heights >= self.MIN_SIDE) sane &= widths * heights <= self.MAX_BOX_AREA_RATIO * float(w * h) sane &= aspect <= self.MAX_ASPECT_RATIO boxes, scores = boxes[sane], scores[sane] order = self._rank(boxes, scores)[: self.MAX_DET] return boxes[order], scores[order] @staticmethod def _clip_boxes(boxes: ndarray, image_size: tuple[int, int]) -> ndarray: w, h = image_size boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, w - 1) boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, h - 1) return boxes def _decode( self, preds: ndarray, ratio: float, pad: tuple[float, float], orig_size: tuple[int, int], ) -> list[BoundingBox]: boxes, scores = self._candidates(preds) boxes, scores = self._screen(boxes, scores) keep = self._nms(boxes, scores, self.NMS_IOU) boxes, scores = boxes[keep], scores[keep] pad_w, pad_h = pad boxes[:, [0, 2]] = (boxes[:, [0, 2]] - pad_w) / ratio boxes[:, [1, 3]] = (boxes[:, [1, 3]] - pad_h) / ratio boxes = self._clip_boxes(boxes, orig_size) boxes, scores = self._suppress(boxes, scores, orig_size) return [ BoundingBox( x1=int(math.floor(float(b[0]))), y1=int(math.floor(float(b[1]))), x2=int(math.ceil(float(b[2]))), y2=int(math.ceil(float(b[3]))), cls_id=PERSON_CLS_ID, conf=float(c), ) for b, c in zip(boxes, scores) if b[2] > b[0] and b[3] > b[1] ] def _predict_single(self, image: ndarray) -> list[BoundingBox]: blob, ratio, pad = self._preprocess(image) preds = self.session.run(None, {self.input_name: blob})[0] return self._decode(preds, ratio, pad, (image.shape[1], image.shape[0])) def predict_batch( self, batch_images: list[ndarray], offset: int, n_keypoints: int, ) -> list[TVFrameResult]: _ = n_keypoints results: list[TVFrameResult] = [] for index, image in enumerate(batch_images): try: boxes = self._predict_single(image) except Exception as exc: print(f"⚠️ Inference failed for frame {offset + index}: {exc}") boxes = [] results.append( TVFrameResult(frame_id=offset + index, boxes=boxes, keypoints=[]) ) return results