scorevision: push artifact

miner.py

@@ -0,0 +1,214 @@
+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