Skip to main content
RFDETR is the base class for all RF-DETR model variants. You do not instantiate it directly; instead, use a concrete variant such as RFDETRMedium or RFDETRSegLarge. Every variant inherits the methods documented here. 
from rfdetr import RFDETRMedium

model = RFDETRMedium()
See Model Variants for the full list of available classes.

​
predict()

Run inference on one or more images and return bounding box (or segmentation mask) predictions. 
detections = model.predict("image.jpg", threshold=0.5)
Batch inference: 
detections_list = model.predict(["image1.jpg", "image2.jpg"], threshold=0.4)
From a PIL Image or NumPy array: 
from PIL import Image
import numpy as np

pil_img = Image.open("image.jpg")
detections = model.predict(pil_img)

np_img = np.array(pil_img)
detections = model.predict(np_img)

​
Parameters

​
images
str | PIL.Image.Image | np.ndarray | torch.Tensor | list
required
A single image or a list of images to run inference on. Accepted types per image:
  • str — file path or HTTP/HTTPS URL
  • PIL.Image.Image — PIL image in RGB channel order
  • np.ndarray — NumPy array in RGB channel order, shape (H, W, 3)
  • torch.Tensor — pre-normalized tensor in [0, 1], shape (C, H, W)
Pass a list to process multiple images in a single batched call.
​
threshold
float
default:"0.5"
Minimum confidence score for a detection to be included in the output. Detections with a score below this value are discarded.
​
shape
tuple[int, int]
Optional (height, width) tuple to override the model’s default inference resolution. Both dimensions must be positive integers divisible by patch_size × num_windows. Defaults to (model.resolution, model.resolution).
​
patch_size
int
Backbone patch size used for shape-divisibility validation. Defaults to the value in model_config.patch_size (typically 14 for large models and 16 for smaller ones). When provided, must match the instantiated model’s patch size exactly.

​
Returns

​
detections
supervision.Detections | list[supervision.Detections]
A single supervision.Detections when one image is passed, or a list when multiple images are passed.

​
train()

Fine-tune the model on your own dataset using the PyTorch Lightning training stack. All keyword arguments are forwarded directly to TrainConfig. You can pass them as plain keyword arguments — you do not need to construct a TrainConfig object yourself. 
model.train(
    dataset_dir="dataset/",
    epochs=50,
    batch_size=8,
    lr=1e-4,
    output_dir="runs/my-experiment",
)
Resume from a checkpoint: 
model.train(
    dataset_dir="dataset/",
    epochs=100,
    resume="output/checkpoint_best_total.pth",
)

​
Parameters

All parameters accepted by TrainConfig are valid here. The most commonly used ones are listed below.
​
dataset_dir
str
required
Path to your dataset directory. Supports COCO (_annotations.coco.json) and YOLO (data.yaml) formats. See dataset formats for details.
​
epochs
int
default:"100"
Number of training epochs.
​
batch_size
int | "auto"
default:"4"
Per-device micro-batch size. Set to "auto" to let RF-DETR probe available memory and choose the largest safe value.
​
lr
float
default:"1e-4"
Base learning rate for the decoder and projection heads.
​
output_dir
str
default:"\"output\""
Directory where checkpoints, logs, and the exported training config are written.
​
resume
str
Path to a checkpoint file to resume training from. When set, PTL restores optimizer state and epoch count automatically.
​
device
str
Torch-style device string (e.g. "cpu", "cuda", "cuda:1", "mps"). Mapped to the appropriate PyTorch Lightning accelerator. When omitted, the best available device is selected automatically.
For the complete list of training parameters, see TrainConfig.

​
export()

Export the trained model to ONNX format. 
model.train(dataset_dir="dataset/", epochs=50)
model.export(output_dir="output/")
Export with a dynamic batch dimension: 
model.export(output_dir="output/", dynamic_batch=True)

​
Parameters

​
output_dir
str
default:"\"output\""
Directory to write the .onnx file to.
​
opset_version
int
default:"17"
ONNX opset version to target.
​
batch_size
int
default:"1"
Static batch size to bake into the ONNX graph.
​
dynamic_batch
boolean
default:"false"
When True, exports with a dynamic batch dimension so the ONNX model accepts variable batch sizes at runtime.
​
shape
tuple[int, int]
(height, width) to override the model’s default export resolution. Both dimensions must be divisible by patch_size × num_windows.
​
backbone_only
boolean
default:"false"
Export only the backbone (feature extractor), omitting the detection head.
​
verbose
boolean
default:"true"
Print export progress to stdout.
ONNX export requires the onnx extras. Install with:
pip install "rfdetr[onnx]"
See the full ONNX export guide for runtime integration examples.

​
deploy_to_roboflow()

