scorevision: push artifact

miner.py

@@ -1,59 +1,11 @@
-"""Strict consumer for an Ultralytics end-to-end-NMS ONNX (person element).
-
-Contract: static batch-1 NCHW input, one ``[1, 300, 6]`` output whose rows are
-``x1, y1, x2, y2, confidence, model_class_id``.  NMS runs inside the graph, so
-this module only validates the tensors, applies the confidence gate, undoes the
-letterbox, drops geometrically impossible boxes and caps the result count.
-
-No ``dataclass`` declarations: ScoreVision's ``spec_from_file_location`` loader
-does not register the module in ``sys.modules`` before executing it, which makes
-dataclass decoration fail at import.
-
-ELEMENT_CONFIG below is written by ``scripts/build_submission.py`` from the
-policy chosen on the validation split -- do not hand-edit it without re-running
-``scripts/optimize_policy.py``, since ``nms_iou`` in particular must match the
-value frozen into the ONNX graph at export time.
-"""
-
-import ast
 import json
-import math
-import os
 from pathlib import Path
-from typing import Any
-
 import cv2
 import numpy as np
 import onnxruntime as ort
 from numpy import ndarray
 from pydantic import BaseModel
 
-API_VERSION = "c2-strict-v1"
-STRICT_INFERENCE_ENV = "SV_STRICT_INFERENCE"
-OUTPUT_ROWS = 300
-OUTPUT_COLUMNS = 6
-
-ELEMENT_CONFIG: dict[str, Any] = {
-    "element_id": "person",
-    "classes": ["person"],
-    "coco_ids": [0],
-    "imgsz": 704,
-    "conf": 0.5,
-    "nms_iou": 0.45,  # frozen into the ONNX graph at export
-    "max_det": 300,
-    "min_side": 3,
-    "max_aspect": 10,
-    "max_area_frac": 0.95,
-    "scale_w": 1,  # box geometry calibration to the validator standard
-    "scale_h": 1,
-    "conf_steps": ((30, 0.5), (60, 0.55), (None, 0.6)),
-}
-
-
-class InferenceInvariantError(RuntimeError):
-    """A model tensor or decoded detection violated the serving contract."""
-
-
 class BoundingBox(BaseModel):
     x1: int
     y1: int
@@ -62,240 +14,132 @@ class BoundingBox(BaseModel):
     cls_id: int
     conf: float
 
-
 class TVFrameResult(BaseModel):
     frame_id: int
     boxes: list[BoundingBox]
     keypoints: list[tuple[int, int]]
 
-
-def _conf_for(n_raw: int) -> float:
-    """Confidence gate, optionally raised on crowded frames.
-
-    The wrong-box penalty is a per-image absolute count that saturates at ten,
-    so a frame holding ninety people can lose the whole term while a frame with
-    six cannot.  When CONF_STEPS is non-empty the gate is spent where that
-    exposure actually is.  Empty steps means a flat gate.
-    """
-    for limit, conf in ELEMENT_CONFIG.get("conf_steps", ()):
-        if limit is None or n_raw < limit:
-            return float(conf)
-    return float(ELEMENT_CONFIG["conf"])
-
-
 class Miner:
-    api_version = API_VERSION
-    strict_inference = False
 
     def __init__(self, path_hf_repo: Path) -> None:
