| """SN44 TurboVision public-track miner for a raw YOLO head (the owner's element_trainer baselines: YOLO11n, detect or |
| seg, output [1, 4 + NC (+ mask coefficients), N]), decoded exactly as the owner's own baseline miner decodes it. |
| |
| The validator hands predict_batch OpenCV BGR uint8 frames. Each frame is converted to RGB, resized (plain stretch, no |
| letterbox) to the model's input size, scaled to 0..1, and run on the CPU with 2 threads. Per anchor: the arg-max of the |
| NC class scores (mask coefficients ignored), kept when it reaches CONF; one class-agnostic greedy NMS at IOU over all |
| kept boxes; boxes mapped back by the stretch factors, floored / ceiled to whole pixels and clipped to the frame. |
| CLASS_MAP maps the model's classes to the element's `objects` indices (-1 drops a class; empty: the same order). |
| """ |
|
|
| import math |
| from pathlib import Path |
|
|
| import cv2 |
| import numpy as np |
| import onnxruntime as ort |
| from numpy import ndarray |
| from pydantic import BaseModel |
|
|
| MODEL_FILE = "model.onnx" |
| NC = 5 |
| CONF = 0.25 |
| IOU = 0.45 |
| THREADS = 2 |
| CLASS_MAP = () |
| WARMUP_RUNS = 3 |
|
|
|
|
| class BoundingBox(BaseModel): |
| x1: int |
| y1: int |
| x2: int |
| y2: int |
| cls_id: int |
| conf: float |
|
|
|
|
| class Polygon(BaseModel): |
| cls_id: int |
| conf: float |
| points: list[tuple[int, int]] |
|
|
|
|
| class TVFrameResult(BaseModel): |
| frame_id: int |
| boxes: list[BoundingBox] | None = None |
| polygons: list[Polygon] | None = None |
| keypoints: list[tuple[int, int]] | None = None |
|
|
|
|
| def _as_bgr(img) -> ndarray: |
| if not isinstance(img, np.ndarray) or img.ndim not in (2, 3) or img.size == 0: |
| raise ValueError(f"unsupported frame: {type(img).__name__} {getattr(img, 'shape', None)}") |
| if img.dtype != np.uint8: |
| img = np.clip(np.nan_to_num(img), 0, 255).astype(np.uint8) |
| if img.ndim == 2 or img.shape[2] == 1: |
| return cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) |
| if img.shape[2] == 4: |
| return cv2.cvtColor(img, cv2.COLOR_BGRA2BGR) |
| return img |
|
|
|
|
| def _nms(boxes: ndarray, scores: ndarray, iou_thr: float) -> list[int]: |
| order = scores.argsort()[::-1] |
| keep = [] |
| while order.size > 0: |
| i = order[0] |
| keep.append(int(i)) |
| 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 / np.maximum(area_i + area_r - inter, 1e-6) |
| order = rest[np.where(iou <= iou_thr)[0]] |
| return keep |
|
|
|
|
| class Miner: |
| def __init__(self, path_hf_repo: Path) -> None: |
| so = ort.SessionOptions() |
| so.intra_op_num_threads = THREADS |
| so.inter_op_num_threads = 1 |
| so.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL |
| so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL |
| so.add_session_config_entry("session.intra_op.allow_spinning", "0") |
| cv2.setNumThreads(1) |
| self.session = ort.InferenceSession(str(Path(path_hf_repo) / MODEL_FILE), sess_options=so, |
| providers=["CPUExecutionProvider"]) |
| inp = self.session.get_inputs()[0] |
| self.input_name = inp.name |
| self.in_h, self.in_w = int(inp.shape[2]), int(inp.shape[3]) |
| dummy = np.zeros((1024, 1024, 3), np.uint8) |
| for _ in range(WARMUP_RUNS): |
| self._frame_boxes(dummy) |
|
|
| def __repr__(self) -> str: |
| return f"{MODEL_FILE} @ {self.in_w}x{self.in_h} stretch, CPU {THREADS} threads, conf {CONF}, agnostic NMS {IOU}" |
|
|
| def _frame_boxes(self, img) -> list[BoundingBox]: |
| img = _as_bgr(img) |
| h, w = img.shape[:2] |
| rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) |
| x = cv2.resize(rgb, (self.in_w, self.in_h)).astype(np.float32) / 255.0 |
| x = np.ascontiguousarray(np.transpose(x, (2, 0, 1))[None]) |
| pred = self.session.run(None, {self.input_name: x})[0][0] |
| if pred.shape[0] < pred.shape[1]: |
| pred = pred.T |
| scores = pred[:, 4:4 + NC] |
| cls = scores.argmax(1) |
| conf = scores.max(1) |
| m = np.isfinite(pred[:, :4]).all(1) & (conf >= CONF) |
| b, cls, conf = pred[m, :4], cls[m], conf[m] |
| if b.shape[0] == 0: |
| return [] |
| sx, sy = w / float(self.in_w), h / float(self.in_h) |
| xyxy = np.stack([(b[:, 0] - b[:, 2] / 2) * sx, (b[:, 1] - b[:, 3] / 2) * sy, |
| (b[:, 0] + b[:, 2] / 2) * sx, (b[:, 1] + b[:, 3] / 2) * sy], 1) |
| boxes = [] |
| for i in _nms(xyxy, conf, IOU): |
| c = int(cls[i]) |
| if CLASS_MAP: |
| if c >= len(CLASS_MAP) or CLASS_MAP[c] < 0: |
| continue |
| c = CLASS_MAP[c] |
| x1, y1, x2, y2 = xyxy[i] |
| ix1, iy1 = max(0, min(w, math.floor(x1))), max(0, min(h, math.floor(y1))) |
| ix2, iy2 = max(0, min(w, math.ceil(x2))), max(0, min(h, math.ceil(y2))) |
| if ix2 - ix1 < 1 or iy2 - iy1 < 1: |
| continue |
| boxes.append(BoundingBox(x1=ix1, y1=iy1, x2=ix2, y2=iy2, cls_id=c, conf=float(max(0.0, min(1.0, conf[i]))))) |
| return boxes |
|
|
| def predict_batch(self, batch_images: list[ndarray], offset: int, n_keypoints: int) -> list[TVFrameResult]: |
| results: list[TVFrameResult] = [] |
| for i, img in enumerate(batch_images): |
| try: |
| boxes = self._frame_boxes(img) |
| except Exception as exc: |
| print(f"inference_failure frame {offset + i}: {type(exc).__name__}: {exc}") |
| boxes = [] |
| results.append(TVFrameResult(frame_id=offset + i, boxes=boxes, polygons=[], |
| keypoints=[(0, 0)] * int(n_keypoints) if n_keypoints else [])) |
| return results |
|
|
