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