-        cfg = ELEMENT_CONFIG
-        self.config = cfg
-        self.class_names = list(cfg["classes"])
-        self.conf_threshold = float(cfg["conf"])
-        self.max_det = int(cfg["max_det"])
-        self.min_side = float(cfg.get("min_side", 3.0))
-        self.max_aspect_ratio = float(cfg.get("max_aspect", 10.0))
-        self.max_area_frac = float(cfg.get("max_area_frac", 0.95))
-        self.conf_steps = tuple(tuple(s) for s in cfg.get("conf_steps", ()))
-        # box geometry calibration: the validator draws wider/taller boxes than
-        # the training sources; scale about the centre before the IoU match.
-        self.scale_w = float(cfg.get("scale_w", 1.0))
-        self.scale_h = float(cfg.get("scale_h", 1.0))
-        raw = os.environ.get(STRICT_INFERENCE_ENV, "")
-        if raw not in ("", "0", "1"):
-            raise ValueError(f"{STRICT_INFERENCE_ENV} must be '1' or unset/'0', got {raw!r}")
-        self.strict_inference = raw == "1"
-        self.inference_failures = 0
-
-        # The referee grants one physical core (two SMT threads).  Let ORT own
-        # both for the convolution graph and keep OpenCV single-threaded: a
-        # second OpenCV worker is slower and adds tail jitter on two siblings.
-        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"
-        try:
-            self.session = ort.InferenceSession(
-                str(model_path), sess_options=options,
-                providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
-            )
-        except Exception as exc:
-            print(f"CUDA session failed, falling back to CPU: {exc}")
-            self.session = ort.InferenceSession(
-                str(model_path), sess_options=options, providers=["CPUExecutionProvider"]
-            )
-
-        self.input_name = self.session.get_inputs()[0].name
-        self.output_names = [o.name for o in self.session.get_outputs()]
-        shape = self.session.get_inputs()[0].shape
-        if len(shape) != 4 or shape[1] != 3:
-            raise ValueError(f"model input must be rank-4 NCHW with 3 channels, got {shape!r}")
-        if not all(isinstance(v, int) and v > 0 for v in shape[2:]):
-            raise ValueError(f"e2e export must have static spatial dims, got {shape[2:]!r}")
-        if (shape[2], shape[3]) != (cfg["imgsz"], cfg["imgsz"]):
-            raise ValueError(f"model input {shape[2]}x{shape[3]} != config imgsz {cfg['imgsz']}")
-        self.input_height = self.input_width = int(cfg["imgsz"])
-        out_shape = self.session.get_outputs()[0].shape
-        if list(out_shape) != [1, OUTPUT_ROWS, OUTPUT_COLUMNS]:
-            raise ValueError(f"e2e output must be [1, {OUTPUT_ROWS}, {OUTPUT_COLUMNS}], got {out_shape!r}")
-        self.model_names = self._read_model_names()
-        self._warmup(3)
-        print(
-            f"api_version={self.api_version} element={cfg['element_id']} imgsz={cfg['imgsz']} "
-            f"conf={self.conf_threshold} steps={self.conf_steps} "
-            f"providers={self.session.get_providers()}"
-        )
-
-    def _read_model_names(self) -> tuple[str, ...]:
-        try:
-            raw = self.session.get_modelmeta().custom_metadata_map["names"]
-            parsed = ast.literal_eval(raw)
-            if not isinstance(parsed, dict):
-                raise TypeError("names metadata is not a dict")
-            keys = sorted(parsed)
-            if keys != list(range(len(keys))):
-                raise ValueError(f"names metadata keys not consecutive from zero: {keys[:8]}")
-            return tuple(str(parsed[i]) for i in keys)
-        except Exception as exc:
-            raise ValueError(f"e2e ONNX requires valid names metadata: {exc}") from exc
-
-    @staticmethod
-    def _letterbox(image: ndarray, new_shape: tuple[int, int], color=(114, 114, 114)):
-        height, width = image.shape[:2]
-        new_width, new_height = new_shape
-        ratio = min(new_width / width, new_height / height)
-        rw, rh = int(round(width * ratio)), int(round(height * ratio))
-        if (rw, rh) != (width, height):
-            image = cv2.resize(
-                image, (rw, rh),
-                interpolation=cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR,
-            )
-        dw, dh = (new_width - rw) / 2.0, (new_height - rh) / 2.0
-        padded = cv2.copyMakeBorder(
-            image,
-            int(round(dh - 0.1)), int(round(dh + 0.1)),
-            int(round(dw - 0.1)), int(round(dw + 0.1)),
-            cv2.BORDER_CONSTANT, value=color,
-        )
-        return padded, ratio, (dw, dh)
-
-    def _decode(self, output, ratio, pad, original_size):
-        rows = output[0]
-        boxes = rows[:, :4].astype(np.float32, copy=True)
-        scores = rows[:, 4].astype(np.float32, copy=False)
-        raw_classes = rows[:, 5]
-        if (scores < 0).any() or (scores > 1).any():
-            raise InferenceInvariantError("e2e confidence outside [0, 1]")
-        rounded = np.rint(raw_classes)
-        if not np.array_equal(raw_classes, rounded):
-            raise InferenceInvariantError("e2e model class ids must be exact integers")
-        model_classes = rounded.astype(np.int64)
-        if (model_classes < 0).any() or (model_classes >= len(self.model_names)).any():
-            raise InferenceInvariantError("e2e model class id outside range")
-
-        active = scores > 0.0
-        reversed_rows = active & ((boxes[:, 2] < boxes[:, 0]) | (boxes[:, 3] < boxes[:, 1]))
-        if reversed_rows.any():
-            raise InferenceInvariantError("e2e output contains reversed boxes")
-
-        # A floor below every reachable gate; nothing under it can be emitted.
-        floor = min([self.conf_threshold] + [float(c) for _, c in self.conf_steps])
-        keep = active & (model_classes == 0) & (scores >= np.float32(floor))
-        if not keep.any():
-            return np.empty((0, 4), np.float32), np.empty((0,), np.float32)
-        boxes, scores = boxes[keep], scores[keep]
-
-        boxes[:, [0, 2]] = (boxes[:, [0, 2]] - pad[0]) / ratio
-        boxes[:, [1, 3]] = (boxes[:, [1, 3]] - pad[1]) / ratio
-        width, height = original_size
-        if self.scale_w != 1.0 or self.scale_h != 1.0:
-            cx = (boxes[:, 0] + boxes[:, 2]) * 0.5
-            cy = (boxes[:, 1] + boxes[:, 3]) * 0.5
-            hw = (boxes[:, 2] - boxes[:, 0]) * 0.5 * np.float32(self.scale_w)
-            hh = (boxes[:, 3] - boxes[:, 1]) * 0.5 * np.float32(self.scale_h)
-            boxes = np.stack([cx - hw, cy - hh, cx + hw, cy + hh], 1).astype(np.float32)
-        boxes[:, [0, 2]] = np.clip(boxes[:, [0, 2]], 0, width - 1)
-        boxes[:, [1, 3]] = np.clip(boxes[:, [1, 3]], 0, height - 1)
-
-        bw = boxes[:, 2] - boxes[:, 0]
-        bh = boxes[:, 3] - boxes[:, 1]
-        aspect = np.maximum(bw / np.maximum(bh, 1e-6), bh / np.maximum(bw, 1e-6))
-        sane = (
-            (bw >= self.min_side)
-            & (bh >= self.min_side)
-            & (bw * bh <= self.max_area_frac * float(width * height))
-            & (aspect <= self.max_aspect_ratio)
-        )
-        boxes, scores = boxes[sane], scores[sane]
-        if not len(scores):
-            return boxes, scores
-        order = np.argsort(-scores, kind="stable")[: self.max_det]
-        boxes, scores = boxes[order], scores[order]
-        gate = _conf_for(len(scores)) if self.conf_steps else self.conf_threshold
-        emit = scores >= np.float32(gate)
-        return boxes[emit], scores[emit]
-
-    def _predict_single(self, image: np.ndarray) -> list[BoundingBox]:
-        if not isinstance(image, np.ndarray) or image.ndim != 3 or image.shape[2] != 3:
-            raise ValueError(f"expected HWC 3-channel image, got {getattr(image, 'shape', type(image))}")
-        if image.dtype != np.uint8:
-            image = image.astype(np.uint8)
-        original_height, original_width = image.shape[:2]
-        padded, ratio, pad = self._letterbox(image, (self.input_width, self.input_height))
-        blob = cv2.dnn.blobFromImage(padded, scalefactor=1 / 255.0, swapRB=True)
-        outputs = self.session.run(self.output_names, {self.input_name: blob})
-        output = outputs[0]
-        if not np.isfinite(output).all():
-            raise InferenceInvariantError("model output must be finite everywhere")
-        if list(output.shape) != [1, OUTPUT_ROWS, OUTPUT_COLUMNS]:
-            raise ValueError(f"unexpected output shape {list(output.shape)}")
-        boxes, scores = self._decode(output, ratio, pad, (original_width, original_height))
-        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, scores)
-            if b[2] > b[0] and b[3] > b[1]
-        ]
-
-    def _warmup(self, iterations: int) -> None:
-        dummy = np.zeros((720, 1280, 3), dtype=np.uint8)
-        original = self.strict_inference
-        self.strict_inference = True
-        try:
-            for _ in range(iterations):
-                self.predict_batch(batch_images=[dummy], offset=0, n_keypoints=0)
-        except Exception as exc:
-            raise RuntimeError(f"warmup failed, refusing to serve: {type(exc).__name__}: {exc}") from exc
-        finally:
-            self.strict_inference = original
-
-    def predict_batch(self, batch_images, offset: int, n_keypoints: int) -> list[TVFrameResult]:
-        _ = n_keypoints
-        results: list[TVFrameResult] = []
-        for index, image in enumerate(batch_images):
-            frame_id = offset + index
-            try:
-                boxes = self._predict_single(image)
-            except Exception as exc:
-                if self.strict_inference:
-                    raise
-                self.inference_failures += 1
-                print(json.dumps({
-                    "event": "inference_failure",
-                    "frame_id": frame_id,
-                    "error_type": type(exc).__name__,
-                    "error": str(exc),
-                    "inference_failures": self.inference_failures,
-                }, sort_keys=True))
-                boxes = []
-            results.append(TVFrameResult(frame_id=frame_id, boxes=boxes, keypoints=[]))
-        return results
+        repo = Path(path_hf_repo)
+        cfg = json.loads((repo / 'calib.json').read_text())
+        self.calib_x = np.asarray(cfg['calib_x'], dtype=np.float32)
+        self.calib_y = np.asarray(cfg['calib_y'], dtype=np.float32)
+        self.k = float(cfg.get('k', 15.0))
+        self.miss_prior = float(cfg.get('miss_prior', 0.0))
+        self.p_floor = float(cfg.get('p_floor', 0.05))
+        self.min_side = float(cfg.get('min_side', 2.0))
+        self.nms_iou = float(cfg.get('nms_iou', 0.6))
+        self.conf_min = float(cfg.get('conf_min', 0.02))
+        self.max_det = int(cfg.get('max_det', 300))
+        self.raw_xyxy = bool(cfg.get('raw_xyxy', True))
+        self.warp = tuple(cfg['warp']) if cfg.get('warp') else None
+        so = ort.SessionOptions()
+        so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+        so.intra_op_num_threads = int(cfg.get('threads', 2))
+        so.inter_op_num_threads = 1
+        so.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
+        self.session = ort.InferenceSession(str(repo / 'weights.onnx'), sess_options=so, providers=['CPUExecutionProvider'])
+        inp = self.session.get_inputs()[0]
+        self.input_name = inp.name
+        self.uint8_nhwc = 'uint8' in inp.type
+        self.size = int(inp.shape[1] if self.uint8_nhwc else inp.shape[2])
+        dummy = np.zeros((1, self.size, self.size, 3), np.uint8) if self.uint8_nhwc else np.zeros((1, 3, self.size, self.size), np.float32)
+        for _ in range(3):
+            self.session.run(None, {self.input_name: dummy})
 
     def __repr__(self) -> str:
