RT-ODLab

Real-time General Object Detection

English | 简体中文

We are trying to build our real-time general target detection code base based on the core concepts of YOLO. We have reproduced most of the YOLO series. In addition, we have also written an introductory tutorial on YOLO. We hope that by learning YOLO, a very popular general target detection framework, beginners can master the basic knowledge necessary to study general target detection.

If you are interested in our book, you can purchase it on e-commerce platforms such as Taobao and JD.com in China.

Requirements

• We recommend you to use Anaconda to create a conda environment:

  conda create -n rtcdet python=3.6
  

• Then, activate the environment:

  conda activate rtcdet
  

• Requirements:

• Install necessary libraies

  pip install -r requirements.txt 
  

1. (optional) Compile MSDeformableAttention ops for DETR series

   cd ./ppdet/modeling/transformers/ext_op/
   
   python setup_ms_deformable_attn_op.py install
   

See details

My environment:

• PyTorch = 1.9.1
• Torchvision = 0.10.1

At least, please make sure your torch is version 1.x.

Experiments

VOC

• Download VOC.

  cd <RT-ODLab>
  cd dataset/scripts/
  sh VOC2007.sh
  sh VOC2012.sh
  

• Check VOC

  cd <RT-ODLab>
  python dataset/voc.py
  

• Train on VOC

For example:

python train.py --cuda -d voc --root path/to/VOCdevkit -m yolov1 -bs 16 --max_epoch 150 --wp_epoch 1 --eval_epoch 10 --fp16 --ema --multi_scale

COCO

• Download COCO.

  cd <RT-ODLab>
  cd dataset/scripts/
  sh COCO2017.sh
  

• Clean COCO

  cd <RT-ODLab>
  cd tools/
  python clean_coco.py --root path/to/coco --image_set val
  python clean_coco.py --root path/to/coco --image_set train
  

• Check COCO

  cd <RT-ODLab>
  python dataset/coco.py
  

• Train on COCO

For example:

python train.py --cuda -d coco --root path/to/COCO -m yolov1 -bs 16 --max_epoch 150 --wp_epoch 1 --eval_epoch 10 --fp16 --ema --multi_scale

We also kindly provide a script train.sh to run the training code. You need to follow the following format to use this script：

bash train.sh <model> <data> <data_path> <batch_size> <num_gpus> <master_port> <resume_weight>

For example, we use this script to train YOLOv3 from the epoch-0:

bash train.sh yolov3 coco path/to/coco 128 4 1699 None

We can also continue training from existing weights by passing the model's weight file to the resume parameter.

bash train.sh yolov3 coco path/to/coco 128 4 1699 path/to/yolov3.pth

Test

python test.py --cuda --dataset coco --root path/to/coco --model yolov1 --weight path/to/yolov1.pth --img_size 640 --show 

Evaluation

python eval.py --cuda --dataset coco --root path/to/coco --model yolov1 --weight path/to/yolov1.pth --img_size 640

Demo

I have provide some images in data/demo/images/, so you can run following command to run a demo with coco pretrained model:

python demo.py --mode image \
               --path_to_img data/demo/images/ \
               --dataset coco \
               --cuda \
               --img_size 640 \
               --model yolov2 \
               --weight path/to/yolov2_coco.pth \
               --show

If you want to try this command with voc pretrained model, you could refer to the following command:

python demo.py --mode image \
               --path_to_img data/demo/images/ \
               --dataset voc \
               --cuda \
               --img_size 640 \
               --model yolov2 \
               --weight path/to/yolov2_voc.pth \
               --show

If you want run a demo of streaming video detection, you need to set --mode to video, and give the path to video --path_to_vid。

python demo.py --mode video \
               --path_to_vid data/demo/videos/your_video \
               --dataset coco \
               --cuda \
               --img_size 640 \
               --model yolov2 \
               --weight path/to/yolov2_coco.pth \
               --show \
               --gif

If you want run video detection with your camera, you need to set --mode to camera。

python demo.py --mode camera \
               --cuda \
               --dataset coco \
               --img_size 640 \
               --model yolov2 \
               --weight path/to/yolov2_coco.pth \
               --show \
               --gif

Detection visualization

• Detector: YOLOv2

Command：

python demo.py --mode video \
               --path_to_vid ./dataset/demo/videos/000006.mp4 \
               --cuda \
               --dataset coco \
               --img_size 640 \
               --model yolov2 \
               --weight path/to/yolov2_coco.pth \
               --show \
               --gif

