yjh0410 / RT-ODLab

YOLO Tutorial
MIT License
136 stars 29 forks source link

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.

image

Requirements

  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:

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

Experiments

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

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

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:

image

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.

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
python track.py --mode camera \
                --cuda \
                --img_size 640 \
                -dt yolov2 \
                -tk byte_tracker \
                --weight path/to/yolov2_coco.pth \
                --show \
                --gif

Tracking visualization

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:

image

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:

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

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', ),
    },
}
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
|  ...
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

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

For example:

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

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