# ------------------------------------------------------------------------ # RF-DETR # Copyright (c) 2025 Roboflow. All Rights Reserved. # Licensed under the Apache License, Version 2.0 [see LICENSE for details] # ------------------------------------------------------------------------ """ONNX Runtime inference helpers for RF-DETR exported models. These functions handle session creation, image preprocessing, and detection decoding without requiring PyTorch or the RF-DETR training stack — only ``onnxruntime``, ``numpy``, ``supervision``, and ``Pillow`` are needed at inference time. """ from __future__ import annotations import contextlib from pathlib import Path from typing import Any import numpy as np from numpy.typing import NDArray from PIL import Image as PILImage from supervision import Detections from rfdetr.export._resize import _bilinear_resize_half_pixel from rfdetr.export._topk import _select_topk_multiclass from rfdetr.utilities.logger import get_logger logger = get_logger() def _create_onnx_session(model_path: str | Path, providers: list[str] | None = None) -> Any: """Load an ONNX model and create an ONNX Runtime inference session. Imports ``onnxruntime`` at call time so that the rest of the package remains usable without it installed. Input and output names / shapes are logged at DEBUG level for troubleshooting. When ``providers`` is ``None``, the session auto-selects the best available backend: CUDA if ``onnxruntime-gpu`` is installed, otherwise CPU (with a warning). Pass an explicit list to pin the backend — useful for benchmarking CPU vs CUDA side-by-side. Args: model_path: Path to the ``.onnx`` model file. providers: Ordered list of ORT execution providers, e.g. ``["CUDAExecutionProvider", "CPUExecutionProvider"]``. When ``None`` (default), the best available provider is selected automatically. Returns: An ``onnxruntime.InferenceSession`` ready for inference. Raises: ImportError: If ``onnxruntime`` is not installed. Examples: .. code-block:: python sess = _create_onnx_session("model.onnx") print(sess.get_inputs()[0].name) """ try: import onnxruntime as ort except ImportError as exc: raise ImportError( "ONNX Runtime inference requires 'onnxruntime'. Install it: `pip install onnxruntime`" ) from exc if providers is None: _preferred = ["CUDAExecutionProvider", "CPUExecutionProvider"] _available = ort.get_available_providers() providers = [p for p in _preferred if p in _available] or ["CPUExecutionProvider"] if providers[0] == "CPUExecutionProvider": logger.warning( "CUDAExecutionProvider not available — running ONNX inference on CPU. " "Install onnxruntime-gpu for GPU acceleration: `pip install onnxruntime-gpu`" ) session = ort.InferenceSession(str(model_path), providers=providers) logger.debug("ONNX Runtime providers in use: %s", session.get_providers()) for inp in session.get_inputs(): logger.debug("Input : name=%s shape=%s type=%s", inp.name, inp.shape, inp.type) for out in session.get_outputs(): logger.debug("Output : name=%s shape=%s type=%s", out.name, out.shape, out.type) return session def _preprocess_pil_to_nchw( image: PILImage.Image, height: int, width: int, channels: int = 3, ) -> NDArray[np.float32]: """Resize and normalise a PIL image to an ``(1, C, H, W)`` float32 NCHW tensor. Resizes with ``RFDETR.predict()``'s exact convention — bilinear, half-pixel centers, ``antialias=False`` — via ``torchvision`` when importable (bit-exact parity) or the pure-NumPy ``_bilinear_resize_half_pixel`` otherwise (float32 op-order noise only). PIL resize is not used: both its BILINEAR and BICUBIC filters apply adaptive antialiasing when downscaling and diverge from predict(), shifting confidence scores. Normalises with ImageNet statistics: ``mean=[0.485, 0.456, 0.406]``, ``std=[0.229, 0.224, 0.225]``. Args: image: Input PIL image; any mode — converted to ``"RGB"`` (3-channel) or ``"L"`` (1-channel) internally. height: Target spatial height expected by the model. width: Target spatial width expected by the model. channels: Number of channels the model expects (``1`` for grayscale, ``3`` for RGB). Returns: Float32 ndarray of shape ``(1, channels, height, width)``. Examples: .. code-block:: python inp = _preprocess_pil_to_nchw(image, height=640, width=640) """ _imagenet_mean = [0.485, 0.456, 0.406] _imagenet_std = [0.229, 0.224, 0.225] pil_mode = "L" if channels == 1 else "RGB" pil_img = image.convert(pil_mode) mean_list = [_imagenet_mean[i % 3] for i in range(channels)] std_list = [_imagenet_std[i % 3] for i in range(channels)] with contextlib.suppress(ImportError): # Match predict() exactly: torchvision to_tensor -> resize(antialias=False) -> normalize. # antialias=False mirrors detr.py's predict(); torchvision's float-tensor default is True. import torch import torchvision.transforms.functional as _F # noqa: N812 with torch.no_grad(): t = _F.to_tensor(pil_img) t = _F.resize(t, [height, width], antialias=False) t = _F.normalize(t, mean_list, std_list) return np.asarray(t.unsqueeze(0).cpu().numpy(), dtype=np.float32) # Torch-free fallback: same antialias-free half-pixel bilinear as predict(), in NumPy. arr = np.asarray(pil_img, dtype=np.float32) / 255.0 if arr.ndim == 2: # "L" → (H, W); needs (H, W, 1) arr = arr[:, :, np.newaxis] chw = _bilinear_resize_half_pixel(arr.transpose(2, 0, 1), height, width) mean = np.array(mean_list, dtype=np.float32)[:, np.newaxis, np.newaxis] std = np.array(std_list, dtype=np.float32)[:, np.newaxis, np.newaxis] chw = (chw - mean) / std return np.expand_dims(chw, axis=0).astype(np.float32) # (1, C, H, W) def _run_inference( session: Any, image_path: str | Path, threshold: float = 0.3, num_select: int | None = None, ) -> tuple[Detections, PILImage.Image]: """Preprocess one image, run ONNX Runtime inference, and decode detections. Reads input shape from the session (NCHW ``float32``), resizes and normalises the image with ImageNet statistics, invokes the model, then decodes the ``dets`` / ``labels`` output tensors into a :class:`supervision.Detections` object with pixel-space ``xyxy`` boxes. **Input contract** (must match ``RFDETR.predict()`` preprocessing exactly): - Image is opened with Pillow and converted to ``"RGB"`` (3-channel) or ``"L"`` (1-channel greyscale) depending on the model's channel count. PIL is used only to decode and convert — never to resize. - Resize follows ``RFDETR.predict()``'s exact convention — bilinear, half-pixel centers, ``antialias=False`` — applied by :func:`_preprocess_pil_to_nchw` via ``torchvision`` when importable (bit-exact) or the pure-NumPy ``_bilinear_resize_half_pixel`` fallback. PIL's own BILINEAR/BICUBIC filters apply adaptive antialiasing when downscaling and would shift pixel values away from predict(), degrading confidence. - Pixel values are scaled to ``[0, 1]`` then normalised with ImageNet statistics: ``mean=[0.485, 0.456, 0.406]``, ``std=[0.229, 0.224, 0.225]``. - The tensor is kept as ``[1, C, H, W]`` (NCHW) — unlike the TFLite helper which uses NHWC because ``onnx2tf`` transposes at export time. ONNX RT consumes the native ONNX NCHW layout directly. Args: session: ONNX Runtime ``InferenceSession`` returned by ``_create_onnx_session``. image_path: Path to the input image (any format supported by Pillow). RGB images are used as-is; RGBA / palette images are converted. threshold: Confidence threshold; detections below this are discarded. num_select: Maximum query/class pairs selected before thresholding. ``None`` uses the exported model's query count, matching shipped RF-DETR configurations; pass an explicit value for custom exports. Returns: A tuple of ``(detections, pil_img)`` where ``detections`` contains pixel-space ``xyxy`` boxes and ``pil_img`` is the original PIL image at its original resolution. Examples: .. code-block:: python sess = _create_onnx_session("model.onnx") dets, img = _run_inference(sess, "photo.jpg", threshold=0.3) print(dets.confidence) """ inputs = session.get_inputs() outputs = session.get_outputs() input_name = inputs[0].name # ONNX NCHW: [batch, channels, height, width] _, channels, height, width = inputs[0].shape with PILImage.open(image_path) as pil_img: inp_tensor = _preprocess_pil_to_nchw(pil_img, height, width, channels) raw_outputs = session.run(None, {input_name: inp_tensor}) # RF-DETR ONNX output names: "dets" = pred_boxes, "labels" = pred_logits. # Match by name so the code is robust to output reordering. output_names = [out.name for out in outputs] boxes_idx = next((i for i, name in enumerate(output_names) if "dets" in name), None) logits_idx = next((i for i, name in enumerate(output_names) if "labels" in name), None) if boxes_idx is None or logits_idx is None: # Fall back to shape-based matching: boxes (*, 4) and logits (*, num_classes+1). logger.warning( "Name-based ONNX output matching failed (available names: %s). Falling back to shape-based matching.", output_names, ) shape_boxes_candidates = [ i for i, arr_out in enumerate(raw_outputs) if arr_out.ndim == 3 and arr_out.shape[-1] == 4 ] shape_logits_candidates = [ i for i, arr_out in enumerate(raw_outputs) if arr_out.ndim == 3 and arr_out.shape[-1] != 4 ] if len(shape_boxes_candidates) == 1 and len(shape_logits_candidates) == 1: boxes_idx = shape_boxes_candidates[0] logits_idx = shape_logits_candidates[0] elif len(raw_outputs) == 2: # Ambiguous shapes (e.g. num_classes==3 → logits dim==4 == boxes dim). # ONNX preserves output order: index 0 = dets (boxes), index 1 = labels (logits). logger.warning( "Shape-based ONNX output matching is ambiguous (both outputs have last dim==4, " "which happens when num_classes==3). Falling back to positional order: " "output 0 = boxes ('dets'), output 1 = logits ('labels'). " "If detections look wrong, inspect output names with _create_onnx_session() " "and set LOG_LEVEL=DEBUG." ) boxes_idx = 0 logits_idx = 1 else: available_shapes = [list(arr_out.shape) for arr_out in raw_outputs] raise ValueError( f"Shape-based ONNX output matching failed. Expected exactly one rank-3 tensor with " f"last dim == 4 (boxes) and one rank-3 tensor with last dim != 4 (logits). " f"Available output shapes: {available_shapes}" ) boxes_cwh = raw_outputs[boxes_idx][0] # (Q, 4) normalised cxcywh # Drop last logit column: RF-DETR adds +1 to num_classes (no-object slot, criterion.py:323). # Keeping it causes class_id == len(class_names) → IndexError at display time. logits = raw_outputs[logits_idx][0, :, :-1] # (Q, num_classes) # RF-DETR uses per-class sigmoid (not softmax) — mirrors PostProcess.forward in postprocess.py. if logits.size: logger.debug( "Logits stats: shape=%s min=%.3f max=%.3f mean=%.3f", logits.shape, float(logits.min()), float(logits.max()), float(logits.mean()), ) else: logger.debug("Logits stats: empty shape=%s", logits.shape) one = np.asarray(1, dtype=logits.dtype) scores_all = one / (one + np.exp(-logits.clip(-88, 88))) # Flatten (Q, C) to Q*C query/class pairs and take the top-scoring ones before thresholding — # mirrors PostProcess._select_topk. A per-query argmax (the previous approach) keeps at most # one class per query, silently dropping legitimate detections whenever a query scores above # threshold on more than one class; see _topk.py for why that happens routinely here. selection_cap = boxes_cwh.shape[0] if num_select is None else num_select scores, cls, query_idx = _select_topk_multiclass(scores_all, threshold, num_select=selection_cap) if scores_all.size: logger.debug( "Scores stats: min=%.3f max=%.3f — detections above threshold %.2f: %d", float(scores_all.min()), float(scores_all.max()), threshold, int(scores.shape[0]), ) else: logger.debug("Scores stats: empty — detections above threshold %.2f: %d", threshold, int(scores.shape[0])) cx, cy, bw, bh = boxes_cwh[query_idx].T ow, oh = pil_img.size xyxy = np.stack([cx - bw / 2, cy - bh / 2, cx + bw / 2, cy + bh / 2], axis=1) xyxy *= np.array([ow, oh, ow, oh], dtype=np.float32) return Detections(xyxy=xyxy, confidence=scores, class_id=cls.astype(int)), pil_img # Benchmarking helper — not part of production inference API; subject to removal. def _onnx_runtime( onnx_path: Path | str, image: PILImage.Image, providers: list[str], warmup: int = 20, runs: int = 100, ) -> tuple[float, float, str]: """Benchmark ONNX Runtime inference for one image and provider list. Creates a fresh ``InferenceSession`` with the requested providers, preprocesses ``image`` once using ImageNet normalisation, then runs timed inference with ``time.perf_counter``. GPU timings may underestimate real latency if the CUDA execution provider is configured for asynchronous execution; for accurate GPU timing use CUDA events. Args: onnx_path: Path to the ``.onnx`` model file. image: Input image (any size); resized to the model's expected spatial resolution. providers: Ordered list of ORT execution providers, e.g. ``["CUDAExecutionProvider", "CPUExecutionProvider"]``. warmup: Number of un-timed warm-up runs before measurement begins. runs: Number of timed runs used to compute statistics. Returns: A ``(mean_ms, std_ms, provider_label)`` tuple where ``provider_label`` is the first active provider with ``"ExecutionProvider"`` stripped, e.g. ``"CUDA"`` or ``"CPU"``. Examples: .. code-block:: python mean_ms, std_ms, label = _onnx_runtime("model.onnx", image, ["CPUExecutionProvider"]) print(f"{label}: {mean_ms:.1f} ms ± {std_ms:.1f}") """ import time sess = _create_onnx_session(onnx_path, providers=providers) active = sess.get_providers()[0] if active != providers[0]: raise RuntimeError( f"Requested provider {providers[0]!r} not active — ORT fell back to {active!r}. " "Install onnxruntime-gpu: `pip install onnxruntime-gpu`" ) input_meta = sess.get_inputs()[0] _, channels, height, width = input_meta.shape inp = _preprocess_pil_to_nchw(image, height, width, channels) feed = {input_meta.name: inp} for _ in range(warmup): sess.run(None, feed) timings: list[float] = [] for _ in range(runs): t0 = time.perf_counter() sess.run(None, feed) timings.append((time.perf_counter() - t0) * 1000.0) arr_t = np.array(timings) provider_label = sess.get_providers()[0].replace("ExecutionProvider", "") return float(arr_t.mean()), float(arr_t.std()), provider_label