Skip to content

This is a PyTorch-based R-YOLOv4 implementation which combines YOLOv4 model and loss function from R3Det for arbitrary oriented object detection.

Notifications You must be signed in to change notification settings

yingkunwu/R-YOLOv4

Repository files navigation

R-YOLOv4

Introduction

The objective of this project is to adapt YOLOv4 model to detecting oriented objects. As a result, modifying the original loss function of the model is required. I got a successful result by increasing the number of anchor boxes with different rotating angle and combining smooth-L1-IoU loss function proposed by R3Det: Refined Single-Stage Detector with Feature Refinement for Rotating Object into the original loss for bounding boxes.

Features


Loss Function (only for x, y, w, h, theta)

loss

angle


*UPDATE*

The followings are updated to the code to improve the performance of the original setting.

  • Circular Smooth Labeling
  • KFIoU
  • Yolov5
  • Yolov7

Setup

  1. Clone repository
    $ git clone https://github.com/kunnnnethan/R-YOLOv4.git
    $ cd R-YOLOv4/
    
  2. Create Environment
  • Conda

    1. Create virual environment
      $ conda create -n ryolo python=3.8
      $ conda activate ryolo
      
    2. Install PyTorch and torchvision following the official instructions, e.g.,
      If you are using CUDA 11.8 version
      $ conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia
      
    3. Install required libraries
      $ pip install -r requirements.txt
      
    4. Install detectron2 for calculating SkewIoU on GPU following the official instructions, e.g.,
      python -m pip install 'git+https://github.com/facebookresearch/detectron2.git'
      
  • Docker

    $ docker build -t ryolo docker/
    $ sudo docker run -it --rm --gpus all --ipc=host --ulimit memlock=-1 -v ${your_path}:/workspace ryolo
    
  1. Download pretrained weights
    weights

  2. Make sure your files arrangment looks like the following
    Note that each of your dataset folder in data should split into three files, namely train, test, and detect.

    R-YOLOv4/
    ├── train.py
    ├── test.py
    ├── detect.py
    ├── xml2txt.py
    ├── environment.xml
    ├── requirements.txt
    ├── model/
    ├── datasets/
    ├── lib/
    ├── outputs/
    ├── weights/
        ├── pretrained/ (for training)
        └── UCAS-AOD/ (for testing and detection)
    └── data/
        └── UCAS-AOD/
            ├── class.names
            ├── train/
                ├── ...png
                └── ...txt
            ├── test/
                ├── ...png
                └── ...txt
            └── detect/
                └── ...png
    

Train

I have implemented methods to load and train three different datasets. They are UCAS-AOD, DOTA, and custom dataset respectively. You can check out how I loaded those dataset into the model at /datasets. The angle of each bounding box is limited in (- pi/2, pi/2], and the height of each bounding box is always longer than it's width.

Check all the setting in .yaml files that you are going to use in the /data folder.

$ python train.py --model_name DOTA_yolov7_csl_800 --config data/hyp.yaml --img_size 800 --data data/DOTA.yaml --epochs 100 --mode csl --ver yolov7

You can run display_inputs.py to visualize whether your data is loaded successfully.

UCAS-AOD dataset

Please refer to this repository to rearrange files so that it can be loaded and trained by this model.

DOTA dataset

Download the official dataset from here. The original files should be able to be loaded and trained by this model.

Train with custom dataset

  1. Use labelImg2 to help label your data. labelImg2 is capable of labeling rotated objects.
  2. Move your data folder into the R-YOLOv4/data folder.
  3. Run xml2txt.py
    1. generate txt files: python xml2txt.py --data_folder your-path --action gen_txt
    2. delete xml files: python xml2txt.py --data_folder your-path --action del_xml

A trash custom dataset that I made are provided for your convenience.

Test

$ python test.py --data data/DOTA.yaml --hyp data/hyp.yaml --weight_path weights/DOTA_yolov7_csl_800/best.pth --batch_size 8 --img_size 800 --mode csl --ver yolov7

detect

$ python detect.py --data data/UCAS_AOD.yaml --hyp data/hyp.yaml --weight_path weights/DOTA_yolov7_csl_800/best.pth --batch_size 8 --img_size 800 --conf_thres 0.3 --mode csl --ver yolov7

Tensorboard

If you would like to use tensorboard for tracking traing process.

  • Open additional terminal in the same folder where you are running program.
  • Run command $ tensorboard --logdir=weights --port=6006
  • Go to http://localhost:6006/

Results

UCAS_AOD

car

plane

DOTA

DOTADOTA

trash (custom dataset)

garbage1

garbage2

References

WongKinYiu/yolov7
ultralytics/yolov5
Tianxiaomo/pytorch-YOLOv4
yangxue0827/RotationDetection
eriklindernoren/PyTorch-YOLOv3

YOLOv4: Optimal Speed and Accuracy of Object Detection

@article{yolov4,
  title={YOLOv4: Optimal Speed and Accuracy of Object Detection},
  author={Alexey Bochkovskiy, Chien-Yao Wang, Hong-Yuan Mark Liao},
  journal = {arXiv},
  year={2020}
}

R3Det: Refined Single-Stage Detector with Feature Refinement for Rotating Object

@article{r3det,
  title={R3Det: Refined Single-Stage Detector with Feature Refinement for Rotating Object},
  author={Xue Yang, Junchi Yan, Ziming Feng, Tao He},
  journal = {arXiv},
  year={2019}
}

About

This is a PyTorch-based R-YOLOv4 implementation which combines YOLOv4 model and loss function from R3Det for arbitrary oriented object detection.

Topics

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published