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