RF-DETR with ONNX
This repository contains code to load an ONNX version of RF-DETR and perform inference, including drawing the results on images. It demonstrates how to convert a PyTorch model to ONNX format and inference with minimal dependencies.
RF-DETR is a transformer-based object detection and instance segmentation architecture developed by Roboflow. For more details on the model, please refer to the impressive work by the Roboflow team here.
| Roboflow |
ONNX Runtime Inference (Object detection) |
ONNX Runtime Inference (Instance segmentation) |
|
|
|
Installation
First, clone the repository:
git clone --depth 1 https://github.com/PierreMarieCurie/rf-detr-onnx.git
Then, install the required dependencies.
Using uv (recommanded)
If not installed, just run (on macOS and Linux):
curl -LsSf https://astral.sh/uv/install.sh | sh
Check Astral documentation if you need alternative installation methods.
Then:
uv sync --extra export-tools
If you only want to use the inference scripts without converting your own model, you don’t need the rfdetr dependencies, so just run:
uv sync
Not using uv (not recommanded)
pip install --upgrade .
Make sure to install Python 3.10+ on your local or virtual environment.
Downloading from Hugging-face
Roboflow provides pre-trained RF-DETR models on the COCO and Objects365 datasets. We have already converted some of these models to the ONNX format for you, which you can directly download from Hugging Face.
Note that this corresponds to rf-detr version 1.4.1:
- Object detection:
- Trained on COCO dataset:
- Trained on Objects365 dataset:
- Instance segmentation (train on COCO dataset)
Converting
If you want to export your own fine-tuned RF-DETR model, we provide a script to help you do it:
uv run export.py --checkpoint path/to/your/file.pth
You don’t need to specify the architecture (Nano, Small, Medium, Base, Large), it is detected automatically.
Additionnal conversion parameters
uv run export.py -h
Use the --model-name argument to specify the output ONNX file, and add the --no-simplify flag if you want to skip simplification.
Inference Script Example
Below is an example showing how to perform inference on a single image:
from rfdetr_onnx import RFDETR_ONNX
# Get model and image
image_path = "https://media.roboflow.com/notebooks/examples/dog-2.jpeg"
model_path = "rf-detr-base.onnx"
# Initialize the model
model = RFDETR_ONNX(model_path)
# Run inference and get detections
_, labels, boxes, masks = model.predict(image_path)
# Draw and display the detections
model.save_detections(image_path, boxes, labels, masks, "output.jpg")
Alternatively, we provide a script to help you do it:
uv run inference.py --model path/to/your/model.onnx --image path/to/your/image
Additionnal inference parameters
uv run inference.py -h
Use the --threshold argument to specify the confidence threshold and the --max_number_boxes argument to limit the maximum number of bounding boxes. Also, add --output option to specify the output file name and extension if needed (default: output.jpg)
Version Compatibility
| Repo Tag |
rfdetr Version |
Status |
| v1.0-rfdetr1.3.0 |
1.3.0 |
Stable |
| main |
1.4.1 |
In progress |
License
This repository is licensed under the MIT License. See license file for more details.
However, some parts of the code are derived from Roboflow software. Below are the details:
- Apache License 2.0 (reference): RF-DETR models and pretrained weights (except
rfdetr-xlarge and rfdetr-2xlarge) and all rfdetr Python package.
- Platform Model License 1.0 (PML-1.0) (reference):
rfdetr-xlarge and rfdetr-2xlarge models and pretrained weights.
More information about Roboflow model licensing here.
Acknowledgements
- Thanks to the Roboflow team and everyone involved in the development of RF-DETR, particularly for sharing a state-of-the-art model under a permissive free software license.
评论