scorevision: push artifact

miner.py

@@ -1,175 +1,339 @@
+"""
+Score Vision SN44 — VehicleDetect miner endpoint.
+
+Class mapping (output indices):
+  0 = car        (COCO class 2)
+  1 = bus        (COCO class 5)
+  2 = truck      (COCO class 7)
+  3 = motorcycle (COCO class 3)
+
+Accepts: base64-encoded image or raw image bytes via chutes cord.
+Returns: list of {bbox: [x1,y1,x2,y2], score: float, class_id: int, class_name: str}
+
+CUDA fix: onnxruntime-gpu finds cuDNN via ldconfig (registered during image build),
+          with ctypes preload as belt-and-suspenders fallback.
+"""
+
+from __future__ import annotations
+
+import base64
+import io
+import os
+import time
 from pathlib import Path
-import math
+from typing import Any
 
+import ctypes
 import cv2
 import numpy as np
-import onnxruntime as ort
-from numpy import ndarray
-from pydantic import BaseModel
-
-
-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:
-    """
-    Auto-generated by subnet_bridge from a Manako element repo.
-    This miner is intentionally self-contained for chute import restrictions.
-    """
-
-    def __init__(self, path_hf_repo: Path) -> None:
-        self.path_hf_repo = path_hf_repo
-        self.class_names = ['person']
-        self.session = ort.InferenceSession(
-            str(path_hf_repo / "weights.onnx"),
-            providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
+from PIL import Image
+
+# ── cuDNN preload (belt-and-suspenders fallback) ──────────────────────────────
+# Primary fix is ldconfig at image build time (see Image builder below).
+# This ctypes preload catches any edge cases where ld.so.cache isn't used.
+def _preload_cuda_libs() -> None:
+    _NVIDIA = "/usr/local/lib/python3.12/dist-packages/nvidia"
+    _LIBS = [
+        "/usr/lib/x86_64-linux-gnu/libcuda.so.1",         # driver stub — must be first
+        f"{_NVIDIA}/cublas/lib/libcublasLt.so.12",
+        f"{_NVIDIA}/cublas/lib/libcublas.so.12",
+        f"{_NVIDIA}/cudnn/lib/libcudnn.so.9",
+    ]
+    for path in _LIBS:
+        if os.path.exists(path):
+            try:
+                ctypes.CDLL(path, mode=ctypes.RTLD_GLOBAL)
+            except OSError:
+                pass
+
+_preload_cuda_libs()
+
+import onnxruntime as ort  # noqa: E402 — must come after preload
+
+# ── Constants ────────────────────────────────────────────────────────────────
+MODEL_DIR   = Path(__file__).parent
+WEIGHTS     = MODEL_DIR / "weights.onnx"
+IMG_SIZE    = 640
+CONF_THRESH = 0.55  # sweep: max composite score (0.60×mAP + 0.40×FP_score) at conf=0.55
+IOU_THRESH  = 0.45
+
+# COCO class index → submission class index
+COCO_TO_OUT: dict[int, int] = {2: 0, 5: 1, 7: 2, 3: 3}
+COCO_VEHICLE_IDX = list(COCO_TO_OUT.keys())
+OUT_NAMES = ["car", "bus", "truck", "motorcycle"]
+
+# ── Model loader (singleton) ─────────────────────────────────────────────────
+_SESSION: ort.InferenceSession | None = None
+
+
+def get_session() -> ort.InferenceSession:
+    global _SESSION
+    if _SESSION is None:
+        opts = ort.SessionOptions()
+        opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+        opts.enable_mem_pattern = True
+        opts.enable_mem_reuse   = True
+        cuda_opts = {
+            "device_id": 0,
+            "arena_extend_strategy": "kNextPowerOfTwo",
+            "gpu_mem_limit": 2 * 1024 ** 3,
+            "cudnn_conv_algo_search": "EXHAUSTIVE",
+            "do_copy_in_default_stream": True,
+        }
+        _SESSION = ort.InferenceSession(
+            str(WEIGHTS),
+            sess_options=opts,
+            providers=[
+                ("CUDAExecutionProvider", cuda_opts),
+                "CPUExecutionProvider",
+            ],
+        )
+        provider = _SESSION.get_providers()[0]
+        print(f"[miner] Model loaded. Provider: {provider}", flush=True)
+    return _SESSION
+
+
+# ── Preprocessing ────────────────────────────────────────────────────────────
+
+def letterbox(img: np.ndarray, size: int = IMG_SIZE) -> tuple[np.ndarray, float, int, int]:
+    h, w = img.shape[:2]
+    r = min(size / h, size / w)
+    new_w, new_h = int(round(w * r)), int(round(h * r))
+    img_r = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
+    dw, dh = size - new_w, size - new_h
+    pad_l, pad_t = dw // 2, dh // 2
+    img_p = cv2.copyMakeBorder(
+        img_r, pad_t, dh - pad_t, pad_l, dw - pad_l,
+        cv2.BORDER_CONSTANT, value=(114, 114, 114),
+    )
+    return img_p, r, pad_l, pad_t
+
+
+def preprocess(img_bgr: np.ndarray) -> tuple[np.ndarray, float, int, int]:
+    img_p, ratio, pad_l, pad_t = letterbox(img_bgr)
+    img_rgb = cv2.cvtColor(img_p, cv2.COLOR_BGR2RGB)
+    inp = img_rgb.transpose(2, 0, 1).astype(np.float32) * (1.0 / 255.0)
+    return np.ascontiguousarray(inp[np.newaxis]), ratio, pad_l, pad_t
+
+
+# ── NMS ──────────────────────────────────────────────────────────────────────
+
+def nms(boxes: np.ndarray, scores: np.ndarray, iou_thresh: float = IOU_THRESH) -> list[int]:
+    if not len(boxes):
+        return []
+    x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
+    areas = (x2 - x1) * (y2 - y1)
+    order = scores.argsort()[::-1]
+    keep: list[int] = []
+    while len(order):
+        i = order[0]
+        keep.append(int(i))
+        xx1 = np.maximum(x1[i], x1[order[1:]])
+        yy1 = np.maximum(y1[i], y1[order[1:]])
+        xx2 = np.minimum(x2[i], x2[order[1:]])
+        yy2 = np.minimum(y2[i], y2[order[1:]])
+        inter = np.maximum(0, xx2 - xx1) * np.maximum(0, yy2 - yy1)
+        iou = inter / (areas[i] + areas[order[1:]] - inter + 1e-7)
+        order = order[1:][iou <= iou_thresh]
+    return keep
+
+
+# ── Postprocessing ───────────────────────────────────────────────────────────
+
+def postprocess(
+    raw: np.ndarray,
+    ratio: float,
+    pad_l: int,
+    pad_t: int,
+    orig_w: int,
+    orig_h: int,
+) -> list[dict[str, Any]]:
+    pred = raw  # [84, 8400]
+    veh_row_idx = np.array([4 + c for c in COCO_VEHICLE_IDX])
+    max_veh_score = pred[veh_row_idx].max(axis=0)
+    mask = max_veh_score > CONF_THRESH
+    if not mask.any():
+        return []
+
+    pred_f = pred[:, mask]
+    cx, cy, bw, bh = pred_f[0], pred_f[1], pred_f[2], pred_f[3]
+
+    x1 = np.clip((cx - bw / 2 - pad_l) / ratio, 0, orig_w)
+    y1 = np.clip((cy - bh / 2 - pad_t) / ratio, 0, orig_h)
+    x2 = np.clip((cx + bw / 2 - pad_l) / ratio, 0, orig_w)
+    y2 = np.clip((cy + bh / 2 - pad_t) / ratio, 0, orig_h)
+    boxes = np.stack([x1, y1, x2, y2], axis=1)
+
+    results: list[dict[str, Any]] = []
+    for coco_cls in COCO_VEHICLE_IDX:
+        scores = pred_f[4 + coco_cls]
+        cls_mask = scores > CONF_THRESH
+        if not cls_mask.any():
+            continue
+        keep = nms(boxes[cls_mask], scores[cls_mask])
+        out_cls = COCO_TO_OUT[coco_cls]
+        for k in keep:
+            box = boxes[cls_mask][k]
+            results.append({
+                "bbox": [
+                    float(box[0]), float(box[1]),
+                    float(box[2]), float(box[3]),
+                ],
+                "score": float(scores[cls_mask][k]),
+                "class_id": out_cls,
+                "class_name": OUT_NAMES[out_cls],
+            })
+    return results
+
+
+# ── Image decoding helpers ───────────────────────────────────────────────────
+
+def decode_image(data: bytes | str) -> np.ndarray:
+    if isinstance(data, str):
+        data = base64.b64decode(data)
+    elif isinstance(data, (bytes, bytearray)):
+        try:
+            data = base64.b64decode(data)
+        except Exception:
+            pass
+    arr = np.frombuffer(data, dtype=np.uint8)
+    img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
+    if img is None:
+        pil = Image.open(io.BytesIO(data)).convert("RGB")
+        img = cv2.cvtColor(np.array(pil), cv2.COLOR_RGB2BGR)
+    return img
+
+
+# ── Core predict function ────────────────────────────────────────────────────
+
+def predict(image_data: bytes | str | np.ndarray) -> dict[str, Any]:
+    sess = get_session()
+
+    if isinstance(image_data, np.ndarray):
+        img_bgr = image_data
+    else:
+        img_bgr = decode_image(image_data)
+
+    orig_h, orig_w = img_bgr.shape[:2]
+    inp, ratio, pad_l, pad_t = preprocess(img_bgr)
+
+    t0 = time.perf_counter()
+    outputs = sess.run(None, {"images": inp})
+    infer_ms = (time.perf_counter() - t0) * 1000.0
+
+    raw = outputs[0][0]   # [84, 8400]
+    detections = postprocess(raw, ratio, pad_l, pad_t, orig_w, orig_h)
+
+    return {
+        "detections": detections,
+        "inference_ms": round(infer_ms, 3),
+        "provider": sess.get_providers()[0],
+    }
+
+
+# ── Chutes cord wrapper ──────────────────────────────────────────────────────
+
+try:
+    from chutes.chute import Chute
+    from chutes.chute.node_selector import NodeSelector
+    from chutes.image import Image as ChuteImage
+
+    chute_image = (
+        ChuteImage(
+            username="lculpitt",
+            name="vehicle-detect-sn44",
+            tag="v4-cuda",
+            readme=(Path(__file__).parent / "README.md").read_text(),
+        )
+        .from_base("parachutes/python:3.12")
+        .run_command("pip install --upgrade setuptools wheel")
+        .run_command(
+            "pip install 'numpy>=1.23' 'onnxruntime-gpu>=1.16' "
+            "'opencv-python-headless>=4.7' 'pillow>=9.5' "
+            "'huggingface_hub>=0.19.4' 'pydantic>=2.0' "
+            "'pyyaml>=6.0' 'aiohttp>=3.9'"
         )
-        self.input_name = self.session.get_inputs()[0].name
-        input_shape = self.session.get_inputs()[0].shape
-        # expected [N, C, H, W]
-        self.input_h = int(input_shape[2])
-        self.input_w = int(input_shape[3])
-        self.conf_threshold = 0.70  # sweep-optimised: max composite 0.65×mAP+0.35×FP_score
-        self.iou_threshold = 0.45
-
-    def __repr__(self) -> str:
-        return f"ONNX Miner session={type(self.session).__name__} classes={len(self.class_names)}"
-
-    def _preprocess(self, image_bgr: ndarray) -> tuple[np.ndarray, tuple[int, int]]:
-        h, w = image_bgr.shape[:2]
-        rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
-        resized = cv2.resize(rgb, (self.input_w, self.input_h))
-        x = resized.astype(np.float32) / 255.0
-        x = np.transpose(x, (2, 0, 1))[None, ...]
-        return x, (h, w)
-
-    def _normalize_predictions(self, raw: np.ndarray) -> np.ndarray:
-        # Common ultralytics export shapes:
-        # - [1, C, N] where C=4+num_classes
-        # - [1, N, C]
-        pred = raw[0]
-        if pred.ndim != 2:
-            raise ValueError(f"Unexpected prediction shape: {raw.shape}")
-        if pred.shape[0] < pred.shape[1]:
-            pred = pred.transpose(1, 0)
-        return pred
-
-    def _nms(self, dets: list[tuple[float, float, float, float, float, int]]) -> list[tuple[float, float, float, float, float, int]]:
-        if not dets:
-            return []
-
-        boxes = np.array([[d[0], d[1], d[2], d[3]] for d in dets], dtype=np.float32)
-        scores = np.array([d[4] for d in dets], dtype=np.float32)
-        order = scores.argsort()[::-1]
-        keep = []
-
-        while order.size > 0:
-            i = order[0]
-            keep.append(i)
-
-            xx1 = np.maximum(boxes[i, 0], boxes[order[1:], 0])
-            yy1 = np.maximum(boxes[i, 1], boxes[order[1:], 1])
-            xx2 = np.minimum(boxes[i, 2], boxes[order[1:], 2])
-            yy2 = np.minimum(boxes[i, 3], boxes[order[1:], 3])
-
-            w = np.maximum(0.0, xx2 - xx1)
-            h = np.maximum(0.0, yy2 - yy1)
-            inter = w * h
-
-            area_i = (boxes[i, 2] - boxes[i, 0]) * (boxes[i, 3] - boxes[i, 1])
-            area_rest = (boxes[order[1:], 2] - boxes[order[1:], 0]) * (boxes[order[1:], 3] - boxes[order[1:], 1])
-            union = np.maximum(area_i + area_rest - inter, 1e-6)
-            iou = inter / union
-
-            remaining = np.where(iou <= self.iou_threshold)[0]
-            order = order[remaining + 1]
-
-        return [dets[idx] for idx in keep]
-
-    def _infer_single(self, image_bgr: ndarray) -> list[BoundingBox]:
-        inp, (orig_h, orig_w) = self._preprocess(image_bgr)
-        out = self.session.run(None, {self.input_name: inp})[0]
-        pred = self._normalize_predictions(out)
-
-        if pred.shape[1] < 5:
-            return []
-
-        boxes = pred[:, :4]
-        cls_scores = pred[:, 4:]
-
-        if cls_scores.shape[1] == 0:
-            return []
-
-        cls_ids = np.argmax(cls_scores, axis=1)
-        confs = np.max(cls_scores, axis=1)
-        keep = confs >= self.conf_threshold
-
-        boxes = boxes[keep]
-        confs = confs[keep]
-        cls_ids = cls_ids[keep]
-
-        if boxes.shape[0] == 0:
-            return []
-
-        sx = orig_w / float(self.input_w)
-        sy = orig_h / float(self.input_h)
-
-        dets: list[tuple[float, float, float, float, float, int]] = []
-        for i in range(boxes.shape[0]):
-            cx, cy, bw, bh = boxes[i].tolist()
-            x1 = (cx - bw / 2.0) * sx
-            y1 = (cy - bh / 2.0) * sy
-            x2 = (cx + bw / 2.0) * sx
-            y2 = (cy + bh / 2.0) * sy
-            dets.append((x1, y1, x2, y2, float(confs[i]), int(cls_ids[i])))
-
-        dets = self._nms(dets)
-
-        out_boxes: list[BoundingBox] = []
-        for x1, y1, x2, y2, conf, cls_id in dets:
-            ix1 = max(0, min(orig_w, math.floor(x1)))
-            iy1 = max(0, min(orig_h, math.floor(y1)))
-            ix2 = max(0, min(orig_w, math.ceil(x2)))
-            iy2 = max(0, min(orig_h, math.ceil(y2)))
-            out_boxes.append(
-                BoundingBox(
-                    x1=ix1,
-                    y1=iy1,
-                    x2=ix2,
-                    y2=iy2,
-                    cls_id=cls_id,
-                    conf=max(0.0, min(1.0, conf)),
-                )
-            )
-        return out_boxes
-
-    def predict_batch(
-        self,
-        batch_images: list[ndarray],
-        offset: int,
-        n_keypoints: int,
-    ) -> list[TVFrameResult]:
-        results: list[TVFrameResult] = []
-        for idx, image in enumerate(batch_images):
-            boxes = self._infer_single(image)
-            keypoints = [(0, 0) for _ in range(max(0, int(n_keypoints)))]
-            results.append(
-                TVFrameResult(
-                    frame_id=offset + idx,
-                    boxes=boxes,
-                    keypoints=keypoints,
-                )
-            )
-        return results
+        # Bake cuDNN/cuBLAS paths into the image as Docker ENV so onnxruntime
+        # CUDAExecutionProvider finds libcudnn.so.9 on every node at container start.
+        .with_env(
+            "LD_LIBRARY_PATH",
+            "/usr/local/lib/python3.12/dist-packages/nvidia/cudnn/lib"
+            ":/usr/local/lib/python3.12/dist-packages/nvidia/cublas/lib",
+        )
+    )
+
+    chute = Chute(
+        username="lculpitt",
+        name="vehicle-detect-sn44",
+        tagline="YOLO11n vehicle detector — car, bus, truck, motorcycle",
+        readme=(Path(__file__).parent / "README.md").read_text(),
+        image=chute_image,
+        concurrency=4,
+        max_instances=5,
+        shutdown_after_seconds=300,
+        scaling_threshold=0.5,
+        node_selector=NodeSelector(
+            gpu_count=1,
+            min_vram_gb_per_gpu=16,
+            # All CUDA 12.x, all $0.40–$0.85/hr (within 2.5× spread from cheapest)
+            include=["4090", "a40", "a6000", "l40", "l40s"],
+        ),
+    )
+
+    @chute.cord(path="/predict", method="POST")
+    async def predict_cord(image_b64: str) -> dict:
+        """
+        POST /predict
+        Body: {"image_b64": "<base64-encoded image>"}
+        Returns detection JSON.
+        """
+        return predict(image_b64)
+
+except ImportError:
+    pass
+
+
+# ── Local test ───────────────────────────────────────────────────────────────
+
+if __name__ == "__main__":
+    import sys
+
+    print("=" * 55)
+    print("  miner.py — local smoke test")
+    print("=" * 55)
+
+    dummy_bgr = np.full((720, 1280, 3), 128, dtype=np.uint8)
+    cv2.rectangle(dummy_bgr, (100, 100), (400, 300), (0, 255, 0), 3)
+
+    if len(sys.argv) > 1:
+        loaded = cv2.imread(sys.argv[1])
+        if loaded is not None:
+            dummy_bgr = loaded
+            print(f"  Using image: {sys.argv[1]}  ({loaded.shape[1]}x{loaded.shape[0]})")
+        else:
+            print(f"  Could not load {sys.argv[1]}, using dummy.")
+    else:
+        print("  Using synthetic 1280x720 dummy image.")
+
+    result = predict(dummy_bgr)
+    print(f"\n  Provider     : {result['provider']}")
+    print(f"  Inference    : {result['inference_ms']:.2f} ms")
+    print(f"  Detections   : {len(result['detections'])}")
+    for d in result["detections"]:
+        x1, y1, x2, y2 = [round(v, 1) for v in d["bbox"]]
+        print(f"    [{d['class_id']}] {d['class_name']:12s}  score={d['score']:.3f}  "
+              f"bbox=[{x1},{y1},{x2},{y2}]")
+
+    print("\n  Latency benchmark (50 runs)...")
+    times = []
+    for _ in range(50):
+        t0 = time.perf_counter()
+        predict(dummy_bgr)
+        times.append((time.perf_counter() - t0) * 1000)
+    times.sort()
+    p50, p95 = times[25], times[47]
+    fps = 1000.0 / p50
+    print(f"  P50={p50:.2f}ms  P95={p95:.2f}ms  FPS={fps:.1f}")
+    print(f"  Target >=30 FPS  : {'PASS' if fps >= 30 else 'FAIL'}")
+    print(f"  Target P95<50ms  : {'PASS' if p95 < 50 else 'FAIL'}")
+    print("=" * 55)