feat: add ONNX support for RF-DETR and warm-up functionality

README.md

@@ -88,6 +88,60 @@ The demo opens a Rerun viewer with the source image + 2D skeleton, a live
 RF-DETR / InstantHMR / total-latency plot, and the 3D scene with the
 predicted camera frustum.
 
+## Recent improvements (RF-DETR ONNX, GPU ORT, warm-up)
+
+These options sit on top of the default PyTorch RF-DETR + InstantHMR stack.
+
+### RF-DETR via ONNX (`--detector-onnx`)
+
+You can run the **person detector** through ONNXRuntime instead of PyTorch
+`rfdetr` — useful when you already have an RF-DETR detection export or want
+the detector on the same ORT path as InstantHMR:
+
+```bash
+python demo.py --video clip.mp4 --detector-onnx models/rf-detr-medium.onnx
+```
+
+Pre-exported RF-DETR ONNX checkpoints (nano through xxlarge, COCO/O365, segmentation variants) are available at **[PierreMarieCurie/rf-detr-onnx](https://huggingface.co/PierreMarieCurie/rf-detr-onnx/tree/main)** on Hugging Face — download the `.onnx` you want and pass its path to `--detector-onnx`.
+
+The ONNX graph must expose a **single** float32 image input (`NCHW`,
+ImageNet-normalised RGB at the export resolution, typically square **576×576**)
+and outputs named **`pred_boxes`** (cxcywh) and **`pred_logits`**, as in
+Roboflow’s RF-DETR ONNX export. **`--detector-variant` is ignored** when this
+flag is set.
+
+### ONNX Runtime on NVIDIA GPUs
+
+For InstantHMR and the optional RF-DETR ONNX detector to use the GPU, you need
+the **`onnxruntime-gpu`** wheel (Linux + NVIDIA: `python install.py` installs
+it together with the matching Torch CUDA stack). **Do not install** the plain
+CPU package **`onnxruntime`** in the same environment as **`onnxruntime-gpu`**
+— they overwrite shared files and you may end up on CPU-only providers or a
+broken install. If that happens: `pip uninstall -y onnxruntime onnxruntime-gpu`
+then reinstall **`onnxruntime-gpu`** (see also the comment above `onnxruntime`
+in `requirements.txt`).
+
+Verify GPU execution providers:
+
+```bash
+python -c "import onnxruntime as ort; print(ort.get_available_providers())"
+```
+
+You should see **`CUDAExecutionProvider`** when CUDA libraries are available.
+
+### Warm-up before inference (`PosePipeline.warmup`)
+
+The first ONNXRuntime **CUDA** run often pays a large one-off cost (kernel
+scheduling, allocator warm-up). **`PosePipeline.warmup(image_rgb, runs=2)`**
+runs the full detector + HMR path before timed frames, resets internal
+stride/caching state, and — if no person is detected on the warm-up frame —
+runs one extra InstantHMR forward with a **synthetic centre crop** so the HMR
+graph always executes once.
+
+The **`demo.py`** runner calls this automatically (you will see
+`Warming up inference …`); **`tools/bench.py`** warm-ups on the first video
+frame and rewinds the capture so benchmarks are not dominated by cold start.
+
 ## MHR body mesh rendering
 
 The demo can render a full dense body mesh for each detected person by
@@ -225,6 +279,10 @@ pipeline = PosePipeline(
     detector_variant="medium",
 )
 
+# Optional: RF-DETR ONNX + explicit GPU warm-up before latency-sensitive loops.
+# pipeline = PosePipeline(..., detector_onnx="models/rf-detr-medium.onnx")
+# pipeline.warmup(image_rgb)
+
 result = pipeline.predict(image_rgb)
 for r in result.persons:
     print(r.joints_3d_cam.shape)   # (70, 3)

demo.py

@@ -72,7 +72,15 @@ def parse_args() -> argparse.Namespace:
     p.add_argument(
         "--detector-variant", type=str, default="medium",
         choices=["nano", "small", "medium", "base", "large"],
-        help="RF-DETR size variant (default: medium).",
+        help="RF-DETR size variant when using PyTorch rfdetr (default: medium). Ignored if --detector-onnx is set.",
+    )
+    p.add_argument(
+        "--detector-onnx", type=str, default=None, metavar="PATH",
+        help=(
+            "Run RF-DETR from this ONNX file via ONNXRuntime instead of PyTorch. "
+            "Must expose outputs pred_boxes and pred_logits (Roboflow RF-DETR export). "
+            "Example: models/rf-detr-medium.onnx"
+        ),
     )
     p.add_argument(
         "--det-confidence", type=float, default=0.5,
@@ -134,15 +142,29 @@ def build_pipeline(args: argparse.Namespace) -> PosePipeline:
         )
         sys.exit(1)
 
-    return PosePipeline(
+    det_onnx: Path | None = None
+    if args.detector_onnx:
+        det_onnx = Path(args.detector_onnx)
+        if not det_onnx.exists():
+            sys.stderr.write(
+                f"\n[error] RF-DETR ONNX not found: {det_onnx}\n\n"
+            )
+            sys.exit(1)
+
+    pipeline = PosePipeline(
         onnx_path=model_path,
         device=args.device,
         detector_variant=args.detector_variant,
         det_confidence=args.det_confidence,
         max_persons=args.max_persons,
         detector_stride=args.detector_stride,
         batch_persons=not args.no_batch_persons,
+        detector_onnx=det_onnx,
     )
+    prov = getattr(pipeline.detector, "active_provider", None)
+    if prov is not None:
+        print(f"RF-DETR ONNX ({det_onnx.name}): {prov}")
+    return pipeline
 
 
 def build_mhr_renderer(args: argparse.Namespace):
@@ -232,9 +254,9 @@ def run_image(args: argparse.Namespace) -> None:
     mhr = build_mhr_renderer(args)
     viz = build_visualizer(args, mhr)
 
-    # Warm-up: first call includes lazy CUDA / ORT initialisation, so a real
-    # latency reading is much more useful from the second call onwards.
-    _ = pipeline.predict(rgb)
+    print("Warming up inference …", flush=True)
+    pipeline.warmup(rgb)
+
     result = pipeline.predict(rgb)
 
     render_ms = viz.log_frame(
@@ -268,6 +290,13 @@ def run_video(args: argparse.Namespace) -> None:
     mhr = build_mhr_renderer(args)
     viz = build_visualizer(args, mhr)
 
+    ret_warm, bgr_warm = cap.read()
+    if ret_warm:
+        rgb_warm = cv2.cvtColor(bgr_warm, cv2.COLOR_BGR2RGB)
+        print("Warming up inference …", flush=True)
+        pipeline.warmup(rgb_warm)
+        cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
+
     frame_idx = 0
     sent = 0
     det_times: list[float] = []
@@ -348,6 +377,9 @@ def run_camera(args: argparse.Namespace) -> None:
     mhr = build_mhr_renderer(args)
     viz = build_visualizer(args, mhr)
 
+    print("Warming up inference …", flush=True)
+    pipeline.warmup(np.zeros((height, width, 3), dtype=np.uint8))
+
     frame_idx = 0
     det_times: list[float] = []
     hmr_times: list[float] = []

instanthmr/__init__.py

@@ -1,14 +1,15 @@
 """InstantHMR — standalone 3D human pose inference + Rerun visualization."""
 
 from .inference import InstantHMR, HMRPrediction
-from .detector import RFDETRDetector
+from .detector import RFDetrONNXDetector, RFDETRDetector
 from .pipeline import PosePipeline, FrameResult
 from .skeleton import JOINT_NAMES, SKELETON_EDGES, NUM_JOINTS, edges_for
 
 __all__ = [
     "InstantHMR",
     "HMRPrediction",
     "RFDETRDetector",
+    "RFDetrONNXDetector",
     "PosePipeline",
     "FrameResult",
     "JOINT_NAMES",

instanthmr/detector.py

@@ -12,11 +12,54 @@
 
 from __future__ import annotations
 
+from pathlib import Path
+
+import cv2
 import numpy as np
 
 
 COCO_PERSON_CLASS_ID = 1  # RF-DETR exposes COCO 1-indexed class ids
 
+_IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32).reshape(1, 1, 3)
+_IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32).reshape(1, 1, 3)
+
+
+def _onnx_providers(device: str) -> list[str]:
+    """Match ``InstantHMR`` ORT provider selection for a given *device* string."""
+    import onnxruntime as ort
+
+    available = set(ort.get_available_providers())
+    wanted: list[str] = []
+    dl = device.lower()
+    if "cuda" in dl and "CUDAExecutionProvider" in available:
+        wanted.append("CUDAExecutionProvider")
+    elif "coreml" in dl and "CoreMLExecutionProvider" in available:
+        wanted.append("CoreMLExecutionProvider")
+    wanted.append("CPUExecutionProvider")
+    return wanted
+
+
+def _sigmoid(x: np.ndarray) -> np.ndarray:
+    """Numerically stable sigmoid."""
+    out = np.empty_like(x, dtype=np.float32)
+    pos = x >= 0
+    out[pos] = 1.0 / (1.0 + np.exp(-x[pos].astype(np.float64))).astype(np.float32)
+    ex = np.exp(x[~pos].astype(np.float64))
+    out[~pos] = (ex / (1.0 + ex)).astype(np.float32)
+    return out
+
+
+def _box_cxcywh_to_xyxy(cxcywh: np.ndarray) -> np.ndarray:
+    """(N, 4) cxcywh in [0, 1] → (N, 4) xyxy in [0, 1]. Same as ``rfdetr`` ``box_ops``."""
+    cxc, yc, bw, bh = np.split(cxcywh.astype(np.float32), 4, axis=-1)
+    bw = np.clip(bw, 0.0, None)
+    bh = np.clip(bh, 0.0, None)
+    x1 = cxc - 0.5 * bw
+    y1 = yc - 0.5 * bh
+    x2 = cxc + 0.5 * bw
+    y2 = yc + 0.5 * bh
+    return np.concatenate([x1, y1, x2, y2], axis=-1)
+
 
 class RFDETRDetector:
     """Person detector backed by RF-DETR.
@@ -97,3 +140,127 @@ def detect(self, image_rgb: np.ndarray) -> list[dict]:
 
         detections.sort(key=lambda d: d["confidence"], reverse=True)
         return detections[: self._max_persons]
+
+
+class RFDetrONNXDetector:
+    """RF-DETR person detector via ONNXRuntime (same I/O contract as ``RFDETRDetector``).
+
+    Expects an RF-DETR *detection* ONNX export with inputs ``input`` shaped
+    ``(1, 3, H, H)`` float32 (ImageNet-normalised RGB) and outputs
+    ``pred_boxes`` / ``pred_logits`` as in Roboflow's RF-DETR ONNX export.
+    Post-processing mirrors ``rfdetr.models.postprocess.PostProcess`` (top-300
+    over all class logits, ``cxcywh`` → ``xyxy``, scale to original image size).
+    """
+
+    def __init__(
+        self,
+        onnx_path: str | Path,
+        device: str = "cuda",
+        confidence: float = 0.5,
+        max_persons: int = 5,
+        num_select: int = 300,
+        providers: list[str] | None = None,
+    ):
+        import onnxruntime as ort
+
+        onnx_path = Path(onnx_path)
+        if not onnx_path.exists():
+            raise FileNotFoundError(f"RF-DETR ONNX not found: {onnx_path}")
+
+        if hasattr(ort, "preload_dlls"):
+            try:
+                ort.preload_dlls()
+            except Exception:
+                pass
+
+        if providers is None:
+            providers = _onnx_providers(device)
+
+        sess_options = ort.SessionOptions()
+        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+        sess_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
+
+        self._session = ort.InferenceSession(
+            str(onnx_path), sess_options=sess_options, providers=providers,
+        )
+        self._active_provider = self._session.get_providers()[0]
+
+        inputs = self._session.get_inputs()
+        if len(inputs) != 1:
+            raise ValueError(
+                f"RF-DETR ONNX expects a single input, got {len(inputs)}: "
+                f"{[i.name for i in inputs]}"
+            )
+        self._in_name = inputs[0].name
+        shape = inputs[0].shape
+        if len(shape) != 4 or shape[1] != 3:
+            raise ValueError(f"Unexpected RF-DETR input shape {shape!r}")
+        h_in, w_in = shape[2], shape[3]
+        if isinstance(h_in, int) and isinstance(w_in, int) and h_in == w_in:
+            self._input_size = int(h_in)
+        else:
+            # Dynamic axes — default matches RF-DETR medium export.
+            self._input_size = 576
+
+        outs = {o.name for o in self._session.get_outputs()}
+        for need in ("pred_boxes", "pred_logits"):
+            if need not in outs:
+                raise ValueError(
+                    f"RF-DETR ONNX missing output {need!r}; have {sorted(outs)}"
+                )
+
+        self._confidence = float(confidence)
+        self._max_persons = int(max_persons)
+        self._num_select = int(num_select)
+
+    @property
+    def variant(self) -> str:
+        return "onnx"
+
+    @property
+    def active_provider(self) -> str:
+        return self._active_provider
+
+    def detect(self, image_rgb: np.ndarray) -> list[dict]:
+        image_rgb = np.ascontiguousarray(image_rgb)
+        h0, w0 = image_rgb.shape[:2]
+        size = self._input_size
+
+        resized = cv2.resize(image_rgb, (size, size), interpolation=cv2.INTER_LINEAR)
+        x = resized.astype(np.float32) / 255.0
+        x = (x - _IMAGENET_MEAN) / _IMAGENET_STD
+        batch = np.transpose(x, (2, 0, 1))[np.newaxis].astype(np.float32)
+
+        raw = self._session.run(None, {self._in_name: batch})
+        name_to = {o.name: v for o, v in zip(self._session.get_outputs(), raw)}
+        pred_boxes = name_to["pred_boxes"][0].astype(np.float32)   # (300, 4) cxcywh
+        pred_logits = name_to["pred_logits"][0].astype(np.float32)  # (300, C)
+
+        prob = _sigmoid(pred_logits)
+        nq, ncls = prob.shape
+        flat = prob.reshape(-1)
+        k = min(self._num_select, flat.size)
+        idx = np.argpartition(-flat, k - 1)[:k]
+        idx = idx[np.argsort(-flat[idx])]
+        scores = flat[idx]
+        topk_boxes = (idx // ncls).astype(np.int64)
+        labels = (idx % ncls).astype(np.int64)
+
+        xyxy_norm = _box_cxcywh_to_xyxy(pred_boxes)
+        boxes_sel = xyxy_norm[topk_boxes]
+        scale = np.array([w0, h0, w0, h0], dtype=np.float32)
+        boxes_abs = boxes_sel * scale
+
+        detections: list[dict] = []
+        for box, lab, sc in zip(boxes_abs, labels, scores):
+            if int(lab) != COCO_PERSON_CLASS_ID:
+                continue
+            if float(sc) < self._confidence:
+                continue
+            detections.append({
+                "bbox": np.asarray(box, dtype=np.float32),
+                "confidence": float(sc),
+            })
+
+        detections.sort(key=lambda d: d["confidence"], reverse=True)
+        return detections[: self._max_persons]