scorevision: push artifact

miner.py

@@ -0,0 +1,556 @@
+"""Strict config-driven consumer for Ultralytics end-to-end NMS ONNX.
+
+The model contract is a static batch-1 input and one fixed ``[1, 300, 6]``
+output whose rows are ``x1, y1, x2, y2, confidence, model_class_id``.  NMS is
+inside the ONNX graph.  Python only validates, applies the element confidence
+threshold and model-class mapping from the hardcoded element config, unletterboxes, filters
+geometrically impossible boxes, sorts, and caps the configured result count.
+
+This module intentionally uses no ``dataclass`` declarations.  ScoreVision's
+current ``spec_from_file_location`` loader does not register the dynamic module
+in ``sys.modules`` before executing it, which makes Python dataclass decoration
+fail during import.
+"""
+
+import ast
+import json
+import math
+import os
+from collections.abc import Mapping
+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
+
+
+# The deployment rail currently recognizes this strict-inference contract
+# version for every cycle-2 miner.  The tensor layout is separately enforced by
+# this module's fixed [1, 300, 6] startup check.
+API_VERSION = "c2-strict-v1"
+STRICT_INFERENCE_ENV = "SV_STRICT_INFERENCE"
+INPUT_CHANNELS = 3
+OUTPUT_ROWS = 300
+OUTPUT_COLUMNS = 6
+IMGSZ_MIN, IMGSZ_MAX, IMGSZ_MULTIPLE = 320, 1280, 32
+MAX_DET_LIMIT = 300
+_ROOT_KEYS = {
+    "element_id", "classes", "emit_order", "coco_ids", "missing_classes", "imgsz",
+    "conf", "nms_iou", "max_det",
+}
+
+
+class InferenceInvariantError(RuntimeError):
+    """A model tensor or decoded detection violated the serving contract."""
+
+
+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]]
+
+
+def strict_inference_from_env(environ: Mapping[str, str] | None = None) -> bool:
+    raw = (os.environ if environ is None else environ).get(STRICT_INFERENCE_ENV, "")
+    if raw in ("", "0"):
+        return False
+    if raw == "1":
+        return True
+    raise ValueError(
+        f"{STRICT_INFERENCE_ENV} must be '1' or unset/'0', got {raw!r}"
+    )
+
+
+def _strict_int(value: Any, where: str, low: int, high: int) -> int:
+    if isinstance(value, bool) or not isinstance(value, int) or not low <= value <= high:
+        raise ValueError(f"element config {where} must be an integer in [{low}, {high}], got {value!r}")
+    return value
+
+
+def _strict_unit(value: Any, where: str) -> float:
+    if isinstance(value, bool) or not isinstance(value, (int, float)):
+        raise ValueError(f"element config {where} must be a finite number in (0, 1], got {value!r}")
+    number = float(value)
+    if not math.isfinite(number) or not 0.0 < number <= 1.0:
+        raise ValueError(f"element config {where} must be a finite number in (0, 1], got {value!r}")
+    return number
+
+
+class ElementConfig:
+    """Strict immutable-by-convention config without dynamic-loader dataclasses."""
+
+    __slots__ = (
+        "element_id", "classes", "emit_order", "output_cls_ids", "coco_ids", "emitted_cls_ids",
+        "missing_classes", "imgsz", "conf", "nms_iou", "max_det",
+    )
+
+    def __init__(
+        self,
+        element_id: str,
+        classes: tuple[str, ...],
+        coco_ids: tuple[int, ...],
+        emitted_cls_ids: tuple[int, ...],
+        missing_classes: tuple[str, ...],
+        imgsz: int,
+        conf: float,
+        nms_iou: float,
+        max_det: int,
+        *,
+        emit_order: tuple[str, ...] | None = None,
+        output_cls_ids: tuple[int, ...] = (),
+    ) -> None:
+        self.element_id = element_id
+        self.classes = classes
+        self.emit_order = emit_order
+        self.output_cls_ids = output_cls_ids
+        self.coco_ids = coco_ids
+        self.emitted_cls_ids = emitted_cls_ids
+        self.missing_classes = missing_classes
+        self.imgsz = imgsz
+        self.conf = conf
+        self.nms_iou = nms_iou
+        self.max_det = max_det
+
+    @classmethod
+    def from_mapping(cls, data: dict[str, Any]) -> "ElementConfig":
+        required = {"element_id", "classes", "coco_ids", "imgsz", "conf", "nms_iou", "max_det"}
+        missing = sorted(required - set(data))
+        if missing:
+            raise ValueError(f"element config missing required field(s): {', '.join(missing)}")
+        unknown = sorted(set(data) - _ROOT_KEYS)
+        if unknown:
+            raise ValueError(f"element config root has unknown key(s): {', '.join(unknown)}")
+        element_id = data["element_id"]
+        classes = data["classes"]
+        coco_ids = data["coco_ids"]
+        missing_classes = data.get("missing_classes", [])
+        if not isinstance(element_id, str) or not element_id:
+            raise ValueError("element config element_id must be a non-empty string")
+        if not isinstance(classes, list) or not classes or not all(isinstance(name, str) and name for name in classes):
+            raise ValueError("element config classes must be a non-empty list of strings")
+        if len(set(classes)) != len(classes):
+            raise ValueError("element config classes must not contain duplicates")
+        emit_order_raw = data.get("emit_order")
+        if "emit_order" not in data:
+            emit_order = None
+        else:
+            if not isinstance(emit_order_raw, list) or not emit_order_raw \
+                    or not all(isinstance(name, str) and name for name in emit_order_raw):
+                raise ValueError("element config emit_order must be a non-empty list of class names")
+            if len(set(emit_order_raw)) != len(emit_order_raw):
+                raise ValueError("element config emit_order must not contain duplicates")
+            unknown = sorted(set(emit_order_raw) - set(classes))
+            omitted = sorted(set(classes) - set(emit_order_raw))
+            if unknown or omitted or len(emit_order_raw) != len(classes):
+                raise ValueError(
+                    "element config emit_order must be an exact permutation of classes; "
+                    f"unknown={unknown}, missing={omitted}"
+                )
+            emit_order = tuple(emit_order_raw)
+        output_order = emit_order if emit_order is not None else tuple(classes)
+        output_cls_ids = tuple(output_order.index(name) for name in classes)
+        if not isinstance(missing_classes, list) or not all(isinstance(name, str) for name in missing_classes):
+            raise ValueError("element config missing_classes must be a list of strings")
+        if len(set(missing_classes)) != len(missing_classes) or any(name not in classes for name in missing_classes):
+            raise ValueError("element config missing_classes must be unique members of classes")
+        if not isinstance(coco_ids, list) or not all(
+            isinstance(class_id, int) and not isinstance(class_id, bool) and class_id >= 0
+            for class_id in coco_ids
+        ):
+            raise ValueError("element config coco_ids must be non-negative integer model class ids")
+        if len(set(coco_ids)) != len(coco_ids):
+            raise ValueError("element config coco_ids must not contain duplicates")
+        emitted_cls_ids = tuple(index for index, name in enumerate(classes) if name not in missing_classes)
+        if len(coco_ids) != len(emitted_cls_ids):
+            raise ValueError(
+                "element config classes/coco_ids length mismatch: "
+                f"{len(classes)} classes minus {len(missing_classes)} missing classes requires "
+                f"{len(emitted_cls_ids)} ids, got {len(coco_ids)}"
+            )
+        imgsz = _strict_int(data["imgsz"], "imgsz", IMGSZ_MIN, IMGSZ_MAX)
+        if imgsz % IMGSZ_MULTIPLE:
+            raise ValueError(f"element config imgsz must be a multiple of {IMGSZ_MULTIPLE}, got {imgsz}")
+        return cls(
+            element_id,
+            tuple(classes),
+            tuple(coco_ids),
+            emitted_cls_ids,
+            tuple(missing_classes),
+            imgsz,
+            _strict_unit(data["conf"], "conf"),
+            _strict_unit(data["nms_iou"], "nms_iou"),
+            _strict_int(data["max_det"], "max_det", 1, MAX_DET_LIMIT),
+            emit_order=emit_order,
+            output_cls_ids=output_cls_ids,
+        )
+
+
+ELEMENT_CONFIG: dict[str, Any] = {
+    "element_id": "person",
+    "classes": ["person"],
+    "coco_ids": [0],
+    "imgsz": 640,
+    "conf": 0.4,
+    "nms_iou": 0.7,
+    "max_det": 300,
+}
+
+
+class Miner:
+    api_version = API_VERSION
+    strict_inference = False
+    inference_failures = 0
+    invariant_failures = 0
+    input_dims_source = "unbound"
+    input_shape_declared: list[Any] = []
+    min_side = 3.0
+    max_aspect_ratio = 10.0
+
+    def __init__(self, path_hf_repo: Path) -> None:
+        repo = Path(path_hf_repo)
+        self.config = ElementConfig.from_mapping(ELEMENT_CONFIG)
+        self.class_names = list(self.config.emit_order or self.config.classes)
+        self.conf_threshold = self.config.conf
+        self.nms_iou = self.config.nms_iou  # retained for API parity; NMS is in the graph
+        self.max_det = self.config.max_det
+        self.strict_inference = strict_inference_from_env()
+        self.inference_failures = 0
+        self.invariant_failures = 0
+
+        # The referee grants one physical core (two SMT threads).  Let ORT own
+        # both threads for the convolution graph, but keep OpenCV's short
+        # resize/blob stages single-threaded: a second OpenCV worker is slower
+        # and adds tail jitter on two sibling logical CPUs.
+        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 = 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 = [output.name for output in self.session.get_outputs()]
+        self._bind_input_dims(self.session.get_inputs()[0].shape)
+        self._check_output_shape(self.session.get_outputs()[0].shape)
+        self.model_names = self._read_model_names()
+        self._bind_class_map()
+        self._warmup(3)
+        print(
+            f"api_version={self.api_version} element={self.config.element_id} "
+            f"imgsz={self.config.imgsz} conf={self.conf_threshold} "
+            f"providers={self.session.get_providers()} warmups=3"
+        )
+
+    @staticmethod
+    def _static_dim(value: Any) -> bool:
+        return isinstance(value, int) and not isinstance(value, bool) and value > 0
+
+    def _bind_input_dims(self, shape: Any) -> None:
+        if not isinstance(shape, (list, tuple)) or len(shape) != 4:
+            raise ValueError(f"model input must be rank-4 NCHW, got {shape!r}")
+        batch, channels, height, width = shape
+        self.input_shape_declared = [value if self._static_dim(value) else str(value) for value in shape]
+        if self._static_dim(batch) and batch != 1:
+            raise ValueError(f"model input batch must be 1, got {batch}")
+        if self._static_dim(channels) and channels != INPUT_CHANNELS:
+            raise ValueError(f"model input channels must be 3, got {channels}")
+        if not (self._static_dim(height) and self._static_dim(width)):
+            raise ValueError(f"e2e shipping export must have static spatial dims, got {height}x{width}")
+        if (height, width) != (self.config.imgsz, self.config.imgsz):
+            raise ValueError(
+                f"model input spatial dims {height}x{width} != config imgsz "
+                f"{self.config.imgsz}x{self.config.imgsz}"
+            )
+        self.input_height = self.input_width = self.config.imgsz
+        self.input_dims_source = "static"
+
+    @staticmethod
+    def _check_output_shape(shape: Any) -> None:
+        if list(shape) != [1, OUTPUT_ROWS, OUTPUT_COLUMNS]:
+            raise ValueError(
+                f"e2e model output must be exactly [1, {OUTPUT_ROWS}, {OUTPUT_COLUMNS}], got {shape!r}"
+            )
+
+    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 are not consecutive from zero: {keys[:8]}")
+            names = tuple(str(parsed[index]) for index in keys)
+        except Exception as exc:
+            raise ValueError(f"e2e ONNX requires valid consecutive names metadata: {exc}") from exc
+        if not names:
+            raise ValueError("e2e ONNX names metadata is empty")
+        return names
+
+    def _bind_class_map(self) -> None:
+        self.model_to_element: dict[int, int] = {}
+        for model_id, element_id in zip(self.config.coco_ids, self.config.emitted_cls_ids):
+            if model_id >= len(self.model_names):
+                raise ValueError(
+                    f"config model class id {model_id} outside ONNX names range 0..{len(self.model_names)-1}"
+                )
+            expected = self.config.classes[element_id]
+            actual = self.model_names[model_id]
+            if actual != expected:
+                raise ValueError(
+                    f"config class mapping mismatch: model class {model_id} is {actual!r}, "
+                    f"element class {element_id} is {expected!r}"
+                )
+            self.model_to_element[model_id] = element_id
+
+    @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)
+        resized_width, resized_height = int(round(width * ratio)), int(round(height * ratio))
+        if (resized_width, resized_height) != (width, height):
+            image = cv2.resize(
+                image, (resized_width, resized_height),
+                interpolation=cv2.INTER_CUBIC if ratio > 1.0 else cv2.INTER_LINEAR,
+            )
+        dw, dh = (new_width - resized_width) / 2.0, (new_height - resized_height) / 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 _preprocess(self, image: ndarray):
+        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)
+        return blob, ratio, pad, (original_width, original_height)
+
+    @staticmethod
+    def _check_image(image: np.ndarray) -> np.ndarray:
+        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))}")
+        return image if image.dtype == np.uint8 else image.astype(np.uint8)
+
+    @staticmethod
+    def _check_finite(tensor: np.ndarray, where: str) -> None:
+        if not np.isfinite(tensor).all():
+            raise InferenceInvariantError(
+                f"{where} must be finite everywhere: NaN={int(np.isnan(tensor).sum())} "
+                f"Inf={int(np.isinf(tensor).sum())}"
+            )
+
+    def _check_outputs(self, outputs: Any) -> np.ndarray:
+        if not isinstance(outputs, (list, tuple)) or len(outputs) != len(self.output_names):
+            raise ValueError(f"session returned invalid output list, expected {len(self.output_names)} tensors")
+        for name, tensor in zip(self.output_names, outputs):
+            if not isinstance(tensor, np.ndarray):
+                raise ValueError(f"model output {name!r} is not an ndarray")
+            if np.issubdtype(tensor.dtype, np.floating):
+                self._check_finite(tensor, f"model output {name!r}")
+        output = outputs[0]
+        if not np.issubdtype(output.dtype, np.floating) or list(output.shape) != [1, OUTPUT_ROWS, OUTPUT_COLUMNS]:
+            raise ValueError(
+                f"e2e output tensor must be float [1, {OUTPUT_ROWS}, {OUTPUT_COLUMNS}], "
+                f"got {output.dtype} {list(output.shape)}"
+            )
+        return output
+
+    def _decode_e2e(
+        self,
+        output: np.ndarray,
+        ratio: float,
+        pad: tuple[float, float],
+        original_size: tuple[int, int],
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        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():
+            bad = model_classes[(model_classes < 0) | (model_classes >= len(self.model_names))]
+            raise InferenceInvariantError(f"e2e model class id outside range: {bad[:8].tolist()}")
+
+        active = scores > 0.0
+        reversed_rows = active & ((boxes[:, 2] < boxes[:, 0]) | (boxes[:, 3] < boxes[:, 1]))
+        if reversed_rows.any():
+            raise InferenceInvariantError(f"e2e output contains reversed boxes: {boxes[reversed_rows][:4].tolist()}")
+        mapped = np.fromiter((int(class_id) in self.model_to_element for class_id in model_classes), bool, len(rows))
+        keep = mapped & (scores >= np.float32(self.conf_threshold))
+        if not keep.any():
+            return (
+                np.empty((0, 4), np.float32),
+                np.empty((0,), np.float32),
+                np.empty((0,), np.int64),
+            )
+        boxes = boxes[keep]
+        scores = scores[keep]
+        model_classes = model_classes[keep]
+        element_classes = np.fromiter(
+            (self.model_to_element[int(class_id)] for class_id in model_classes),
+            dtype=np.int64,
+            count=len(model_classes),
+        )
+        boxes[:, [0, 2]] = (boxes[:, [0, 2]] - pad[0]) / ratio
+        boxes[:, [1, 3]] = (boxes[:, [1, 3]] - pad[1]) / ratio
+        width, height = original_size
+        boxes[:, [0, 2]] = np.clip(boxes[:, [0, 2]], 0, width - 1)
+        boxes[:, [1, 3]] = np.clip(boxes[:, [1, 3]], 0, height - 1)
+        self._check_decoded(boxes, scores, element_classes, original_size)
+
+        box_widths = boxes[:, 2] - boxes[:, 0]
+        box_heights = boxes[:, 3] - boxes[:, 1]
+        areas = box_widths * box_heights
+        image_area = float(width * height)
+        sane = (
+            (box_widths >= self.min_side)
+            & (box_heights >= self.min_side)
+            & (areas <= 0.95 * image_area)
+            & (
+                np.maximum(
+                    box_widths / np.maximum(box_heights, 1e-6),
+                    box_heights / np.maximum(box_widths, 1e-6),
+                )
+                <= self.max_aspect_ratio
+            )
+        )
+        boxes, scores, element_classes = boxes[sane], scores[sane], element_classes[sane]
+        if len(scores):
+            order = np.argsort(-scores, kind="stable")[: self.max_det]
+            boxes, scores, element_classes = boxes[order], scores[order], element_classes[order]
+        # Thresholds, model-name checks and all head semantics above stay in
+        # ``classes`` order.  Only the validator-facing ids are remapped here.
+        if len(element_classes):
+            element_classes = np.asarray(self.config.output_cls_ids, dtype=np.int64)[element_classes]
+        return boxes, scores, element_classes
+
+    def _check_decoded(
+        self,
+        boxes: np.ndarray,
+        scores: np.ndarray,
+        class_ids: np.ndarray,
+        original_size: tuple[int, int],
+    ) -> None:
+        if not len(boxes):
+            return
+        width, height = original_size
+        if not np.isfinite(boxes).all():
+            raise InferenceInvariantError("decoded boxes contain non-finite coordinates")
+        if (boxes[:, 2] < boxes[:, 0]).any() or (boxes[:, 3] < boxes[:, 1]).any():
+            raise InferenceInvariantError("decoded boxes are reversed")
+        xs, ys = boxes[:, [0, 2]], boxes[:, [1, 3]]
+        if (xs < 0).any() or (xs > width - 1).any() or (ys < 0).any() or (ys > height - 1).any():
+            raise InferenceInvariantError(f"decoded boxes outside {width}x{height} image")
+        if not np.isfinite(scores).all() or (scores < 0).any() or (scores > 1).any():
+            raise InferenceInvariantError("decoded confidence outside [0, 1]")
+        if (class_ids < 0).any() or (class_ids >= len(self.class_names)).any():
+            raise InferenceInvariantError("decoded cls_id outside element class range")
+
+    @staticmethod
+    def _to_bounding_boxes(
+        boxes: np.ndarray,
+        scores: np.ndarray,
+        class_ids: np.ndarray,
+    ) -> list[BoundingBox]:
+        return [
+            BoundingBox(
+                x1=int(math.floor(box[0])), y1=int(math.floor(box[1])),
+                x2=int(math.ceil(box[2])), y2=int(math.ceil(box[3])),
+                cls_id=int(class_id), conf=float(score),
+            )
+            for box, score, class_id in zip(boxes, scores, class_ids)
+            if box[2] > box[0] and box[3] > box[1]
+        ]
+
+    def _predict_single(self, image: np.ndarray) -> list[BoundingBox]:
+        image = self._check_image(image)
+        blob, ratio, pad, original_size = self._preprocess(image)
+        outputs = self.session.run(self.output_names, {self.input_name: blob})
+        output = self._check_outputs(outputs)
+        arrays = self._decode_e2e(output, ratio, pad, original_size)
+        return self._to_bounding_boxes(*arrays)
+
+    def _warmup(self, iterations: int) -> None:
+        dummy = np.zeros((720, 1280, 3), dtype=np.uint8)
+        original_strict = 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_strict
+
+    def _record_failure(self, frame_id: int, exc: Exception) -> None:
+        self.inference_failures += 1
+        invariant = isinstance(exc, InferenceInvariantError)
+        if invariant:
+            self.invariant_failures += 1
+        print(json.dumps({
+            "event": "inference_failure",
+            "frame_id": frame_id,
+            "error_type": type(exc).__name__,
+            "error": str(exc),
+            "invariant": invariant,
+            "inference_failures": self.inference_failures,
+            "invariant_failures": self.invariant_failures,
+        }, sort_keys=True))
+
+    def predict_batch(
+        self,
+        batch_images: list[np.ndarray],
+        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._record_failure(frame_id, exc)
+                boxes = []
+            results.append(TVFrameResult(frame_id=frame_id, boxes=boxes, keypoints=[]))
+        return results
+
+    def __repr__(self) -> str:
+        return (
+            f"E2EONNXRuntime(session={type(self.session).__name__}, "
+            f"providers={self.session.get_providers()})"
+        )