Spaces:
Running on Zero
Running on Zero
File size: 4,458 Bytes
2874635 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | #!/usr/bin/env python3
"""
infer_torch.py — ZeroGPU detection backend for the HF Space (v10 .pt).
Why a raw forward instead of ultralytics' predict():
spatial_logic.postprocess_onnx does conf-filter -> class-agnostic NMS ->
**top-3** extraction, and those top-3 slots are what sphinx_corrector's
Viterbi consumes. model.predict() returns top-1 only, after its own NMS,
which would gut the entire NLP correction layer. DetectionModel's raw
forward returns the same [1, 4+nc, N] tensor the ONNX export (nms=False)
produces — pixel cxcywh in rows 0-3, independent sigmoid class scores in
rows 4.. — so postprocess_onnx is reused UNCHANGED.
Verified against artifacts/best_modelv10.pt: ultralytics 8.4.102, yolo11l,
imgsz 1024, nc=150, class order identical to class_map50_v9.json (0
mismatches, transmutations m4/f34/o29 still at idx 30/61/70). So the tensor
is [1, 154, 21504], exactly as v9.
ZeroGPU rules honoured here:
* `import spaces` before torch.
* The model is built on CPU at module scope and only moved to CUDA INSIDE
the @spaces.GPU function — ZeroGPU patches CUDA init, so a module-scope
.to('cuda') fails.
* postprocess_onnx runs OUTSIDE the GPU function: quota is duration-based,
so we hold the GPU only for the forward pass, and we avoid pickling
Detection dataclasses back across the fork boundary.
"""
from __future__ import annotations
import os
from pathlib import Path
from typing import Callable
import numpy as np
import spaces # must precede torch (ZeroGPU)
import torch
from ultralytics import YOLO
import spatial_logic as SL
ROOT = Path(__file__).parent
WEIGHTS = Path(os.getenv('SPHINX_WEIGHTS',
ROOT / 'artifacts' / 'best_modelv10.pt'))
GPU_SECONDS = int(os.getenv('SPHINX_GPU_SECONDS', '60'))
# Built once, on CPU.
# .eval() — the detect head only emits the concatenated inference tensor in
# eval mode; training mode returns the per-stride feature maps.
# .float() — best_modelv10.pt stores a MIX of HalfStorage and FloatStorage
# tensors (verified by inspecting the checkpoint), so feeding a
# float32 input would raise "expected scalar type Half but found
# Float". Upcasting everything to fp32 removes the dtype mismatch
# and keeps this backend numerically comparable to the fp32 ONNX
# path the pipeline was validated against.
_model = YOLO(str(WEIGHTS)).model.float().eval()
for _p in _model.parameters():
_p.requires_grad_(False)
def _canvas_to_tensor(canvas: np.ndarray) -> torch.Tensor:
"""BGR uint8 HWC -> RGB float32 NCHW in [0,1] (matches make_onnx_infer_fn)."""
x = canvas[:, :, ::-1].astype(np.float32) / 255.0 # BGR -> RGB
x = np.ascontiguousarray(x.transpose(2, 0, 1))[None] # HWC -> NCHW
return torch.from_numpy(x)
@spaces.GPU(duration=GPU_SECONDS)
def _forward(canvas: np.ndarray) -> np.ndarray:
"""Raw forward on the letterboxed canvas. Returns [1, 4+nc, N] as numpy."""
dev = 'cuda' if torch.cuda.is_available() else 'cpu'
model = _model.to(dev)
x = _canvas_to_tensor(canvas).to(dev)
with torch.inference_mode():
out = model(x)
# DetectionModel in eval mode returns either the tensor or a
# (tensor, feature_maps) tuple depending on version — normalise.
if isinstance(out, (list, tuple)):
out = out[0]
return out.float().cpu().numpy()
def make_torch_infer_fn(
class_names : list[str],
conf_thresh : float = SL.CONF_THRESHOLD,
iou_thresh : float = SL.NMS_IOU,
imgsz : int = 1024,
) -> Callable[[np.ndarray], list]:
"""
Drop-in replacement for SL.make_onnx_infer_fn — same signature, same
Callable[[bgr], list[Detection]] contract, bboxes in original-image pixels.
Letterbox (never stretch): fragile breakpoint #9. Reuses SL.letterbox so
preprocessing is provably identical to the ONNX path.
"""
def infer(bgr: np.ndarray) -> list:
canvas, scale, dx, dy = SL.letterbox(bgr, imgsz)
raw = _forward(canvas)
dets = SL.postprocess_onnx(raw, class_names, conf_thresh, iou_thresh)
# undo letterbox: original = (model - d) / scale
for d in dets:
x1, y1, x2, y2 = d.bbox
d.bbox = ((x1 - dx) / scale, (y1 - dy) / scale,
(x2 - dx) / scale, (y2 - dy) / scale)
return dets
return infer
|