Upload a trained model to Roboflow for hosted inference via the Roboflow API and Workflows. 
model.train(dataset_dir="dataset/", epochs=50)
model.deploy_to_roboflow(
    workspace="my-workspace",
    project_id="my-project",
    version=1,
    api_key="YOUR_API_KEY",
)

​
Parameters

​
workspace
str
required
The slug of the Roboflow workspace to deploy to.
​
project_id
str
required
The project ID within the workspace.
​
version
int | str
required
The project version number to deploy the model to.
​
api_key
str
Your Roboflow API key. When omitted, the value of the ROBOFLOW_API_KEY environment variable is used.
Set ROBOFLOW_API_KEY in your environment to avoid passing the key in code:
export ROBOFLOW_API_KEY=your_key_here
See the Roboflow deployment guide for a full walkthrough.

​
from_checkpoint() class method

Load a trained or fine-tuned model from a checkpoint file. The correct model subclass is inferred automatically from the pretrain_weights field stored in the checkpoint. 
from rfdetr import RFDETRMedium

model = RFDETRMedium.from_checkpoint("output/checkpoint_best_total.pth")
detections = model.predict("image.jpg")
You can also call it on the base RFDETR class if you do not know the variant ahead of time: 
from rfdetr.detr import RFDETR

model = RFDETR.from_checkpoint("output/checkpoint_best_total.pth")

​
Parameters

​
path
str | os.PathLike
required
Path to a .pth checkpoint file produced by model.train().
​
**kwargs
any
Additional keyword arguments forwarded to the model constructor. Useful for overriding num_classes or passing device.

​
Returns

​
model
RFDETR
An instance of the appropriate RFDETR subclass loaded from the checkpoint, ready for inference or further training.
from_checkpoint() calls torch.load with weights_only=False, which deserialises arbitrary Python objects. Only load checkpoints from trusted sources.

​
optimize_for_inference()

Optimize the model for faster inference by creating a copy of the model in eval mode with the detection head exported. Optionally compiles the model with torch.jit.trace for maximum speed. 
model = RFDETRMedium()
model.optimize_for_inference(compile=True, batch_size=1)

# Run inference as normal — the optimized model is used automatically
detections = model.predict("image.jpg")

​
Parameters

​
compile
boolean
default:"true"
When True, compiles the optimized model with torch.jit.trace for maximum inference speed. The compiled model is fixed to a specific batch_size and resolution — both must match exactly when calling predict().
​
batch_size
int
default:"1"
Batch size to use when compiling the traced model. Only relevant when compile=True.
​
dtype
torch.dtype
default:"torch.float32"
Floating-point precision for the optimized model.
After calling optimize_for_inference(compile=True, batch_size=N), you must call predict() with exactly N images. Call model.remove_optimized_model() to revert to the standard model.

​
remove_optimized_model()

Remove a previously created optimized inference model and revert to standard inference. 
model.optimize_for_inference()
# ... fast inference ...
model.remove_optimized_model()
# back to standard inference

​
class_names property

Returns the class names associated with the loaded model. 
model = RFDETRMedium()
print(model.class_names)  # ['person', 'bicycle', 'car', ...]
Returns the 80 COCO class names by default. When a fine-tuned checkpoint is loaded (via pretrain_weights or from_checkpoint()), returns the class names embedded in that checkpoint.

​
Module-level from_checkpoint()

A convenience wrapper that delegates to RFDETR.from_checkpoint(). Identical behaviour, slightly shorter import: 
from rfdetr import from_checkpoint

model = from_checkpoint("output/checkpoint_best_total.pth")

​
ModelContext

ModelContext is the lightweight model wrapper returned by RFDETR.get_model(). It holds the underlying nn.Module, the post-processor, the device, and the input resolution. You rarely need to interact with it directly, but it is available if you need low-level access. 
model = RFDETRMedium()
ctx = model.model  # ModelContext instance

print(ctx.device)      # e.g. "cuda"
print(ctx.resolution)  # e.g. 576
print(ctx.class_names) # list of class name strings

​
Attributes

​
model
torch.nn.Module
The underlying LWDETR nn.Module.
​
postprocess
PostProcess
Post-processing module that converts raw logits and boxes to scored detections.
​
device
torch.device
Device the model lives on.
​
resolution
int
Square input resolution in pixels (e.g. 576 for RFDETRMedium).
​
class_names
list[str] | None
Class names loaded from the checkpoint. None when no custom names are embedded.

Model variants

All available detection and segmentation model classes with their default configs.

TrainConfig

Complete reference for every training parameter accepted by model.train().

ONNX export

Export your trained model to ONNX for deployment.

Roboflow deployment

Deploy to Roboflow for hosted inference.

Build docs developers (and LLMs) love

Get started for freeTalk to us