diff --git a/docs/learn/export.md b/docs/learn/export.md index d05355ac0..697565bba 100644 --- a/docs/learn/export.md +++ b/docs/learn/export.md @@ -1,5 +1,5 @@ --- -description: Export RF-DETR models to ONNX, TensorRT, TFLite, ExecuTorch, and native CoreML (FP32/FP16/INT8) for high-performance inference on GPUs, mobile, and edge devices. +description: Export RF-DETR models to ONNX, TensorRT, TFLite, ExecuTorch, native CoreML and OpenVINO IR (FP32/FP16/INT8) for high-performance inference on GPUs, mobile, and edge devices. --- # Export RF-DETR Model @@ -7,6 +7,7 @@ description: Export RF-DETR models to ONNX, TensorRT, TFLite, ExecuTorch, and na !!! tip "Key Takeaways" - Export to ONNX for cross-platform inference with ONNX Runtime, OpenVINO, or TensorRT + - Export to OpenVINO IR for optimized inference on CPU (x86, ARM), GPU (Intel integrated & discrete GPU) and AI accelerators (Intel NPU) - Export to TFLite (FP32, FP16, INT8) for mobile and edge deployment - TensorRT conversion delivers lowest latency on NVIDIA GPUs (2.3 ms for Nano) - INT8 quantization requires calibration data from your dataset for accurate results @@ -14,7 +15,7 @@ description: Export RF-DETR models to ONNX, TensorRT, TFLite, ExecuTorch, and na - Export to ExecuTorch for on-device PyTorch inference (XNNPACK, CoreML, QNN) - Export directly to native CoreML (`.mlpackage`) for Xcode / Apple-platform deployment — see [Native CoreML Export](#native-coreml-export-mlpackage) -RF-DETR supports exporting models to ONNX, TFLite, ExecuTorch, and native CoreML formats, enabling deployment across a wide range of inference frameworks, edge devices, and hardware accelerators. +RF-DETR supports exporting models to ONNX, TFLite, ExecuTorch, native CoreML and OpenVINO IR formats, enabling deployment across a wide range of inference frameworks, edge devices, and hardware accelerators. ## Installation @@ -24,6 +25,9 @@ Install the export dependencies you need: # ONNX export only pip install "rfdetr[onnx]" +# OpenVINO IR export +pip install "rfdetr[openvino]" + # TFLite export pip install "rfdetr[tflite]" @@ -67,7 +71,7 @@ The `export()` method accepts several parameters to customize the export process | Parameter | Default | Description | | ------------------ | ---------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | | `output_dir` | `"output"` | Directory where the exported model will be saved. | -| `format` | `"onnx"` | Export format: `"onnx"`, `"tflite"`, `"tensorrt"` (alias: `"trt"`), `"executorch"`, or `"coreml"`. | +| `format` | `"onnx"` | Export format: `"onnx"`, `"tflite"`, `"tensorrt"` (alias: `"trt"`), `"executorch"`, `"openvino"` or `"coreml"`. | | `quantization` | `None` | TFLite quantization mode: `None`/`"fp32"`, `"fp16"`, or `"int8"`. Only used when `format="tflite"`. | | `calibration_data` | `None` | Calibration data for TFLite export. Image directory, `.npy` file path, NumPy array, or `None`. See [TFLite Export](#tflite-export). | | `max_images` | `100` | Maximum number of images to load from a calibration directory for TFLite INT8 quantization. Ignored for other calibration data formats. | @@ -464,6 +468,104 @@ boxes = interpreter.get_tensor(boxes_detail["index"]) labels = interpreter.get_tensor(labels_detail["index"]) ``` +## OpenVINO IR Export + +OpenVINO IR (Intermediate Representation) is a proprietary model format used by the OpenVINO Toolkit to optimize and deploy deep learning models. + +### Prerequisites + +```bash +pip install "rfdetr[openvino]" +``` + +### Basic OpenVINO Export + +=== "Object Detection" + + ```python + from rfdetr import RFDETRMedium + + model = RFDETRMedium(pretrain_weights="") + + model.export(format="openvino", output_dir="output") + ``` + +=== "Image Segmentation" + + ```python + from rfdetr import RFDETRSegMedium + + model = RFDETRSegMedium(pretrain_weights="") + + model.export(format="openvino", output_dir="output") + ``` + +This produces two files (named after the model's variant): + +- `output/.xml` - The model structure (Intermediate Representation) +- `output/.bin` - The model weights + +### OpenVINO Export with Custom Resolution + +```python +from rfdetr import RFDETRMedium + +model = RFDETRMedium(pretrain_weights="") + +model.export(format="openvino", shape=(608, 608)) +``` + +### OpenVINO Inference Example + +```python +import numpy as np +from PIL import Image +from rfdetr.export._openvino.inference import OpenVINOInference + +# Load the exported model +model = OpenVINOInference("output/rfdetr-medium.xml") + +# Prepare input image (NCHW format, ImageNet normalized) +image = Image.open("image.jpg").convert("RGB").resize((576, 576)) +image_array = np.array(image).astype(np.float32) / 255.0 + +# Apply ImageNet normalization +mean = np.array([0.485, 0.456, 0.406]) +std = np.array([0.229, 0.224, 0.225]) +image_array = (image_array - mean) / std + +# Convert to NCHW format +image_array = np.transpose(image_array, (2, 0, 1)) +image_array = np.expand_dims(image_array, axis=0) + +# Run inference +outputs = model(image_array) +boxes, labels = outputs +``` + +### Benchmark OpenVINO Model + +Use OpenVINO's `benchmark_app` tool to measure performance: + +```bash +benchmark_app -m output/rfdetr-medium.xml -data_shape [1,3,576,576] +``` + +### OpenVINO Model Outputs + +The exported OpenVINO IR model produces the following outputs: + +- **Object Detection Models**: + + - Output 0: Bounding boxes `[batch, 300, 4]` (x, y, w, h in normalized coordinates) + - Output 1: Class logits `[batch, 300, num_classes]` + +- **Segmentation Models**: + + - Output 0: Bounding boxes `[batch, 300, 4]` + - Output 1: Class logits `[batch, 300, num_classes]` + - Output 2: Instance masks (if segmentation head is present) + ## ExecuTorch Export !!! warning "Experimental — Use with Caution" diff --git a/pyproject.toml b/pyproject.toml index ff97e3155..f01d01748 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -138,6 +138,9 @@ coreml = [ # cause is understood/fixed upstream; don't treat it as a guaranteed-deterministic fix. "torch<2.12", ] +openvino = [ + "openvino>=2024.0", # Direct PyTorch -> OpenVINO IR conversion +] loggers = [ "tensorboard>=2.13.0", "protobuf>=3.20.0", # Cap <4.0.0 removed — tensorflow>=2.16.0 (required by [tflite]) needs protobuf>=3.20.3 and is incompatible with <4.0.0 (see #1041) diff --git a/src/rfdetr/detr.py b/src/rfdetr/detr.py index 71de03648..60aff9d2f 100644 --- a/src/rfdetr/detr.py +++ b/src/rfdetr/detr.py @@ -1569,125 +1569,129 @@ def export( ) -> Path: """Export the trained model to ONNX, TFLite, TensorRT, ExecuTorch, or CoreML format. - See the `export documentation `_ for more information. - - Args: - output_dir: Directory to write the exported model to. - infer_dir: Optional directory of sample images for dynamic-axes inference. - backbone_only: Export only the backbone (feature extractor). - opset_version: ONNX opset version to target. - verbose: Print export progress information. - shape: ``(height, width)`` tuple; defaults to square at model resolution. - Both dimensions must be divisible by ``patch_size * num_windows``. - batch_size: Static batch size to bake into the ONNX graph. - dynamic_batch: If True, export with a dynamic batch dimension so the model accepts variable batch sizes - at runtime (spatial dimensions always stay fixed). Applies to the ONNX and TFLite graphs. Not - supported for ExecuTorch export on executorch 1.3.1 (raises ``NotImplementedError``): the runtime - cannot resize RF-DETR's windowed-attention reshapes, so a dynamic ``.pte`` runs only at the traced - batch — export one ``.pte`` per batch size instead. Also unsupported for native CoreML - (``format="coreml"``): fixed shapes are required for reliable ANE / GPU scheduling. - patch_size: Backbone patch size. Defaults to the value stored in - ``model_config.patch_size`` (typically 14 or 16). When provided explicitly it must match the - instantiated model's patch size. Shape divisibility is validated against ``patch_size * num_windows``. - format: Export format — ``"onnx"`` (default), ``"tflite"``, ``"tensorrt"`` (alias: ``"trt"``), - ``"executorch"`` (alias: ``"pte"``), or ``"coreml"``. - ``"tflite"`` and ``"tensorrt"`` both first export to ONNX, then convert: ``"tflite"`` via - ``onnx2tf`` (requires ``pip install rfdetr[tflite]``); ``"tensorrt"`` via the TensorRT - Python API (requires ``pip install rfdetr[tensorrt]``). Unlike ``"onnx"``/ - ``"tflite"`` portable serialization, ``"tensorrt"`` performs target-specific compilation at export - time and produces a non-portable ``.trt`` engine tied to the build machine's GPU and TensorRT version. - When ``"executorch"`` is selected the model is exported directly via ``torch.export`` to an ExecuTorch - ``.pte`` file (no ONNX step), configured by *backend* / *soc* below. Requires - ``pip install rfdetr[executorch]``. - When ``"coreml"`` is selected the model is exported via ``torch.export`` + ``coremltools`` to a - native ``.mlpackage`` (no ONNX step; requires ``pip install rfdetr[coreml]``). This is distinct from - ExecuTorch's ``format="executorch", backend="coreml"`` path, which still produces a ``.pte``. If - you know that ExecuTorch delegate and expect ``format="coreml"`` to mean the same thing: it does - not — pass ``format="executorch", backend="coreml"`` for the ``.pte`` route instead. Passing both - ``format="coreml"`` and ``backend="coreml"`` together does **not** fall through to the ExecuTorch - delegate; ``backend`` is ignored (with a warning) and the native ``.mlpackage`` path always runs. - Keypoint models are untested with ``format="coreml"`` — detection and segmentation have - registry-clean and numerical-parity test coverage (see - ``tests/export/test_coreml_op_coverage.py`` / ``test_coreml_export.py``), keypoint models - currently do not. - - .. warning:: - TFLite, ExecuTorch, and CoreML export are experimental and subject to change; upstream dependency - instabilities (``onnx2tf``, ``ai_edge_litert``, ``executorch``, ``coremltools``) may affect results. - quantization: TFLite quantization mode (ignored when - ``format="onnx"``). One of ``None``, ``"fp32"``, ``"fp16"``, ``"int8"``. ``None`` / ``"fp32"`` / - ``"fp16"`` produce FP32 + FP16 ``.tflite`` files; ``"int8"`` additionally produces an INT8-quantized - model. - calibration_data: Representative images for INT8 calibration and ``onnx2tf`` output validation. Accepts: - - * ``None`` — auto-generate random data (sufficient for fp32/fp16; warns for int8). - * A **directory path** (``str``) containing JPEG/PNG - images — the converter automatically loads, resizes, and prepares them. This is the simplest - approach. - * A path (``str``) to a ``.npy`` file of shape ``(N, H, W, 3)``, dtype float32, values in ``[0, 1]``. - * A :class:`numpy.ndarray` with the same format. - - For INT8 quantization, provide 20–100 representative images from your training/validation set for best - accuracy. - max_images: Maximum number of images to load from a calibration directory. Defaults to ``100``. Only used - when *calibration_data* is a directory path. - backend: Hardware backend to specialize the export for. Required when ``format="executorch"`` and - ignored — with a warning — for any other format. Accepted values for ExecuTorch: - ``"xnnpack"`` (portable CPU, fp32), ``"coreml"`` (Apple devices, fp16; requires ``coremltools``), - and ``"qnn"`` (Qualcomm Snapdragon HTP, fp16; requires an ExecuTorch source build against the - QAIRT SDK — not available via pip). - soc: Target SoC (System on Chip) — the specific Qualcomm Snapdragon chip the exported model will run - on. Required when ``backend="qnn"``: the QNN backend compiles the ``.pte`` ahead-of-time for one - chip's Hexagon Tensor Processor (HTP), unlike ``"xnnpack"``/``"coreml"`` which run on any device of - their platform, so the target chip must be known at export time. Ignored — with a warning — for - any other backend or format. Must be a - :class:`~executorch.backends.qualcomm.serialization.qc_schema.QcomChipset` name, e.g. ``"SM8650"`` - (Snapdragon 8 Gen 3); see that enum for the full list of supported chips. Has no effect for - ``"xnnpack"`` or ``"coreml"``. - fp16: Build the TensorRT engine with FP16 precision. Only applies when ``format="tensorrt"`` - (alias ``"trt"``); ignored for every other format. Defaults to ``True`` for lowest latency - on NVIDIA GPUs. Pass ``False`` to build an FP32 engine — required on TensorRT builds that do - not expose the FP16 builder flag (``export()`` otherwise aborts while configuring FP16). - notes: Optional user-defined metadata (string, dict, list, or - any JSON-serialisable value) to embed in the exported ONNX model under the ``"rfdetr_notes"`` metadata - property. When ``None`` no metadata entry is written. String values are stored verbatim; all other - types are JSON-encoded so consumers must call ``json.loads()`` to recover a dict or list. The same - value can be passed to :meth:`train` so the checkpoint and the ONNX file share the same provenance - information. **Ignored for ``format="executorch"`` and ``format="coreml"``**: those artifacts have - no ONNX-style metadata slot, and a non-``None`` value emits a ``UserWarning`` instead of being - embedded. - coreml_precision: ``ct.convert`` compute precision for ``format="coreml"`` — ``None`` (default) or - ``"float32"`` selects FP32 (tight CPU parity with eager PyTorch); ``"float16"`` selects a smaller - ANE-oriented bundle (expect larger numeric drift). Ignored for every other format. - output_name: Full filename override (without extension), e.g. ``"my-model"``. When set, takes - precedence over the model's variant name (``self.size``) and the exported file is named - ``{output_name}.{ext}`` verbatim — this also suppresses the ``_fp32``/``_fp16``/``_{backend}`` - detail suffix that would otherwise be appended to encode the resolved precision/backend/SoC - (see *format* / *coreml_precision* / *backend* / *soc* / *fp16* above). Sanitized against path - traversal (only the basename, extension stripped, is used). Exception: ``format="tflite"`` - always writes multiple files (one per precision/quantization mode), so the ``_fp32``/``_fp16``/ - ``_dynamic_range_quant`` suffix is unavoidable even with *output_name* set — it becomes the stem - instead of the model's variant name. - Exceptions: ``format="onnx"`` with ``backbone_only=True`` appends ``-backbone`` to the filename - (``{output_name}-backbone.onnx``); ``format="tflite"`` writes separate per-precision files instead - of a single ``{output_name}.tflite`` file. The TFLite filenames may include a ``_gs_patched`` infix - before the precision suffix when GridSample ops are patched, e.g. - ``{output_name}_gs_patched_fp32.tflite``; this is the standard RF-DETR path. - - Returns: - Path to the exported model file (``.onnx``, ``.tflite``, ``.trt``, ``.pte``, or ``.mlpackage``). - - Raises: - ValueError: If ``format`` is unrecognized; if ``format="executorch"`` and ``backend`` is missing, - unrecognized, or (for ``backend="qnn"``) ``soc`` is missing; or if the resolved export shape is - not divisible by ``patch_size * num_windows``. - NotImplementedError: If ``dynamic_batch=True`` is combined with ``format="executorch"`` or - ``format="coreml"`` — those paths require a fixed batch size. - ImportError: If the optional dependencies for the requested ``format``/``backend`` are not installed - (e.g. ``rfdetr[onnx]``, ``rfdetr[executorch]``, ``rfdetr[coreml]``, ``coremltools`` for ExecuTorch - ``backend="coreml"``, or an ExecuTorch source build against the QAIRT SDK for ``backend="qnn"``). - RuntimeError: If called after the model has undergone in-place inference optimization (the original - model has been cleared; instantiate a new :class:`RFDETR` to export). + See the `export documentation `_ for more information. + + Args: + output_dir: Directory to write the exported model to. + infer_dir: Optional directory of sample images for dynamic-axes inference. + backbone_only: Export only the backbone (feature extractor). + opset_version: ONNX opset version to target. + verbose: Print export progress information. + shape: ``(height, width)`` tuple; defaults to square at model resolution. + Both dimensions must be divisible by ``patch_size * num_windows``. + batch_size: Static batch size to bake into the ONNX graph. + dynamic_batch: If True, export with a dynamic batch dimension so the model accepts variable batch sizes + at runtime (spatial dimensions always stay fixed). Applies to the ONNX and TFLite graphs. Not + supported for ExecuTorch export on executorch 1.3.1 (raises ``NotImplementedError``): the runtime + cannot resize RF-DETR's windowed-attention reshapes, so a dynamic ``.pte`` runs only at the traced + batch — export one ``.pte`` per batch size instead. Also unsupported for native CoreML + (``format="coreml"``): fixed shapes are required for reliable ANE / GPU scheduling. + patch_size: Backbone patch size. Defaults to the value stored in + ``model_config.patch_size`` (typically 14 or 16). When provided explicitly it must match the + instantiated model's patch size. Shape divisibility is validated against ``patch_size * num_windows``. + format: Export format — ``"onnx"`` (default), ``"tflite"``, ``"tensorrt"`` (alias: ``"trt"``), + ``"executorch"`` (alias: ``"pte"``), ``"coreml"`` or ``"openvino"``. + ``"tflite"`` and ``"tensorrt"`` both first export to ONNX, then convert: ``"tflite"`` via + ``onnx2tf`` (requires ``pip install rfdetr[tflite]``); ``"tensorrt"`` via the TensorRT + Python API (requires ``pip install rfdetr[tensorrt]``). Unlike ``"onnx"``/ + ``"tflite"`` portable serialization, ``"tensorrt"`` performs target-specific compilation at export + time and produces a non-portable ``.trt`` engine tied to the build machine's GPU and TensorRT version. + When ``"executorch"`` is selected the model is exported directly via ``torch.export`` to an ExecuTorch + ``.pte`` file (no ONNX step), configured by *backend* / *soc* below. Requires + ``pip install rfdetr[executorch]``. ``"openvino"`` converts directly from PyTorch to OpenVINO IR + format (requires ``pip install rfdetr[openvino]``). + ``pip install rfdetr[executorch]``. + When ``"coreml"`` is selected the model is exported via ``torch.export`` + ``coremltools`` to a + native ``.mlpackage`` (no ONNX step; requires ``pip install rfdetr[coreml]``). This is distinct from + ExecuTorch's ``format="executorch", backend="coreml"`` path, which still produces a ``.pte``. If + you know that ExecuTorch delegate and expect ``format="coreml"`` to mean the same thing: it does + not — pass ``format="executorch", backend="coreml"`` for the ``.pte`` route instead. Passing both + ``format="coreml"`` and ``backend="coreml"`` together does **not** fall through to the ExecuTorch + delegate; ``backend`` is ignored (with a warning) and the native ``.mlpackage`` path always runs. + Keypoint models are untested with ``format="coreml"`` — detection and segmentation have + registry-clean and numerical-parity test coverage (see + ``tests/export/test_coreml_op_coverage.py`` / ``test_coreml_export.py``), keypoint models + currently do not. + + .. warning:: + TFLite, ExecuTorch, and CoreML export are experimental and subject to change; upstream dependency + instabilities (``onnx2tf``, ``ai_edge_litert``, ``executorch``, ``coremltools``) may affect results. + quantization: TFLite quantization mode (ignored when + ``format="onnx"``, ``format="openvino"``, or ``format="executorch"``). One of ``None``, ``"fp32"``, ``"fp16"``, ``"int8"``. + ``None`` / ``"fp32"`` / ``"fp16"`` produce FP32 + FP16 ``.tflite`` files; ``"int8"`` additionally + produces an INT8-quantized model. + calibration_data: Representative images for INT8 calibration and ``onnx2tf`` output validation. Accepts: + + * ``None`` — auto-generate random data (sufficient for fp32/fp16; warns for int8). + * A **directory path** (``str``) containing JPEG/PNG + images — the converter automatically loads, resizes, and prepares them. This is the simplest + approach. + * A path (``str``) to a ``.npy`` file of shape ``(N, H, W, 3)``, dtype float32, values in ``[0, 1]``. + * A :class:`numpy.ndarray` with the same format. + + For INT8 quantization, provide 20–100 representative images from your training/validation set for best + accuracy. + max_images: Maximum number of images to load from a calibration directory. Defaults to ``100``. Only used + when *calibration_data* is a directory path. + backend: Hardware backend to specialize the export for. Required when ``format="executorch"`` and + ignored — with a warning — for any other format. Accepted values for ExecuTorch: + ``"xnnpack"`` (portable CPU, fp32), ``"coreml"`` (Apple devices, fp16; requires ``coremltools``), + and ``"qnn"`` (Qualcomm Snapdragon HTP, fp16; requires an ExecuTorch source build against the + QAIRT SDK — not available via pip). + soc: Target SoC (System on Chip) — the specific Qualcomm Snapdragon chip the exported model will run + on. Required when ``backend="qnn"``: the QNN backend compiles the ``.pte`` ahead-of-time for one + chip's Hexagon Tensor Processor (HTP), unlike ``"xnnpack"``/``"coreml"`` which run on any device of + their platform, so the target chip must be known at export time. Ignored — with a warning — for + any other backend or format. Must be a + :class:`~executorch.backends.qualcomm.serialization.qc_schema.QcomChipset` name, e.g. ``"SM8650"`` + (Snapdragon 8 Gen 3); see that enum for the full list of supported chips. Has no effect for + ``"xnnpack"`` or ``"coreml"``. + fp16: Build the TensorRT engine with FP16 precision. Only applies when ``format="tensorrt"`` + (alias ``"trt"``); ignored for every other format. Defaults to ``True`` for lowest latency + on NVIDIA GPUs. Pass ``False`` to build an FP32 engine — required on TensorRT builds that do + not expose the FP16 builder flag (``export()`` otherwise aborts while configuring FP16). + notes: Optional user-defined metadata (string, dict, list, or + any JSON-serialisable value) to embed in the exported ONNX model under the ``"rfdetr_notes"`` metadata + property. When ``None`` no metadata entry is written. String values are stored verbatim; all other + types are JSON-encoded so consumers must call ``json.loads()`` to recover a dict or list. The same + value can be passed to :meth:`train` so the checkpoint and the ONNX file share the same provenance + information. **Ignored for ``format="executorch"`` and ``format="coreml"``**: those artifacts have + no ONNX-style metadata slot, and a non-``None`` value emits a ``UserWarning`` instead of being + embedded. + coreml_precision: ``ct.convert`` compute precision for ``format="coreml"`` — ``None`` (default) or + ``"float32"`` selects FP32 (tight CPU parity with eager PyTorch); ``"float16"`` selects a smaller + ANE-oriented bundle (expect larger numeric drift). Ignored for every other format. + output_name: Full filename override (without extension), e.g. ``"my-model"``. When set, takes + precedence over the model's variant name (``self.size``) and the exported file is named + ``{output_name}.{ext}`` verbatim — this also suppresses the ``_fp32``/``_fp16``/``_{backend}`` + detail suffix that would otherwise be appended to encode the resolved precision/backend/SoC + (see *format* / *coreml_precision* / *backend* / *soc* / *fp16* above). Sanitized against path + traversal (only the basename, extension stripped, is used). Exception: ``format="tflite"`` + always writes multiple files (one per precision/quantization mode), so the ``_fp32``/``_fp16``/ + ``_dynamic_range_quant`` suffix is unavoidable even with *output_name* set — it becomes the stem + instead of the model's variant name. + Exceptions: ``format="onnx"`` with ``backbone_only=True`` appends ``-backbone`` to the filename + (``{output_name}-backbone.onnx``); ``format="tflite"`` writes separate per-precision files instead + of a single ``{output_name}.tflite`` file. The TFLite filenames may include a ``_gs_patched`` infix + before the precision suffix when GridSample ops are patched, e.g. + ``{output_name}_gs_patched_fp32.tflite``; this is the standard RF-DETR path. + + Returns: + + + Path to the exported model file (``.onnx``, ``.tflite``, ``.trt``, ``.pte``, ``.mlpackage`` or ``.xml`` for OpenVINO). + + Raises: + ValueError: If ``format`` is unrecognized; if ``format="executorch"`` and ``backend`` is missing, + unrecognized, or (for ``backend="qnn"``) ``soc`` is missing; or if the resolved export shape is + not divisible by ``patch_size * num_windows``. + NotImplementedError: If ``dynamic_batch=True`` is combined with ``format="executorch"`` or + ``format="coreml"`` — those paths require a fixed batch size. + ImportError: If the optional dependencies for the requested ``format``/``backend`` are not installed + (e.g. ``rfdetr[onnx]``, ``rfdetr[executorch]``, ``rfdetr[coreml]``, ``coremltools`` for ExecuTorch ``backend="coreml"``, ``openvino`` + for OpenVINO export, or an ExecuTorch source build against the QAIRT SDK for ``backend="qnn"``). + RuntimeError: If called after the model has undergone in-place inference optimization (the original + model has been cleared; instantiate a new :class:`RFDETR` to export). """ if format == "trt": # "trt" is an alias for "tensorrt" format = "tensorrt" @@ -1715,14 +1719,24 @@ def export( "ANE / GPU scheduling). Export one .mlpackage per batch size instead." ) logger.info(f"Exporting model to {format} format") - try: - from rfdetr.export.main import export_onnx, make_infer_image - except ImportError: - logger.error( - "It seems some dependencies for ONNX export are missing." - " Please run `pip install rfdetr[onnx]` and try again.", - ) - raise + # OpenVINO and ExecuTorch use direct conversion, others may need ONNX + if format not in ("openvino", "executorch"): + try: + from rfdetr.export.main import export_onnx, make_infer_image + except ImportError: + logger.error( + "It seems some dependencies for ONNX export are missing." + " Please run `pip install rfdetr[onnx]` and try again.", + ) + raise + else: + try: + from rfdetr.export.main import make_infer_image + except ImportError: + logger.error( + "It seems some dependencies for export are missing. Please run `pip install rfdetr` and try again.", + ) + raise device = self.model.device @@ -1818,6 +1832,28 @@ def export( model.cpu() input_tensors = input_tensors.cpu() + if format == "openvino": + try: + from rfdetr.export._openvino.exporter import export_openvino + except ImportError: + logger.error( + "It seems OpenVINO is not installed." + ' Please run `pip install "rfdetr[openvino]"` and try again.', + ) + raise + + output_file = export_openvino( + output_dir=str(output_dir_path), + model=model, + input_tensors=input_tensors, + backbone_only=backbone_only, + verbose=verbose, + variant_name=getattr(self, "size", None), + output_names=output_names, + ) + logger.info(f"Successfully exported OpenVINO model to: {output_file}") + return Path(output_file) + if format == "executorch": from rfdetr.export._backend import _export_executorch_format diff --git a/src/rfdetr/export/_openvino/README.md b/src/rfdetr/export/_openvino/README.md new file mode 100644 index 000000000..7244c8fe3 --- /dev/null +++ b/src/rfdetr/export/_openvino/README.md @@ -0,0 +1,134 @@ +# OpenVINO Export for RF-DETR + +This module provides direct PyTorch to OpenVINO IR export for RF-DETR models, without requiring an intermediate ONNX conversion. + +## Installation + +Install RF-DETR with OpenVINO support: + +```bash +pip install "rfdetr[openvino]" +``` + +Or if you already have RF-DETR installed, just add OpenVINO: + +```bash +pip install openvino +``` + +## Basic Usage + +Export your trained model to OpenVINO IR format: + +```python +from rfdetr import RFDETRMedium + +model = RFDETRMedium(pretrain_weights="") + +# Export to OpenVINO IR +model.export(format="openvino", output_dir="output") +``` + +This will create two files (named after the model's variant): + +- `output/rfdetr-medium.xml` - The OpenVINO IR model +- `output/rfdetr-medium.bin` - The model weights + +## Export Parameters + +The `export()` method with `format="openvino"` accepts these parameters: + +- **output_dir** (str, default: `"output"`): Directory where the exported model will be saved +- **format** (str): Set to `"openvino"` for OpenVINO IR export +- **backbone_only** (bool, default: `False`): Export only the backbone feature extractor +- **verbose** (bool, default: `True`): Print export progress information +- **shape** (tuple, optional): Input shape as `(height, width)`. If not provided, uses model's default resolution +- **batch_size** (int, default: `1`): Static batch size for the exported model + +## Advanced Examples + +### Export with Custom Output Directory + +```python +from rfdetr import RFDETRMedium + +model = RFDETRMedium(pretrain_weights="") +model.export(format="openvino", output_dir="exports/my_model") +``` + +### Export with Custom Resolution + +```python +from rfdetr import RFDETRMedium + +model = RFDETRMedium(pretrain_weights="") +model.export(format="openvino", shape=(608, 608)) +``` + +### Export Backbone Only + +```python +from rfdetr import RFDETRMedium + +model = RFDETRMedium(pretrain_weights="") +model.export(format="openvino", backbone_only=True) +``` + +## Using the Exported Model + +### Python Inference + +```python +import numpy as np +from PIL import Image +from rfdetr.export._openvino.inference import OpenVINOInference + +# Load the exported model (use the .xml file from your export) +model = OpenVINOInference("output/rfdetr-medium.xml") + +# Prepare input image (NCHW format, ImageNet normalized) +image = Image.open("image.jpg").convert("RGB").resize((576, 576)) +image_array = np.array(image).astype(np.float32) / 255.0 + +# Apply ImageNet normalization +mean = np.array([0.485, 0.456, 0.406]) +std = np.array([0.229, 0.224, 0.225]) +image_array = (image_array - mean) / std + +# Convert to NCHW format +image_array = np.transpose(image_array, (2, 0, 1)) +image_array = np.expand_dims(image_array, axis=0) + +# Run inference +outputs = model(image_array) +boxes, labels = outputs +``` + +### OpenVINO benchmark_app + +Test the exported model performance: + +```bash +benchmark_app -m output/rfdetr-medium.xml -data_shape [1,3,576,576] +``` + +## Model Outputs + +The exported OpenVINO model produces the following outputs: + +- **Object Detection Models**: + + - Output 0: Bounding boxes `[batch, 300, 4]` (x, y, w, h in normalized coordinates) + - Output 1: Class logits `[batch, 300, num_classes]` + +- **Segmentation Models**: + + - Output 0: Bounding boxes `[batch, 300, 4]` + - Output 1: Class logits `[batch, 300, num_classes]` + - Output 2: Instance masks (if segmentation head is present) + +- **Keypoint Models**: + + - Output 0: Bounding boxes `[batch, 300, 4]` + - Output 1: Class logits `[batch, 300, num_classes]` + - Output 2: Keypoints (if keypoint head is present) diff --git a/src/rfdetr/export/_openvino/__init__.py b/src/rfdetr/export/_openvino/__init__.py new file mode 100644 index 000000000..e9db55b0d --- /dev/null +++ b/src/rfdetr/export/_openvino/__init__.py @@ -0,0 +1,7 @@ +# ------------------------------------------------------------------------ +# RF-DETR +# Copyright (c) 2025 Roboflow. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ + +"""OpenVINO IR export utilities.""" diff --git a/src/rfdetr/export/_openvino/exporter.py b/src/rfdetr/export/_openvino/exporter.py new file mode 100644 index 000000000..8d996bdfe --- /dev/null +++ b/src/rfdetr/export/_openvino/exporter.py @@ -0,0 +1,149 @@ +# ------------------------------------------------------------------------ +# RF-DETR +# Copyright (c) 2025 Roboflow. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ + +"""Direct PyTorch → OpenVINO IR export.""" + +from __future__ import annotations + +import os + +import torch +from torch import nn + +from rfdetr.utilities.logger import get_logger + +logger = get_logger() + + +def export_openvino( + output_dir: str, + model: nn.Module, + input_tensors: torch.Tensor, + backbone_only: bool = False, + verbose: bool = True, + variant_name: str | None = None, + output_names: list[str] | None = None, +) -> str: + """Export a PyTorch model directly to OpenVINO IR format. + + Args: + output_dir: Directory where the exported model will be saved. + model: PyTorch model to export. + input_tensors: Example input tensor(s) for tracing. + backbone_only: Whether exporting only the backbone. + verbose: Whether to print verbose export information. + variant_name: Optional model variant name (e.g., "nano", "small", "medium"). + output_names: List of output names (e.g., ["dets", "labels", "masks"]). + + Returns: + Path to the exported OpenVINO IR model (.xml file). + + Raises: + ImportError: If OpenVINO is not installed. + RuntimeError: If export fails. + """ + try: + from openvino import convert_model, save_model + except ImportError: + logger.error( + 'OpenVINO is not installed. Please run `pip install "rfdetr[openvino]"` and try again.', + ) + raise + + os.makedirs(output_dir, exist_ok=True) + + # Determine output filename + if variant_name: + # Sanitize against path traversal (e.g. "foo/bar" -> "bar", "/tmp/x" -> "x") + variant_name = os.path.splitext(os.path.basename(variant_name))[0] + export_name = f"{variant_name}-backbone" if backbone_only else variant_name + else: + export_name = "backbone_model" if backbone_only else "inference_model" + + output_xml = os.path.join(output_dir, f"{export_name}.xml") + output_bin = os.path.join(output_dir, f"{export_name}.bin") + + if verbose: + logger.info("Converting PyTorch model to OpenVINO IR...") + logger.info(f"Input shape: {input_tensors.shape}") + + # Ensure model is in eval mode and on CPU + model = model.eval().cpu() + if hasattr(model, "export") and callable(getattr(model, "export")): + model.export() + input_tensors = input_tensors.cpu() + + # Determine output names if not provided + if output_names is None: + output_names = ["dets", "labels"] + + try: + # Create a wrapper to handle dictionary outputs + class ModelWrapper(nn.Module): + """Wrapper to convert dictionary outputs to tuple of tensors.""" + + def __init__(self, model: nn.Module, output_names: list[str]): + super().__init__() + self.model = model + self.output_names = output_names + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, ...]: + output = self.model(x) + + if isinstance(output, tuple): + return output + + # Handle backbone-only case (returns tensor directly) + if not isinstance(output, dict): + return (output,) + + # Extract outputs in the specified order + result = [] + for name in self.output_names: + # Map output names to model's keys + if name == "dets": + result.append(output.get("pred_boxes", output.get("dets"))) + elif name == "labels": + result.append(output.get("pred_logits", output.get("labels"))) + elif name == "masks": + masks = output.get("pred_masks", output.get("masks")) + if isinstance(masks, torch.Tensor): + result.append(masks) + elif name == "keypoints": + keypoints = output.get("pred_keypoints", output.get("keypoints")) + if isinstance(keypoints, torch.Tensor): + result.append(keypoints) + elif name == "features": + result.append(output) + + # Filter out None values + result = [r for r in result if r is not None] + return tuple(result) + + # Wrap the model + wrapped_model = ModelWrapper(model, output_names) + wrapped_model.eval() + + # Convert directly to OpenVINO + with torch.no_grad(): + ov_model = convert_model(wrapped_model, example_input=input_tensors) + + # Save the model + save_model(ov_model, output_xml) + + if verbose: + logger.info(f"✓ OpenVINO IR model saved to {output_xml}") + logger.info(f"✓ Model binary saved to {output_bin}") + + return output_xml + + except Exception as e: + logger.error(f"OpenVINO export failed: {e}") + import traceback + + if verbose: + traceback.print_exc() + raise RuntimeError(f"Failed to export model to OpenVINO IR: {e}") from e diff --git a/src/rfdetr/export/_openvino/inference.py b/src/rfdetr/export/_openvino/inference.py new file mode 100644 index 000000000..b30f58181 --- /dev/null +++ b/src/rfdetr/export/_openvino/inference.py @@ -0,0 +1,89 @@ +# ------------------------------------------------------------------------ +# RF-DETR +# Copyright (c) 2025 Roboflow. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 [see LICENSE for details] +# ------------------------------------------------------------------------ + +"""OpenVINO inference utilities for exported RF-DETR models.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from numpy.typing import NDArray + +from rfdetr.utilities.logger import get_logger + +logger = get_logger() + + +class OpenVINOInference: + """Inference wrapper for OpenVINO IR models. + + Example: + >>> from rfdetr.export._openvino.inference import OpenVINOInference + >>> + >>> model = OpenVINOInference("output/inference_model.xml") + >>> # Prepare input image (NCHW format, ImageNet normalized) + >>> outputs = model.infer(image_array) + >>> boxes, labels = outputs + """ + + def __init__(self, model_path: str | Path): + """Initialize OpenVINO inference session. + + Args: + model_path: Path to the OpenVINO IR model (.xml file). + + Raises: + ImportError: If OpenVINO is not installed. + FileNotFoundError: If the model file doesn't exist. + """ + try: + import openvino as ov + except ImportError: + logger.error( + 'OpenVINO is not installed. Please run `pip install "rfdetr[openvino]"` and try again.', + ) + raise + + model_path = Path(model_path) + if not model_path.exists(): + raise FileNotFoundError(f"Model file not found: {model_path}") + + # Initialize OpenVINO runtime + core = ov.Core() + self.model = core.read_model(model_path) + self.compiled_model = core.compile_model(self.model, "AUTO") + self.infer_request = self.compiled_model.create_infer_request() + + # Get input/output info + self.input_layer = self.compiled_model.input(0) + self.output_layers = [self.compiled_model.output(i) for i in range(len(self.compiled_model.outputs))] + + logger.info(f"Loaded OpenVINO model from {model_path}") + logger.info(f"Input shape: {self.input_layer.partial_shape}") + logger.info(f"Number of outputs: {len(self.output_layers)}") + + def infer(self, input_data: NDArray[Any]) -> tuple[NDArray[Any], ...]: + """Run inference on input data. + + Args: + input_data: Input tensor in NCHW format (batch, channels, height, width). + Should be ImageNet normalized [0.485, 0.456, 0.406] mean, + [0.229, 0.224, 0.225] std. + + Returns: + Tuple of output tensors (typically boxes, labels, and optionally masks/keypoints). + """ + # Run inference + self.infer_request.infer({self.input_layer: input_data}) + + # Get outputs + outputs = tuple(self.infer_request.get_output_tensor(i).data for i in range(len(self.output_layers))) + return outputs + + def __call__(self, input_data: NDArray[Any]) -> tuple[NDArray[Any], ...]: + """Alias for infer() to match typical model calling convention.""" + return self.infer(input_data)