Mask R-CNN for object detection and instance segmentation on Pytorch
This is an implementation of the instance segmentation model Mask R-CNN on Pytorch, based on the previous work of Matterport and lasseha. Matterport's repository is an implementation on Keras and TensorFlow while lasseha's repository is an implementation on Pytorch.
Compared with other PyTorch implementations, this repository has the following features:
The instructions come from lasseha's repository.
We use the Non-Maximum Suppression from ruotianluo and the RoiAlign from longcw. Please follow the instructions below to build the functions.
cd nms/src/cuda/
nvcc -c -o nms_kernel.cu.o nms_kernel.cu -x cu -Xcompiler -fPIC -arch=arch
cd ../../
python build.py
cd ../
cd roialign/roi_align/src/cuda/
nvcc -c -o crop_and_resize_kernel.cu.o crop_and_resize_kernel.cu -x cu -Xcompiler -fPIC -arch=arch
cd ../../
python build.py
cd ../../
where 'arch' is determined by your GPU model:
GPU | TitanX | GTX 960M | GTX 1070 | GTX 1080 (Ti) |
---|---|---|---|---|
arch | sm_52 | sm_50 | sm_61 | sm_61 |
If you want to train the network on the COCO dataset, please install the Python COCO API and create a symlink.
ln -s /path/to/coco/cocoapi/PythonAPI/pycocotools/ pycocotools
The pretrained models on COCO and ImageNet are available here.
The training and evaluation is based on COCO Dataset 2014. To understand the indicators below, please have a look at pycocotools. Notably, I only used one GTX 1080 (Ti). I think the performance could be improved if more GPUs are available.
Indicator | IoU | area | maxDets | Value |
---|---|---|---|---|
Average Precision (AP) | 0.50:0.95 | all | 100 | 0.392 |
Average Precision (AP) | 0.50 | all | 100 | 0.574 |
Average Precision (AP) | 0.75 | all | 100 | 0.434 |
Average Precision (AP) | 0.50:0.95 | small | 100 | 0.199 |
Average Precision (AP) | 0.50:0.95 | medium | 100 | 0.448 |
Average Precision (AP) | 0.50:0.95 | large | 100 | 0.575 |
Average Recall (AR) | 0.50:0.95 | all | 1 | 0.321 |
Average Recall (AR) | 0.50:0.95 | all | 10 | 0.445 |
Average Recall (AR) | 0.50:0.95 | all | 100 | 0.457 |
Average Recall (AR) | 0.50:0.95 | small | 100 | 0.231 |
Average Recall (AR) | 0.50:0.95 | medium | 100 | 0.508 |
Average Recall (AR) | 0.50:0.95 | large | 100 | 0.645 |