Differentiable IoU of Oriented Boxes
(image source: here)
Introduction
This repo is an unofficial implementation of IoU Loss for 2D/3D Object Detection. It contains the Pytorch function which calculates the intersection area of oriented rectangles using GPU.
Note
The CUDA extension is modified recently to cover some corner cases. Please consider to update the code and re-compile the extension.
Check List
- Pytorch function to find intersection points of oriented rectangles
- Pytorch function to check if corners of one rectangle lie in another
- CUDA extension to anti-clockwise sort vertices of the intersection polygon of two rectangles
- Pytorch function to calculate the intersection of area of rectangles using functions above
- Test cases
- Rotated 2d/3d GIoU and DIoU loss
- Demo to validate the back-propagation
- Validate 2d/3d IoU loss in Object detection
Requirements
Code is tested on Ubuntu 18.04. Following dependencies are needed
cudatoolkit=10.2
pytorch=1.5 # newer version should work as well
numpy
matplotlib
argparse
Usage
First, compile the CUDA extension.
cd cuda_op
python setup.py install
Then, run a demo which validate the Pytorch functions and CUDA extension.
cd ..
python demo.py
This demo trains a network which takes N set of box corners and predicts the x, y, w, h and angle of each rotated boxes. In order to do the back-prop, the predicted box parameters and the GT are converted to coordinates of box corners. The area of intersection is calculated using the Pytorch function with CUDA extension. Then, the GIoU loss or DIoU loss can be calculated. This demo first generates data and then do the training.
You are expected to see some information like followings:
... generating 819200 boxes, please wait...
... generating 81920 boxes, please wait...
data saved in: ./data
[Epoch 1: 10/200] train loss: 0.6721 mean_iou: 0.3703
[Epoch 1: 20/200] train loss: 0.4931 mean_iou: 0.5211
[Epoch 1: 30/200] train loss: 0.4532 mean_iou: 0.5546
[Epoch 1: 40/200] train loss: 0.4249 mean_iou: 0.5805
[Epoch 1: 50/200] train loss: 0.4361 mean_iou: 0.5713
[Epoch 1: 60/200] train loss: 0.4148 mean_iou: 0.5910
[Epoch 1: 70/200] train loss: 0.4272 mean_iou: 0.5803
[Epoch 1: 80/200] train loss: 0.4283 mean_iou: 0.5801
[Epoch 1: 90/200] train loss: 0.4203 mean_iou: 0.5879
Note the train loss drops and the mean_iou increases, which shows the functions are differentiable.
GIoU and DIoU
This repo implements both GIoU-loss and DIoU-loss for rotated bounding boxes. In the demo, they can be chosen with
python demo.py --loss giou
python demo.py --loss diou # [default]
Both losses need the smallest enclosing box of two boxes. Note there are different choices to determin the enclosing box.
- axis-aligned box: the enclosing box is axis-aligned. This version is simple and fast but theortically non-optimal.
- rotated box (approximated): the enclosing box is rotated as well. The size of rotated enclosing box can be estimated using PCA. The calculation if relatively simple but the result is not accurate. In the demo, this methode seems work well.
- rotated box (accurate): real smallest enclosing bounding box. Since the brutal force search is used to get the minimum bounding box, the computational cost is high.
The three types of enclosing box can be chosen with:
python demo.py --enclosing aligned # simple and naive. bad performance.
python demo.py --enclosing pca # approximated smallest box. slightly worse performance.
python demo.py --enclosing smallest # [default]. smallest box. best performance.
Acknowledgement
The idea of calculating intersection area is inspired by this paper:
@INPROCEEDINGS{8886046,
author={D. {Zhou} and J. {Fang} and X. {Song} and C. {Guan} and J. {Yin} and Y. {Dai} and R. {Yang}},
booktitle={2019 International Conference on 3D Vision (3DV)},
title={IoU Loss for 2D/3D Object Detection},
year={2019},
pages={85-94},}
Some code for CUDA extension is modified from:
@article{pytorchpointnet++,
Author = {Erik Wijmans},
Title = {Pointnet++ Pytorch},
Journal = {https://github.com/erikwijmans/Pointnet2_PyTorch},
Year = {2018}
}