Results:

Tracking

Our project also supports multi-object tracking tasks. We use the YOLO of this project as the detector, following the "tracking-by-detection" framework, and use the simple and efficient ByteTrack as the tracker.

• images tracking

  python track.py --mode image \
              --path_to_img path/to/images/ \
              --cuda \
              --img_size 640 \
              -dt yolov2 \
              -tk byte_tracker \
              --weight path/to/yolov2_coco.pth \
              --show \
              --gif
  

• video tracking

  python track.py --mode video \
              --path_to_img path/to/video/ \
              --cuda \
              --img_size 640 \
              -dt yolov2 \
              -tk byte_tracker \
              --weight path/to/yolov2_coco.pth \
              --show \
              --gif
  

• camera tracking

  python track.py --mode camera \
              --cuda \
              --img_size 640 \
              -dt yolov2 \
              -tk byte_tracker \
              --weight path/to/yolov2_coco.pth \
              --show \
              --gif
  

Tracking visualization

• Detector: YOLOv2
• Tracker: ByteTracker
• Device: i5-12500H CPU

Command：

python track.py --mode video \
                --path_to_img ./dataset/demo/videos/000006.mp4 \
                --img_size 640 \
                -dt yolov2 \
                -tk byte_tracker \
                --weight path/to/yolov2_coco.pth \
                --show \
                --gif

Results:

Train on custom dataset

Besides the popular datasets, we can also train the model on ourself dataset. To achieve this goal, you should follow these steps:

• Step-1: Prepare the images (JPG/JPEG/PNG ...) and use labelimg to make XML format annotation files.

  CustomedDataset
  |_ train
  |  |_ images     
  |     |_ 0.jpg
  |     |_ 1.jpg
  |     |_ ...
  |  |_ annotations
  |     |_ 0.xml
  |     |_ 1.xml
  |     |_ ...
  |_ val
  |  |_ images     
  |     |_ 0.jpg
  |     |_ 1.jpg
  |     |_ ...
  |  |_ annotations
  |     |_ 0.xml
  |     |_ 1.xml
  |     |_ ...
  |  ...
  

• Step-2: Make the configuration for our dataset.

  cd <RT-ODLab>
  cd config/data_config
  

  You need to edit the dataset_cfg defined in dataset_config.py. You can refer to the customed defined in dataset_cfg to modify the relevant parameters, such as num_classes, classes_names, to adapt to our dataset.

For example:

dataset_cfg = {
    'customed':{
        'data_name': 'AnimalDataset',
        'num_classes': 9,
        'class_indexs': (0, 1, 2, 3, 4, 5, 6, 7, 8),
        'class_names': ('bird', 'butterfly', 'cat', 'cow', 'dog', 'lion', 'person', 'pig', 'tiger', ),
    },
}

• Step-3: Convert customed to COCO format.

  cd <RT-ODLab>
  cd tools
  # convert train split
  python convert_ours_to_coco.py --root path/to/dataset/ --split train
  # convert val split
  python convert_ours_to_coco.py --root path/to/dataset/ --split val
  

Then, we can get a train.json file and a val.json file, as shown below.

CustomedDataset
|_ train
|  |_ images     
|     |_ 0.jpg
|     |_ 1.jpg
|     |_ ...
|  |_ annotations
|     |_ 0.xml
|     |_ 1.xml
|     |_ ...
|     |_ train.json
|_ val
|  |_ images     
|     |_ 0.jpg
|     |_ 1.jpg
|     |_ ...
|  |_ annotations
|     |_ 0.xml
|     |_ 1.xml
|     |_ ...
|     |_ val.json
|  ...

• Step-4 Check the data.

  cd <RT-ODLab>
  cd dataset
  # convert train split
  python customed.py --root path/to/dataset/ --split train
  # convert val split
  python customed.py --root path/to/dataset/ --split val
  

• Step-5 Train

For example:

cd <RT-ODLab>
python train.py --root path/to/dataset/ -d customed -m yolov1 -bs 16 --max_epoch 100 --wp_epoch 1 --eval_epoch 5 -p path/to/yolov1_coco.pth

• Step-6 Test

For example:

cd <RT-ODLab>
python test.py --root path/to/dataset/ -d customed -m yolov1 --weight path/to/checkpoint --show

• Step-7 Eval

For example:

cd <RT-ODLab>
python eval.py --root path/to/dataset/ -d customed -m yolov1 --weight path/to/checkpoint

Deployment

1. ONNX export and an ONNXRuntime