From f3147496ec23d0481acf3a0732d98abd6a195d60 Mon Sep 17 00:00:00 2001 From: Seeeeeyo Date: Tue, 5 May 2026 12:03:56 -0600 Subject: [PATCH] feat: add ONNX support for RF-DETR and warm-up functionality in PosePipeline --- README.md | 58 ++++++++++++++ demo.py | 42 +++++++++-- instanthmr/__init__.py | 3 +- instanthmr/detector.py | 167 +++++++++++++++++++++++++++++++++++++++++ instanthmr/pipeline.py | 52 +++++++++++-- requirements.txt | 9 ++- tools/bench.py | 25 +++++- 7 files changed, 339 insertions(+), 17 deletions(-) diff --git a/README.md b/README.md index 8306b89..4cd276d 100644 --- a/README.md +++ b/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) diff --git a/demo.py b/demo.py index a46efe7..cb04a61 100755 --- a/demo.py +++ b/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,7 +142,16 @@ 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, @@ -142,7 +159,12 @@ def build_pipeline(args: argparse.Namespace) -> PosePipeline: 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] = [] diff --git a/instanthmr/__init__.py b/instanthmr/__init__.py index 1194ccc..b7bcbc0 100644 --- a/instanthmr/__init__.py +++ b/instanthmr/__init__.py @@ -1,7 +1,7 @@ """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 @@ -9,6 +9,7 @@ "InstantHMR", "HMRPrediction", "RFDETRDetector", + "RFDetrONNXDetector", "PosePipeline", "FrameResult", "JOINT_NAMES", diff --git a/instanthmr/detector.py b/instanthmr/detector.py index c828eb4..029eeab 100644 --- a/instanthmr/detector.py +++ b/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] diff --git a/instanthmr/pipeline.py b/instanthmr/pipeline.py index a4c23a7..25f9308 100644 --- a/instanthmr/pipeline.py +++ b/instanthmr/pipeline.py @@ -14,7 +14,7 @@ import numpy as np -from .detector import RFDETRDetector +from .detector import RFDetrONNXDetector, RFDETRDetector from .inference import InstantHMR, HMRPrediction @@ -56,6 +56,10 @@ class PosePipeline: onnx_path: Path to ``instanthmr.onnx``. device: ``"cuda"``, ``"coreml"`` or ``"cpu"`` for InstantHMR. detector_variant: ``"nano" | "small" | "medium" | "base" | "large"``. + Ignored when *detector_onnx* is set. + detector_onnx: Optional path to an RF-DETR **detection** ONNX model + (``pred_boxes`` / ``pred_logits`` outputs). When set, runs the + detector in ONNXRuntime instead of PyTorch ``rfdetr``. det_confidence: RF-DETR confidence threshold. max_persons: Maximum number of persons returned per frame. detector_stride: Run the detector every ``N`` frames; on the @@ -86,16 +90,25 @@ def __init__( max_persons: int = 5, detector_stride: int = 1, batch_persons: bool = True, + detector_onnx: str | Path | None = None, ): if detector_stride < 1: raise ValueError("detector_stride must be >= 1") self.hmr = InstantHMR(onnx_path, device=device) - self.detector = RFDETRDetector( - variant=detector_variant, - confidence=det_confidence, - max_persons=max_persons, - ) + if detector_onnx is not None: + self.detector = RFDetrONNXDetector( + onnx_path=detector_onnx, + device=device, + confidence=det_confidence, + max_persons=max_persons, + ) + else: + self.detector = RFDETRDetector( + variant=detector_variant, + confidence=det_confidence, + max_persons=max_persons, + ) self._detector_stride = detector_stride self._batch_persons = batch_persons @@ -103,6 +116,33 @@ def __init__( self._last_detections: list[dict] = [] self._last_bbox: Optional[np.ndarray] = None + def warmup(self, image_rgb: np.ndarray, *, runs: int = 2) -> None: + """Compile GPU / ONNXRuntime kernels before timed inference. + + Runs the full detector + HMR path *runs* times (default 2). If no + person is detected, performs one extra InstantHMR forward with a + synthetic centre crop bbox so the HMR ONNX graph still executes — + otherwise cold-start cost would hit the first real detection frame. + + Resets stride state (``_frame_idx`` and cached bboxes) afterwards so + video/camera timing aligns with frame 0. + """ + image_rgb = np.ascontiguousarray(image_rgb) + last: FrameResult | None = None + for _ in range(max(1, runs)): + last = self.predict(image_rgb) + if last is None or not last.persons: + h, w = image_rgb.shape[:2] + fb = np.array( + [w * 0.25, h * 0.25, w * 0.75, h * 0.75], + dtype=np.float32, + ) + self.hmr.predict(image_rgb, bbox=fb, confidence=1.0) + + self._frame_idx = 0 + self._last_detections = [] + self._last_bbox = None + def predict(self, image_rgb: np.ndarray) -> FrameResult: """Detect all persons in *image_rgb* and run InstantHMR on each.""" t_total_start = time.perf_counter() diff --git a/requirements.txt b/requirements.txt index 1644ca1..0454ab1 100644 --- a/requirements.txt +++ b/requirements.txt @@ -16,9 +16,12 @@ rerun-sdk>=0.21 torch>=2.2 torchvision -# Mac ships CoreML EP inside plain onnxruntime; Linux GPU wants -# onnxruntime-gpu (which install.py picks). The CPU wheel below is the -# safe fallback. +# Mac ships CoreML EP inside plain onnxruntime; Linux + NVIDIA needs +# onnxruntime-gpu (see install.py). Do **not** install ``onnxruntime`` and +# ``onnxruntime-gpu`` in the same env — they overwrite each other's files +# and you can end up with a broken import or CPU-only providers. If that +# happens: ``pip uninstall -y onnxruntime onnxruntime-gpu`` then reinstall +# one of them (``onnxruntime-gpu`` for GPU). onnxruntime>=1.19 # ----- MHR body mesh rendering ----- diff --git a/tools/bench.py b/tools/bench.py index 115ee0b..9f57bfc 100644 --- a/tools/bench.py +++ b/tools/bench.py @@ -33,7 +33,12 @@ def parse_args() -> argparse.Namespace: p.add_argument("--model", type=str, default="models/instanthmr.onnx") p.add_argument("--device", type=str, default="cuda") p.add_argument("--detector-variant", type=str, default="medium", - choices=["nano", "small", "medium", "base", "large"]) + choices=["nano", "small", "medium", "base", "large"], + help="PyTorch RF-DETR variant (ignored when --detector-onnx is set).") + p.add_argument( + "--detector-onnx", type=str, default=None, metavar="PATH", + help="RF-DETR detection ONNX (pred_boxes, pred_logits); ONNXRuntime.", + ) p.add_argument("--det-confidence", type=float, default=0.5) p.add_argument("--max-persons", type=int, default=5) p.add_argument("--detector-stride", type=int, default=1, @@ -74,6 +79,10 @@ def main() -> None: height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) print(f"video: {args.video} {width}x{height} {total} frames @ {src_fps:.1f} fps") + det_onnx = Path(args.detector_onnx) if args.detector_onnx else None + if det_onnx is not None and not det_onnx.exists(): + sys.exit(f"[error] RF-DETR ONNX not found: {det_onnx}") + pipeline = PosePipeline( onnx_path=args.model, device=args.device, @@ -82,9 +91,21 @@ def main() -> None: max_persons=args.max_persons, detector_stride=args.detector_stride, batch_persons=not args.no_batch_persons, + detector_onnx=det_onnx, ) print(f"InstantHMR EP: {pipeline.hmr.active_provider}") - print(f"detector variant: {args.detector_variant}") + dprov = getattr(pipeline.detector, "active_provider", None) + if dprov is not None: + print(f"RF-DETR ONNX EP: {dprov} ({det_onnx})") + else: + print(f"detector variant: {args.detector_variant}") + + ret0, bgr0 = cap.read() + if ret0: + rgb0 = cv2.cvtColor(bgr0, cv2.COLOR_BGR2RGB) + print("GPU warmup …", flush=True) + pipeline.warmup(rgb0) + cap.set(cv2.CAP_PROP_POS_FRAMES, 0) det: list[float] = [] hmr: list[float] = []