scorevision: push artifact

miner.py

@@ -0,0 +1,178 @@
+"""TurboVision public-track miner for manak0/Detect-worksite-ppe-compliance-detection-probe-set.
+
+Model: YOLO26n exported to ONNX (input 1x3x704x704, NMS embedded, output 1x300x6
+rows of x1, y1, x2, y2, score, class in letterboxed pixel space).
+
+cls_id in the response is the index into the element's manifest `objects` list.
+"""
+
+from pathlib import Path
+
+import cv2
+import numpy as np
+import onnxruntime as ort
+from numpy import ndarray
+from pydantic import BaseModel
+
+MANIFEST_CLASSES = [
+    "safety helmet",
+    "high-visibility safety vest",
+    "protective gloves",
+    "protective eyewear",
+]
+
+# Per-class confidence thresholds, indexed like MANIFEST_CLASSES.
+CLASS_CONF = np.array([0.40, 0.40, 0.40, 0.30], dtype=np.float32)
+
+MODEL_FILENAME = "weights.onnx"
+LETTERBOX_FILL = 114
+
+
+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 _letterbox(image: ndarray, size: int) -> tuple[ndarray, float, int, int]:
+    """Ultralytics-style centred letterbox to a size x size canvas."""
+    height, width = image.shape[:2]
+    ratio = min(size / height, size / width)
+    new_w = int(round(width * ratio))
+    new_h = int(round(height * ratio))
+    if (new_w, new_h) != (width, height):
+        image = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
+    pad_w = (size - new_w) / 2
+    pad_h = (size - new_h) / 2
+    top, bottom = int(round(pad_h - 0.1)), int(round(pad_h + 0.1))
+    left, right = int(round(pad_w - 0.1)), int(round(pad_w + 0.1))
+    canvas = cv2.copyMakeBorder(
+        image,
+        top,
+        bottom,
+        left,
+        right,
+        cv2.BORDER_CONSTANT,
+        value=(LETTERBOX_FILL, LETTERBOX_FILL, LETTERBOX_FILL),
+    )
+    return canvas, ratio, left, top
+
+
+class Miner:
+    def __init__(self, path_hf_repo: Path) -> None:
+        model_path = Path(path_hf_repo) / MODEL_FILENAME
+        if not model_path.is_file():
+            raise FileNotFoundError(f"{MODEL_FILENAME} not found in {path_hf_repo}")
+
+        options = ort.SessionOptions()
+        options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+        options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
+        options.intra_op_num_threads = 2
+        options.inter_op_num_threads = 1
+
+        providers = ["CPUExecutionProvider"]
+        if "CUDAExecutionProvider" in ort.get_available_providers():
+            providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
+        try:
+            self.session = ort.InferenceSession(
+                str(model_path), sess_options=options, providers=providers
+            )
+        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"]
+            )
+
+        model_input = self.session.get_inputs()[0]
+        self.input_name = model_input.name
+        self.input_size = int(model_input.shape[2])
+        self.output_name = self.session.get_outputs()[0].name
+        self._check_class_order()
+
+        warmup = np.zeros((self.input_size, self.input_size, 3), dtype=np.uint8)
+        for _ in range(2):
+            self._infer_one(warmup)
+        print(f"Miner ready: providers={self.session.get_providers()} input={self.input_size}")
+
+    def __repr__(self) -> str:
+        return (
+            f"PPE YOLO26n ONNX input={self.input_size} "
+            f"providers={self.session.get_providers()}"
+        )
+
+    def _check_class_order(self) -> None:
+        names = self.session.get_modelmeta().custom_metadata_map.get("names", "")
+        for idx, name in enumerate(MANIFEST_CLASSES):
+            if f"{idx}: '{name}'" not in names:
+                raise ValueError(f"ONNX class order does not match manifest: {names}")
+
+    def _infer_one(self, image: ndarray) -> list[BoundingBox]:
+        canvas, ratio, pad_x, pad_y = _letterbox(image, self.input_size)
+        rgb = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB)
+        tensor = np.ascontiguousarray(rgb.transpose(2, 0, 1)[None], dtype=np.float32) / 255.0
+        raw = self.session.run([self.output_name], {self.input_name: tensor})[0][0]
+
+        scores = raw[:, 4]
+        cls_ids = raw[:, 5].astype(np.int64)
+        known = (cls_ids >= 0) & (cls_ids < len(MANIFEST_CLASSES))
+        keep = known & (scores >= CLASS_CONF[np.clip(cls_ids, 0, len(CLASS_CONF) - 1)])
+        if not keep.any():
+            return []
+
+        boxes = raw[keep, :4].astype(np.float64)
+        boxes[:, [0, 2]] = (boxes[:, [0, 2]] - pad_x) / ratio
+        boxes[:, [1, 3]] = (boxes[:, [1, 3]] - pad_y) / ratio
+        height, width = image.shape[:2]
+        boxes[:, [0, 2]] = boxes[:, [0, 2]].clip(0, width)
+        boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, height)
+
+        results: list[BoundingBox] = []
+        for (x1, y1, x2, y2), score, cls_id in zip(boxes, scores[keep], cls_ids[keep]):
+            x1, y1, x2, y2 = (int(round(v)) for v in (x1, y1, x2, y2))
+            if x2 <= x1 or y2 <= y1:
+                continue
+            results.append(
+                BoundingBox(
+                    x1=x1, y1=y1, x2=x2, y2=y2, cls_id=int(cls_id), conf=float(score)
+                )
+            )
+        return results
+
+    def predict_batch(
+        self,
+        batch_images: list[ndarray],
+        offset: int,
+        n_keypoints: int,
+    ) -> list[TVFrameResult]:
+        results: list[TVFrameResult] = []
+        for index, image in enumerate(batch_images):
+            try:
+                boxes = self._infer_one(image)
+            except Exception as exc:
+                print(f"Inference failed on frame {offset + index}: {exc}")
+                boxes = []
+            results.append(
+                TVFrameResult(
+                    frame_id=offset + index,
+                    boxes=boxes,
+                    polygons=[],
+                    keypoints=[(0, 0)] * max(0, int(n_keypoints)),
+                )
+            )
+        return results