Skip to content

Repository files navigation

Drone Tracker

Hugging Face Spaces Hugging Face Models

A robust, real-time drone detection and tracking system designed to identify and track small Unmanned Aerial Vehicles (UAVs) in various environments. Fine-tuned on the Anti-UAV410 dataset, the system specializes in thermal IR imagery and tiny-object detection, ensuring high accuracy even when drones appear as mere pixels on the screen.

sample

Stack: Python 3.13 · PyTorch 2.11 · Ultralytics 8.4 · Albumentations · ByteTrack
Hardware: Optimized for efficiency, capable of running on standard GPUs (e.g. 6GB VRAM) and edge devices via ONNX/TensorRT.

Windows PowerShell note: Prefix every command with $env:PYTHONUTF8=1; to avoid Unicode encoding errors in the terminal.

Product Overview

Purpose

The proliferation of consumer and commercial drones has introduced new challenges for airspace security, privacy, and infrastructure protection. The Drone Tracker provides an automated, computer-vision-based solution to detect, track, and monitor unauthorized drones in real-time.

Key Features

  • Tiny Object Detection: Utilizing a customized YOLOv11 architecture with an additional high-resolution P2 detection head to catch targets as small as 4x4 pixels.
  • Robust Tracking: Integrated ByteTrack algorithm provides stable tracking across frames, even when the drone is temporarily occluded or moves erratically.
  • Environment Agnostic: Trained with aggressive augmentations (motion blur, fog, rain, noise) and optimized for thermal IR to ensure reliability in challenging lighting and weather conditions.
  • Flexible Deployment: Processes pre-recorded video files, connects to live camera feeds, and exports to ONNX/TensorRT for deployment on low-power edge devices like the Raspberry Pi or NVIDIA Jetson.

sample

⚡ Quick Setup

# Install all dependencies (torch cu128 is locked in pyproject.toml)
uv sync

# Verify CUDA is working
uv run python -c "import torch; print(torch.__version__, torch.cuda.get_device_name(0))"

If torch is not CUDA-enabled after uv sync:

uv pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128 --reinstall

Project Structure

drone_tracker/
├── Anti-UAV410/                  <- raw dataset (train / val / test sequences)
├── datasets/                     <- YOLO-format dataset (created by convert step)
│   ├── images/train|val|test/
│   ├── labels/train|val|test/
│   └── data.yaml
├── configs/
│   ├── drone_yolo11s.yaml        <- YOLOv11s + P2 head architecture
│   ├── train.yaml                <- all training hyperparameters
│   └── bytetrack.yaml            <- ByteTrack tracker config
├── src/
│   ├── data/
│   │   ├── convert_dataset.py    <- Anti-UAV410 JSON → YOLO flat-file
│   │   ├── augmentation.py       <- 10-technique Albumentations pipeline
│   │   └── dataset_stats.py      <- dataset analysis and plots
│   ├── training/
│   │   ├── train.py              <- training orchestrator + callbacks
│   │   └── loss.py               <- VarifocalLoss + SIoU
│   ├── eval/
│   │   ├── evaluate_detection.py <- mAP50/95, precision, recall
│   │   └── evaluate_tracking.py  <- MOTA, IDF1, ID switches, FPS
│   ├── inference/
│   │   ├── tracker.py            <- production inference engine
│   │   ├── overlay.py            <- HUD renderer
│   │   ├── app.py                <- video file + live camera app
│   │   └── demo_compare.py       <- GT vs prediction comparison video
│   └── export/
│       └── export_model.py       <- ONNX export + benchmark
├── runs/
│   ├── drone_tracker/            <- training outputs (weights, plots, metrics)
│   └── demo/                     <- demo comparison videos
├── main.py                       <- unified CLI entry point
└── pyproject.toml

Workflow

The complete end-to-end workflow is available here.

Performance Metrics

Metric Achieved
mAP@50 0.8515
mAP@50-95 0.4868
Precision 0.9507
Recall 0.8500
MOTA 0.7375
IDF1 0.3086
ID Switches 6333
FPS (end-to-end) 58.9 *

*Tested on a laptop with Intel i5 13th gen, 16GB RAM, and NVIDIA GeForce RTX 4050 Laptop GPU with torch cu128

Architecture

Component Choice Reason
Backbone YOLOv11s Best tiny-object accuracy with efficient VRAM usage
Detection heads P2 + P3 + P4 + P5 P2 (stride 4) catches drones ≤ 16 px
Classification loss Varifocal Loss Handles extreme foreground/background imbalance
Box loss CIoU + DFL Built into YOLOv11, superior to vanilla IoU
Tracker ByteTrack Fast, appearance-free, handles occlusions well
Augmentation Albumentations Motion blur, fog, rain, noise, occlusion patches
Frame sampling stride=10 Adjacent video frames are ~95% identical

Troubleshooting

UnicodeEncodeError in terminal

$env:PYTHONUTF8=1; uv run python main.py <command>

CUDA not detected after uv sync

uv pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128 --reinstall

CUDA out of memory during training
configs/train.yaml → set batch: 8 or 4

Run all commands from the drone_tracker/ directory (where Anti-UAV410/ lives).

Dataset Acknowledgment

This project trains on the Anti-UAV410 dataset from IEEE Xplore. I only gathered and formatted their dataset for this project; all credit for the dataset collection, annotation, and baseline research goes to the original authors.

Citation:

@article{huang2023anti,
  title={Anti-UAV410: A Thermal Infrared Benchmark and Customized Scheme for Tracking Drones in the Wild},
  author={Huang, Bo and Li, Jianan and Chen, Junjie and Wang, Gang and Zhao, Jian and Xu, Tingfa},
  journal={IEEE Transactions on Pattern Analysis and Machine Intelligence},
  year={2023},
  publisher={IEEE}
}

About

Real-time drone detection and tracking system designed to identify and track small UAVs

Topics

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages