RF-DETR Drone Detection Training

Fine-tuning RF-DETR (Real-time Fast DEtection TRansformer) for drone detection as part of the AIRHOUND UAV perception pipeline.

Status: not actively maintained. Built for the AIRHOUND UAV project and left up as a reference training pipeline. Issues and pull requests may not get a response.

Overview

This repository trains an RF-DETR model on a drone detection dataset and exports it for deployment on NVIDIA Jetson Orin. The trained weights integrate with the main AIRHOUND perception system.

Hardware targets:

• Training: RTX 4070 Laptop (12GB VRAM)
• Deployment: NVIDIA Jetson Orin 16GB

Dataset: Roboflow Drone Detection

• 13,869 training images
• 1,983 validation images
• 1 class: drone

Installation

Prerequisites

• Python 3.11+
• CUDA 11.8+ (for GPU training)
• pip or uv package manager

Install from source

# Clone the repository
git clone https://github.com/rylanmalarchick/rf-detr-training.git
cd rf-detr-training

# Install in development mode
pip install -e ".[dev]"

# Or using uv (faster)
uv pip install -e ".[dev]"

Verify installation

# Check CLI is available
rf-detr-train --help
rf-detr-export --help

# Or run as module
python -m rf_detr_drone --help

Usage

Quick Start

# Run a quick training test (5 epochs)
python scripts/train.py --epochs 5 --batch-size 4 --no-wandb

# Full training run
python scripts/train.py --epochs 50 --batch-size 8 --wandb-project rf-detr-drone

# Export to ONNX
python scripts/export.py weights/drone_rfdetr_best.pt --format onnx

Training Options

python scripts/train.py \
    --data-dir data \
    --epochs 50 \
    --batch-size 8 \
    --lr 1e-4 \
    --device cuda \
    --output-dir weights \
    --tensorboard \
    --wandb \
    --wandb-project rf-detr-drone

Key arguments:

Argument	Default	Description
--data-dir	data	Path to dataset directory
--epochs	50	Number of training epochs
--batch-size	8	Training batch size (8-16 for RTX 4070)
--lr	1e-4	Learning rate
--device	auto	Device (cuda, cpu, or auto)
--output-dir	weights	Directory to save model weights
--tensorboard	True	Enable TensorBoard logging
--wandb	True	Enable Weights & Biases logging
--no-wandb	-	Disable W&B (for testing)

Export Options

# Export to ONNX
python scripts/export.py weights/model.pt --format onnx --fp16

# Export to TensorRT (run on Jetson)
python scripts/export.py weights/model.onnx --format tensorrt

Export arguments:

Argument	Default	Description
--format	onnx	Export format (onnx or tensorrt)
--fp16	True	Use FP16 precision
--opset	17	ONNX opset version
--simplify	True	Simplify ONNX graph
--output-dir	weights	Output directory

Python API

from pathlib import Path
from rf_detr_drone import DroneTrainer, TrainingConfig, DataConfig

# Configure training
training_config = TrainingConfig(
    epochs=50,
    batch_size=8,
    learning_rate=1e-4,
)

data_config = DataConfig(
    data_dir=Path("data"),
)

# Train
trainer = DroneTrainer(
    training_config=training_config,
    data_config=data_config,
)
result = trainer.train()

print(f"Training completed in {result.training_time_seconds:.0f}s")
print(f"Best weights saved to: {result.best_weights_path}")

Project Structure

rf-detr-training/
├── src/rf_detr_drone/     # Main package
│   ├── __init__.py        # Public API
│   ├── config.py          # Frozen dataclass configs
│   ├── train.py           # DroneTrainer class
│   ├── export.py          # ONNX/TensorRT export
│   └── cli.py             # CLI entrypoints
├── scripts/
│   ├── train.py           # Standalone training script
│   └── export.py          # Standalone export script
├── tests/
│   ├── conftest.py        # Pytest fixtures
│   └── test_config.py     # Config tests
├── data/                  # Dataset (gitignored)
│   ├── train/images/      # Training images
│   ├── train/labels/      # Training labels (YOLO format)
│   ├── valid/images/      # Validation images
│   ├── valid/labels/      # Validation labels
│   └── data.yaml          # Dataset configuration
├── weights/               # Saved models (gitignored)
├── runs/                  # TensorBoard logs (gitignored)
├── pyproject.toml         # Project configuration
└── README.md              # This file

Deployment to Jetson

1. Train Model (Any GPU; below done on RTX 4070 Mobile):

   python scripts/train.py --epochs 50 --batch-size 8
   

2. Export to ONNX:

   python scripts/export.py weights/drone_rfdetr_best.pt --format onnx --fp16
   

3. Copy to Jetson:

   scp weights/drone_rfdetr_best.onnx jetson@<ip>:/path/to/airhound/weights/
   

4. Convert to TensorRT on Jetson:

   # On Jetson
   python scripts/export.py weights/drone_rfdetr_best.onnx --format tensorrt
   

   Or using trtexec directly:

   /usr/src/tensorrt/bin/trtexec \
       --onnx=weights/drone_rfdetr_best.onnx \
       --saveEngine=weights/drone_rfdetr_best.engine \
       --fp16 \
       --workspace=4096
   

Monitoring

TensorBoard

tensorboard --logdir runs/tensorboard

Weights & Biases

Training logs are automatically uploaded to W&B. View at: https://wandb.ai//rf-detr-drone

Development

Run tests

pytest tests/ -v

Format code

black src/ tests/ scripts/
isort src/ tests/ scripts/

Lint

ruff check src/ tests/ scripts/
mypy src/

References

• RF-DETR Paper
• RF-DETR GitHub
• Roboflow RF-DETR Docs
• AIRHOUND Perception

License

MIT License - See LICENSE for details.