-        return f"E2EONNXRuntime(providers={self.session.get_providers()})"
+        return f'PersonMiner(size={self.size}, k={self.k}, providers={self.session.get_providers()})'
+
+    def _prep(self, img: ndarray):
+        h, w = img.shape[:2]
+        s = self.size
+        if h == w:
+            x = cv2.resize(img, (s, s), interpolation=cv2.INTER_LINEAR)
+            r, left, top = (s / w, 0, 0)
+        else:
+            r = min(s / h, s / w)
+            nh, nw = (int(round(h * r)), int(round(w * r)))
+            x = np.full((s, s, 3), 114, np.uint8)
+            top, left = ((s - nh) // 2, (s - nw) // 2)
+            x[top:top + nh, left:left + nw] = cv2.resize(img, (nw, nh), interpolation=cv2.INTER_LINEAR)
+        if self.uint8_nhwc:
+            x = x[None]
+        else:
+            x = x[:, :, ::-1].transpose(2, 0, 1)[None].astype(np.float32) * (1.0 / 255.0)
+        return (np.ascontiguousarray(x), r, left, top)
+
+    def _warp_prep(self, img: ndarray):
+        h, w = img.shape[:2]
+        s = self.size
+        hs, hd = (int(round(self.warp[0] * h)), int(round(self.warp[1] * s)))
+        x = np.concatenate([cv2.resize(img[:hs], (s, hd), interpolation=cv2.INTER_LINEAR), cv2.resize(img[hs:], (s, s - hd), interpolation=cv2.INTER_LINEAR)], 0)
+        return (np.ascontiguousarray(x[None]), hs, hd)
+
+    def _unwarp(self, b: ndarray, w: int, h: int, hs: int, hd: int) -> ndarray:
+        s = self.size
+        b = b.copy()
+        b[:, [0, 2]] *= w / s
+        for j in (1, 3):
+            y = b[:, j]
+            b[:, j] = np.where(y < hd, y * hs / hd, hs + (y - hd) * (h - hs) / (s - hd))
+        return b
+
+    def _nms_raw(self, y: ndarray, warp_args=None) -> ndarray:
+        y = y.T
+        sc = y[:, 4:].max(1)
+        m = sc >= self.conf_min
+        y, sc = (y[m], sc[m])
+        if sc.size == 0:
+            return np.zeros((0, 6), np.float32)
+        if self.raw_xyxy:
+            b = y[:, :4].copy()
+        else:
+            b = np.stack([y[:, 0] - y[:, 2] / 2, y[:, 1] - y[:, 3] / 2, y[:, 0] + y[:, 2] / 2, y[:, 1] + y[:, 3] / 2], 1)
+        if warp_args is not None:
+            b = self._unwarp(b, *warp_args)
+        xywh = np.concatenate([b[:, :2], b[:, 2:] - b[:, :2]], 1).astype(np.float64)
+        k = np.asarray(cv2.dnn.NMSBoxes(xywh.tolist(), sc.astype(np.float32).tolist(), self.conf_min, self.nms_iou), dtype=int).reshape(-1)[:self.max_det]
+        return np.concatenate([b[k], sc[k, None], np.zeros((k.size, 1), b.dtype)], 1)
+
+    def _decide(self, det: ndarray, w: int, h: int, r: float, left: int, top: int) -> list[BoundingBox]:
+        det = det[det[:, 4] > 0.0]
+        if det.size == 0:
+            return []
+        p = np.interp(det[:, 4], self.calib_x, self.calib_y).astype(np.float64)
+        n_hat = float(p.sum()) + self.miss_prior
+        thr = max(n_hat / (self.k + n_hat), self.p_floor)
+        keep = p > thr
+        det, p = (det[keep], p[keep])
+        order = np.argsort(-p, kind='stable')
+        det, p = (det[order], p[order])
+        b = det[:, :4].copy()
+        b[:, [0, 2]] = (b[:, [0, 2]] - left) / r
+        b[:, [1, 3]] = (b[:, [1, 3]] - top) / r
+        b[:, [0, 2]] = np.clip(b[:, [0, 2]], 0, w)
+        b[:, [1, 3]] = np.clip(b[:, [1, 3]], 0, h)
+        out = []
+        for (x1, y1, x2, y2), pi in zip(b, p):
+            if x2 - x1 < self.min_side or y2 - y1 < self.min_side:
+                continue
+            out.append(BoundingBox(x1=int(round(x1)), y1=int(round(y1)), x2=int(round(x2)), y2=int(round(y2)), cls_id=0, conf=float(round(pi, 5))))
+        return out
+
+    def predict_batch(self, batch_images: list[ndarray], offset: int, n_keypoints: int) -> list[TVFrameResult]:
+        results = []
+        for i, img in enumerate(batch_images):
+            h, w = img.shape[:2]
+            if self.warp is not None and h == w and self.uint8_nhwc:
+                x, hs, hd = self._warp_prep(img)
+                det = self.session.run(None, {self.input_name: x})[0][0]
+                det = self._nms_raw(det, (w, h, hs, hd))
+                boxes = self._decide(det, w, h, 1.0, 0, 0)
+            else:
+                x, r, left, top = self._prep(img)
+                det = self.session.run(None, {self.input_name: x})[0][0]
+                if det.shape[-1] != 6:
+                    det = self._nms_raw(det)
+                boxes = self._decide(det, w, h, r, left, top)
+            results.append(TVFrameResult(frame_id=offset + i, boxes=boxes, keypoints=[]))
+        return results