Repository: jwyang/faster-rcnn.pytorch Branch: master Commit: f9d984d27b48 Files: 122 Total size: 444.5 KB Directory structure: gitextract_3vzfjfkd/ ├── .gitignore ├── LICENSE ├── README.md ├── _init_paths.py ├── cfgs/ │ ├── res101.yml │ ├── res101_ls.yml │ ├── res50.yml │ └── vgg16.yml ├── demo.py ├── lib/ │ ├── datasets/ │ │ ├── VOCdevkit-matlab-wrapper/ │ │ │ ├── get_voc_opts.m │ │ │ ├── voc_eval.m │ │ │ └── xVOCap.m │ │ ├── __init__.py │ │ ├── coco.py │ │ ├── ds_utils.py │ │ ├── factory.py │ │ ├── imagenet.py │ │ ├── imdb.py │ │ ├── pascal_voc.py │ │ ├── pascal_voc_rbg.py │ │ ├── tools/ │ │ │ └── mcg_munge.py │ │ ├── vg.py │ │ ├── vg_eval.py │ │ └── voc_eval.py │ ├── make.sh │ ├── model/ │ │ ├── __init__.py │ │ ├── faster_rcnn/ │ │ │ ├── __init__.py │ │ │ ├── faster_rcnn.py │ │ │ ├── resnet.py │ │ │ └── vgg16.py │ │ ├── nms/ │ │ │ ├── .gitignore │ │ │ ├── __init__.py │ │ │ ├── _ext/ │ │ │ │ ├── __init__.py │ │ │ │ └── nms/ │ │ │ │ └── __init__.py │ │ │ ├── build.py │ │ │ ├── make.sh │ │ │ ├── nms_cpu.py │ │ │ ├── nms_gpu.py │ │ │ ├── nms_kernel.cu │ │ │ ├── nms_wrapper.py │ │ │ └── src/ │ │ │ ├── nms_cuda.h │ │ │ ├── nms_cuda_kernel.cu │ │ │ └── nms_cuda_kernel.h │ │ ├── roi_align/ │ │ │ ├── __init__.py │ │ │ ├── _ext/ │ │ │ │ ├── __init__.py │ │ │ │ └── roi_align/ │ │ │ │ └── __init__.py │ │ │ ├── build.py │ │ │ ├── functions/ │ │ │ │ ├── __init__.py │ │ │ │ └── roi_align.py │ │ │ ├── make.sh │ │ │ ├── modules/ │ │ │ │ ├── __init__.py │ │ │ │ └── roi_align.py │ │ │ └── src/ │ │ │ ├── roi_align.c │ │ │ ├── roi_align.h │ │ │ ├── roi_align_cuda.c │ │ │ ├── roi_align_cuda.h │ │ │ ├── roi_align_kernel.cu │ │ │ └── roi_align_kernel.h │ │ ├── roi_crop/ │ │ │ ├── __init__.py │ │ │ ├── _ext/ │ │ │ │ ├── __init__.py │ │ │ │ ├── crop_resize/ │ │ │ │ │ └── __init__.py │ │ │ │ └── roi_crop/ │ │ │ │ └── __init__.py │ │ │ ├── build.py │ │ │ ├── functions/ │ │ │ │ ├── __init__.py │ │ │ │ ├── crop_resize.py │ │ │ │ ├── gridgen.py │ │ │ │ └── roi_crop.py │ │ │ ├── make.sh │ │ │ ├── modules/ │ │ │ │ ├── __init__.py │ │ │ │ ├── gridgen.py │ │ │ │ └── roi_crop.py │ │ │ └── src/ │ │ │ ├── roi_crop.c │ │ │ ├── roi_crop.h │ │ │ ├── roi_crop_cuda.c │ │ │ ├── roi_crop_cuda.h │ │ │ ├── roi_crop_cuda_kernel.cu │ │ │ └── roi_crop_cuda_kernel.h │ │ ├── roi_pooling/ │ │ │ ├── __init__.py │ │ │ ├── _ext/ │ │ │ │ ├── __init__.py │ │ │ │ └── roi_pooling/ │ │ │ │ └── __init__.py │ │ │ ├── build.py │ │ │ ├── functions/ │ │ │ │ ├── __init__.py │ │ │ │ └── roi_pool.py │ │ │ ├── modules/ │ │ │ │ ├── __init__.py │ │ │ │ └── roi_pool.py │ │ │ └── src/ │ │ │ ├── roi_pooling.c │ │ │ ├── roi_pooling.h │ │ │ ├── roi_pooling_cuda.c │ │ │ ├── roi_pooling_cuda.h │ │ │ ├── roi_pooling_kernel.cu │ │ │ └── roi_pooling_kernel.h │ │ ├── rpn/ │ │ │ ├── __init__.py │ │ │ ├── anchor_target_layer.py │ │ │ ├── bbox_transform.py │ │ │ ├── generate_anchors.py │ │ │ ├── proposal_layer.py │ │ │ ├── proposal_target_layer_cascade.py │ │ │ └── rpn.py │ │ └── utils/ │ │ ├── .gitignore │ │ ├── __init__.py │ │ ├── bbox.pyx │ │ ├── blob.py │ │ ├── config.py │ │ ├── logger.py │ │ └── net_utils.py │ ├── pycocotools/ │ │ ├── UPSTREAM_REV │ │ ├── __init__.py │ │ ├── _mask.pyx │ │ ├── coco.py │ │ ├── cocoeval.py │ │ ├── license.txt │ │ ├── mask.py │ │ ├── maskApi.c │ │ └── maskApi.h │ ├── roi_data_layer/ │ │ ├── __init__.py │ │ ├── minibatch.py │ │ ├── roibatchLoader.py │ │ └── roidb.py │ └── setup.py ├── requirements.txt ├── test_net.py └── trainval_net.py ================================================ FILE CONTENTS ================================================ ================================================ FILE: .gitignore ================================================ data/* # READ THIS BEFORE YOU REFACTOR ME # # setup.py uses the list of patterns in this file to decide # what to delete, but it's not 100% sound. So, for example, # if you delete aten/build/ because it's redundant with build/, # aten/build/ will stop being cleaned. So be careful when # refactoring this file! ## PyTorch .mypy_cache *.pyc */*.pyc */*.so* */**/__pycache__ */**/*.dylib* */**/*.pyc */**/*.pyd */**/*.so* */**/**/*.pyc */**/**/**/*.pyc */**/**/**/**/*.pyc aten/build/ aten/src/ATen/Config.h aten/src/ATen/cuda/CUDAConfig.h build/ dist/ docs/src/**/* test/.coverage test/cpp/api/mnist test/data/gpu_tensors.pt test/data/legacy_modules.t7 test/data/legacy_serialized.pt test/data/linear.pt test/htmlcov third_party/build/ tools/shared/_utils_internal.py torch.egg-info/ torch/csrc/autograd/generated/* torch/csrc/cudnn/cuDNN.cpp torch/csrc/generated torch/csrc/generic/TensorMethods.cpp torch/csrc/jit/generated/* torch/csrc/nn/THCUNN.cpp torch/csrc/nn/THCUNN.cwrap torch/csrc/nn/THNN_generic.cpp torch/csrc/nn/THNN_generic.cwrap torch/csrc/nn/THNN_generic.h torch/csrc/nn/THNN.cpp torch/csrc/nn/THNN.cwrap torch/lib/*.a* torch/lib/*.dll* torch/lib/*.dylib* torch/lib/*.h torch/lib/*.lib torch/lib/*.so* torch/lib/build torch/lib/cmake torch/lib/include torch/lib/pkgconfig torch/lib/protoc torch/lib/tmp_install torch/lib/torch_shm_manager torch/version.py # IPython notebook checkpoints .ipynb_checkpoints # Editor temporaries *.swn *.swo *.swp *.swm *~ # macOS dir files .DS_Store # Symbolic files tools/shared/cwrap_common.py # Ninja files .ninja_deps .ninja_log compile_commands.json *.egg-info/ docs/source/scripts/activation_images/ ## General # Compiled Object files *.slo *.lo *.o *.cuo *.obj # Compiled Dynamic libraries *.so *.dylib *.dll # Compiled Static libraries *.lai *.la *.a *.lib # Compiled protocol buffers *.pb.h *.pb.cc *_pb2.py # Compiled python *.pyc *.pyd # Compiled MATLAB *.mex* # IPython notebook checkpoints .ipynb_checkpoints # Editor temporaries *.swn *.swo *.swp *~ # Sublime Text settings *.sublime-workspace *.sublime-project # Eclipse Project settings *.*project .settings # QtCreator files *.user # PyCharm files .idea # Visual Studio Code files .vscode .vs # OSX dir files .DS_Store ## Caffe2 # build, distribute, and bins (+ python proto bindings) build build_host_protoc build_android build_ios /build_* .build_debug/* .build_release/* distribute/* *.testbin *.bin cmake_build .cmake_build gen .setuptools-cmake-build .pytest_cache aten/build/* # Bram plsdontbreak # Generated documentation docs/_site docs/gathered _site doxygen docs/dev # LevelDB files *.sst *.ldb LOCK LOG* CURRENT MANIFEST-* # generated version file caffe2/version.py # setup.py intermediates .eggs caffe2.egg-info # Atom/Watchman required file .watchmanconfig # cython generated files lib/model/utils/bbox.c lib/pycocotools/_mask.c ================================================ FILE: LICENSE ================================================ MIT License Copyright (c) 2017 Jianwei Yang Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. ================================================ FILE: README.md ================================================ # A *Faster* Pytorch Implementation of Faster R-CNN ## Write at the beginning [05/29/2020] This repo was initaited about two years ago, developed as the first open-sourced object detection code which supports multi-gpu training. It has been integrating tremendous efforts from many people. However, we have seen many high-quality repos emerged in the last years, such as: * [maskrcnn-benchmark](https://github.com/facebookresearch/maskrcnn-benchmark) * [detectron2](https://github.com/facebookresearch/detectron2) * [mmdetection](https://github.com/open-mmlab/mmdetection) **At this point, I think this repo is out-of-data in terms of the pipeline and coding style, and will not maintain actively. Though you can still use this repo as a playground, I highly recommend you move to the above repos to delve into west world of object detection!** ## Introduction ### :boom: Good news! This repo supports pytorch-1.0 now!!! We borrowed some code and techniques from [maskrcnn-benchmark](https://github.com/facebookresearch/maskrcnn-benchmark). Just go to pytorch-1.0 branch! This project is a *faster* pytorch implementation of faster R-CNN, aimed to accelerating the training of faster R-CNN object detection models. Recently, there are a number of good implementations: * [rbgirshick/py-faster-rcnn](https://github.com/rbgirshick/py-faster-rcnn), developed based on Pycaffe + Numpy * [longcw/faster_rcnn_pytorch](https://github.com/longcw/faster_rcnn_pytorch), developed based on Pytorch + Numpy * [endernewton/tf-faster-rcnn](https://github.com/endernewton/tf-faster-rcnn), developed based on TensorFlow + Numpy * [ruotianluo/pytorch-faster-rcnn](https://github.com/ruotianluo/pytorch-faster-rcnn), developed based on Pytorch + TensorFlow + Numpy During our implementing, we referred the above implementations, especailly [longcw/faster_rcnn_pytorch](https://github.com/longcw/faster_rcnn_pytorch). However, our implementation has several unique and new features compared with the above implementations: * **It is pure Pytorch code**. We convert all the numpy implementations to pytorch! * **It supports multi-image batch training**. We revise all the layers, including dataloader, rpn, roi-pooling, etc., to support multiple images in each minibatch. * **It supports multiple GPUs training**. We use a multiple GPU wrapper (nn.DataParallel here) to make it flexible to use one or more GPUs, as a merit of the above two features. * **It supports three pooling methods**. We integrate three pooling methods: roi pooing, roi align and roi crop. More importantly, we modify all of them to support multi-image batch training. * **It is memory efficient**. We limit the image aspect ratio, and group images with similar aspect ratios into a minibatch. As such, we can train resnet101 and VGG16 with batchsize = 4 (4 images) on a single Titan X (12 GB). When training with 8 GPU, the maximum batchsize for each GPU is 3 (Res101), totaling 24. * **It is faster**. Based on the above modifications, the training is much faster. We report the training speed on NVIDIA TITAN Xp in the tables below. ### What we are doing and going to do - [x] Support both python2 and python3 (great thanks to [cclauss](https://github.com/cclauss)). - [x] Add deformable pooling layer (mainly supported by [Xander](https://github.com/xanderchf)). - [x] Support pytorch-0.4.0 (this branch). - [x] Support tensorboardX. - [x] Support pytorch-1.0 (go to pytorch-1.0 branch). ## Other Implementations * [Feature Pyramid Network (FPN)](https://github.com/jwyang/fpn.pytorch) * [Mask R-CNN](https://github.com/roytseng-tw/mask-rcnn.pytorch) (~~ongoing~~ already implemented by [roytseng-tw](https://github.com/roytseng-tw)) * [Graph R-CNN](https://github.com/jwyang/graph-rcnn.pytorch) (extension to scene graph generation) ## Tutorial * [Blog](http://www.telesens.co/2018/03/11/object-detection-and-classification-using-r-cnns/) by [ankur6ue](https://github.com/ankur6ue) ## Benchmarking We benchmark our code thoroughly on three datasets: pascal voc, coco and visual genome, using two different network architectures: vgg16 and resnet101. Below are the results: 1). PASCAL VOC 2007 (Train/Test: 07trainval/07test, scale=600, ROI Align) model   | #GPUs | batch size | lr       | lr_decay | max_epoch     | time/epoch | mem/GPU | mAP ---------|--------|-----|--------|-----|-----|-------|--------|----- [VGG-16](https://www.dropbox.com/s/6ief4w7qzka6083/faster_rcnn_1_6_10021.pth?dl=0)     | 1 | 1 | 1e-3 | 5   | 6   | 0.76 hr | 3265MB   | 70.1 [VGG-16](https://www.dropbox.com/s/cpj2nu35am0f9hp/faster_rcnn_1_9_2504.pth?dl=0)     | 1 | 4 | 4e-3 | 8  | 9  | 0.50 hr | 9083MB   | 69.6 [VGG-16](https://www.dropbox.com/s/1a31y7vicby0kvy/faster_rcnn_1_10_625.pth?dl=0)     | 8 | 16| 1e-2 | 8   | 10 | 0.19 hr | 5291MB | 69.4 [VGG-16](https://www.dropbox.com/s/hkj7i6mbhw9tq4k/faster_rcnn_1_11_416.pth?dl=0)     | 8 | 24| 1e-2 | 10 | 11 | 0.16 hr | 11303MB | 69.2 [Res-101](https://www.dropbox.com/s/4v3or0054kzl19q/faster_rcnn_1_7_10021.pth?dl=0) | 1 | 1 | 1e-3 | 5 | 7 | 0.88 hr | 3200 MB | 75.2 [Res-101](https://www.dropbox.com/s/8bhldrds3mf0yuj/faster_rcnn_1_10_2504.pth?dl=0)   | 1 | 4 | 4e-3 | 8   | 10 | 0.60 hr | 9700 MB | 74.9 [Res-101](https://www.dropbox.com/s/5is50y01m1l9hbu/faster_rcnn_1_10_625.pth?dl=0)   | 8 | 16| 1e-2 | 8   | 10 | 0.23 hr | 8400 MB | 75.2  [Res-101](https://www.dropbox.com/s/cn8gneumg4gjo9i/faster_rcnn_1_12_416.pth?dl=0)   | 8 | 24| 1e-2 | 10 | 12 | 0.17 hr | 10327MB | 75.1   2). COCO (Train/Test: coco_train+coco_val-minival/minival, scale=800, max_size=1200, ROI Align) model | #GPUs | batch size |lr | lr_decay | max_epoch | time/epoch | mem/GPU | mAP ---------|--------|-----|--------|-----|-----|-------|--------|----- VGG-16     | 8 | 16   |1e-2| 4 | 6 | 4.9 hr | 7192 MB | 29.2 [Res-101](https://www.dropbox.com/s/5if6l7mqsi4rfk9/faster_rcnn_1_6_14657.pth?dl=0)   | 8 | 16   |1e-2| 4   | 6 | 6.0 hr |10956 MB | 36.2 [Res-101](https://www.dropbox.com/s/be0isevd22eikqb/faster_rcnn_1_10_14657.pth?dl=0)   | 8 | 16   |1e-2| 4   | 10 | 6.0 hr |10956 MB | 37.0 **NOTE**. Since the above models use scale=800, you need add "--ls" at the end of test command. 3). COCO (Train/Test: coco_train+coco_val-minival/minival, scale=600, max_size=1000, ROI Align) model | #GPUs | batch size |lr | lr_decay | max_epoch | time/epoch | mem/GPU | mAP ---------|--------|-----|--------|-----|-----|-------|--------|----- [Res-101](https://www.dropbox.com/s/y171ze1sdw1o2ph/faster_rcnn_1_6_9771.pth?dl=0)   | 8 | 24   |1e-2| 4   | 6 | 5.4 hr   |10659 MB | 33.9 [Res-101](https://www.dropbox.com/s/dpq6qv0efspelr3/faster_rcnn_1_10_9771.pth?dl=0)   | 8 | 24   |1e-2| 4   | 10 | 5.4 hr   |10659 MB | 34.5 4). Visual Genome (Train/Test: vg_train/vg_test, scale=600, max_size=1000, ROI Align, category=2500) model | #GPUs | batch size |lr | lr_decay | max_epoch | time/epoch | mem/GPU | mAP ---------|--------|-----|--------|-----|-----|-------|--------|----- [VGG-16](http://data.lip6.fr/cadene/faster-rcnn.pytorch/faster_rcnn_1_19_48611.pth)   | 1 P100 | 4   |1e-3| 5   | 20 | 3.7 hr   |12707 MB | 4.4 Thanks to [Remi](https://github.com/Cadene) for providing the pretrained detection model on visual genome! * Click the links in the above tables to download our pre-trained faster r-cnn models. * If not mentioned, the GPU we used is NVIDIA Titan X Pascal (12GB). ## Preparation First of all, clone the code ``` git clone https://github.com/jwyang/faster-rcnn.pytorch.git ``` Then, create a folder: ``` cd faster-rcnn.pytorch && mkdir data ``` ### prerequisites * Python 2.7 or 3.6 * Pytorch 0.4.0 (**now it does not support 0.4.1 or higher**) * CUDA 8.0 or higher ### Data Preparation * **PASCAL_VOC 07+12**: Please follow the instructions in [py-faster-rcnn](https://github.com/rbgirshick/py-faster-rcnn#beyond-the-demo-installation-for-training-and-testing-models) to prepare VOC datasets. Actually, you can refer to any others. After downloading the data, create softlinks in the folder data/. * **COCO**: Please also follow the instructions in [py-faster-rcnn](https://github.com/rbgirshick/py-faster-rcnn#beyond-the-demo-installation-for-training-and-testing-models) to prepare the data. * **Visual Genome**: Please follow the instructions in [bottom-up-attention](https://github.com/peteanderson80/bottom-up-attention) to prepare Visual Genome dataset. You need to download the images and object annotation files first, and then perform proprecessing to obtain the vocabulary and cleansed annotations based on the scripts provided in this repository. ### Pretrained Model We used two pretrained models in our experiments, VGG and ResNet101. You can download these two models from: * VGG16: [Dropbox](https://www.dropbox.com/s/s3brpk0bdq60nyb/vgg16_caffe.pth?dl=0), [VT Server](https://filebox.ece.vt.edu/~jw2yang/faster-rcnn/pretrained-base-models/vgg16_caffe.pth) * ResNet101: [Dropbox](https://www.dropbox.com/s/iev3tkbz5wyyuz9/resnet101_caffe.pth?dl=0), [VT Server](https://filebox.ece.vt.edu/~jw2yang/faster-rcnn/pretrained-base-models/resnet101_caffe.pth) Download them and put them into the data/pretrained_model/. **NOTE**. We compare the pretrained models from Pytorch and Caffe, and surprisingly find Caffe pretrained models have slightly better performance than Pytorch pretrained. We would suggest to use Caffe pretrained models from the above link to reproduce our results. **If you want to use pytorch pre-trained models, please remember to transpose images from BGR to RGB, and also use the same data transformer (minus mean and normalize) as used in pretrained model.** ### Compilation As pointed out by [ruotianluo/pytorch-faster-rcnn](https://github.com/ruotianluo/pytorch-faster-rcnn), choose the right `-arch` in `make.sh` file, to compile the cuda code: | GPU model | Architecture | | ------------- | ------------- | | TitanX (Maxwell/Pascal) | sm_52 | | GTX 960M | sm_50 | | GTX 1080 (Ti) | sm_61 | | Grid K520 (AWS g2.2xlarge) | sm_30 | | Tesla K80 (AWS p2.xlarge) | sm_37 | More details about setting the architecture can be found [here](https://developer.nvidia.com/cuda-gpus) or [here](http://arnon.dk/matching-sm-architectures-arch-and-gencode-for-various-nvidia-cards/) Install all the python dependencies using pip: ``` pip install -r requirements.txt ``` Compile the cuda dependencies using following simple commands: ``` cd lib sh make.sh ``` It will compile all the modules you need, including NMS, ROI_Pooing, ROI_Align and ROI_Crop. The default version is compiled with Python 2.7, please compile by yourself if you are using a different python version. **As pointed out in this [issue](https://github.com/jwyang/faster-rcnn.pytorch/issues/16), if you encounter some error during the compilation, you might miss to export the CUDA paths to your environment.** ## Train Before training, set the right directory to save and load the trained models. Change the arguments "save_dir" and "load_dir" in trainval_net.py and test_net.py to adapt to your environment. To train a faster R-CNN model with vgg16 on pascal_voc, simply run: ``` CUDA_VISIBLE_DEVICES=$GPU_ID python trainval_net.py \ --dataset pascal_voc --net vgg16 \ --bs $BATCH_SIZE --nw $WORKER_NUMBER \ --lr $LEARNING_RATE --lr_decay_step $DECAY_STEP \ --cuda ``` where 'bs' is the batch size with default 1. Alternatively, to train with resnet101 on pascal_voc, simple run: ``` CUDA_VISIBLE_DEVICES=$GPU_ID python trainval_net.py \ --dataset pascal_voc --net res101 \ --bs $BATCH_SIZE --nw $WORKER_NUMBER \ --lr $LEARNING_RATE --lr_decay_step $DECAY_STEP \ --cuda ``` Above, BATCH_SIZE and WORKER_NUMBER can be set adaptively according to your GPU memory size. **On Titan Xp with 12G memory, it can be up to 4**. If you have multiple (say 8) Titan Xp GPUs, then just use them all! Try: ``` python trainval_net.py --dataset pascal_voc --net vgg16 \ --bs 24 --nw 8 \ --lr $LEARNING_RATE --lr_decay_step $DECAY_STEP \ --cuda --mGPUs ``` Change dataset to "coco" or 'vg' if you want to train on COCO or Visual Genome. ## Test If you want to evaluate the detection performance of a pre-trained vgg16 model on pascal_voc test set, simply run ``` python test_net.py --dataset pascal_voc --net vgg16 \ --checksession $SESSION --checkepoch $EPOCH --checkpoint $CHECKPOINT \ --cuda ``` Specify the specific model session, checkepoch and checkpoint, e.g., SESSION=1, EPOCH=6, CHECKPOINT=416. ## Demo If you want to run detection on your own images with a pre-trained model, download the pretrained model listed in above tables or train your own models at first, then add images to folder $ROOT/images, and then run ``` python demo.py --net vgg16 \ --checksession $SESSION --checkepoch $EPOCH --checkpoint $CHECKPOINT \ --cuda --load_dir path/to/model/directoy ``` Then you will find the detection results in folder $ROOT/images. **Note the default demo.py merely support pascal_voc categories. You need to change the [line](https://github.com/jwyang/faster-rcnn.pytorch/blob/530f3fdccaa60d05fa068bc2148695211586bd88/demo.py#L156) to adapt your own model.** Below are some detection results:
## Webcam Demo You can use a webcam in a real-time demo by running ``` python demo.py --net vgg16 \ --checksession $SESSION --checkepoch $EPOCH --checkpoint $CHECKPOINT \ --cuda --load_dir path/to/model/directoy \ --webcam $WEBCAM_ID ``` The demo is stopped by clicking the image window and then pressing the 'q' key. ## Authorship This project is equally contributed by [Jianwei Yang](https://github.com/jwyang) and [Jiasen Lu](https://github.com/jiasenlu), and many others (thanks to them!). ## Citation @article{jjfaster2rcnn, Author = {Jianwei Yang and Jiasen Lu and Dhruv Batra and Devi Parikh}, Title = {A Faster Pytorch Implementation of Faster R-CNN}, Journal = {https://github.com/jwyang/faster-rcnn.pytorch}, Year = {2017} } @inproceedings{renNIPS15fasterrcnn, Author = {Shaoqing Ren and Kaiming He and Ross Girshick and Jian Sun}, Title = {Faster {R-CNN}: Towards Real-Time Object Detection with Region Proposal Networks}, Booktitle = {Advances in Neural Information Processing Systems ({NIPS})}, Year = {2015} } ================================================ FILE: _init_paths.py ================================================ import os.path as osp import sys def add_path(path): if path not in sys.path: sys.path.insert(0, path) this_dir = osp.dirname(__file__) # Add lib to PYTHONPATH lib_path = osp.join(this_dir, 'lib') add_path(lib_path) coco_path = osp.join(this_dir, 'data', 'coco', 'PythonAPI') add_path(coco_path) ================================================ FILE: cfgs/res101.yml ================================================ EXP_DIR: res101 TRAIN: HAS_RPN: True BBOX_NORMALIZE_TARGETS_PRECOMPUTED: True RPN_POSITIVE_OVERLAP: 0.7 RPN_BATCHSIZE: 256 PROPOSAL_METHOD: gt BG_THRESH_LO: 0.0 DISPLAY: 20 BATCH_SIZE: 128 WEIGHT_DECAY: 0.0001 DOUBLE_BIAS: False LEARNING_RATE: 0.001 TEST: HAS_RPN: True POOLING_SIZE: 7 POOLING_MODE: align CROP_RESIZE_WITH_MAX_POOL: False ================================================ FILE: cfgs/res101_ls.yml ================================================ EXP_DIR: res101 TRAIN: HAS_RPN: True BBOX_NORMALIZE_TARGETS_PRECOMPUTED: True RPN_POSITIVE_OVERLAP: 0.7 RPN_BATCHSIZE: 256 PROPOSAL_METHOD: gt BG_THRESH_LO: 0.0 DISPLAY: 20 BATCH_SIZE: 128 WEIGHT_DECAY: 0.0001 SCALES: [800] DOUBLE_BIAS: False LEARNING_RATE: 0.001 TEST: HAS_RPN: True SCALES: [800] MAX_SIZE: 1200 RPN_POST_NMS_TOP_N: 1000 POOLING_SIZE: 7 POOLING_MODE: align CROP_RESIZE_WITH_MAX_POOL: False ================================================ FILE: cfgs/res50.yml ================================================ EXP_DIR: res50 TRAIN: HAS_RPN: True # IMS_PER_BATCH: 1 BBOX_NORMALIZE_TARGETS_PRECOMPUTED: True RPN_POSITIVE_OVERLAP: 0.7 RPN_BATCHSIZE: 256 PROPOSAL_METHOD: gt BG_THRESH_LO: 0.0 DISPLAY: 20 BATCH_SIZE: 256 WEIGHT_DECAY: 0.0001 DOUBLE_BIAS: False SNAPSHOT_PREFIX: res50_faster_rcnn TEST: HAS_RPN: True POOLING_MODE: crop ================================================ FILE: cfgs/vgg16.yml ================================================ EXP_DIR: vgg16 TRAIN: HAS_RPN: True BBOX_NORMALIZE_TARGETS_PRECOMPUTED: True RPN_POSITIVE_OVERLAP: 0.7 RPN_BATCHSIZE: 256 PROPOSAL_METHOD: gt BG_THRESH_LO: 0.0 BATCH_SIZE: 256 LEARNING_RATE: 0.01 TEST: HAS_RPN: True POOLING_MODE: align CROP_RESIZE_WITH_MAX_POOL: False ================================================ FILE: demo.py ================================================ # -------------------------------------------------------- # Tensorflow Faster R-CNN # Licensed under The MIT License [see LICENSE for details] # Written by Jiasen Lu, Jianwei Yang, based on code from Ross Girshick # -------------------------------------------------------- from __future__ import absolute_import from __future__ import division from __future__ import print_function import _init_paths import os import sys import numpy as np import argparse import pprint import pdb import time import cv2 import imutils import torch from torch.autograd import Variable import torch.nn as nn import torch.optim as optim import torchvision.transforms as transforms import torchvision.datasets as dset from scipy.misc import imread from roi_data_layer.roidb import combined_roidb from roi_data_layer.roibatchLoader import roibatchLoader from model.utils.config import cfg, cfg_from_file, cfg_from_list, get_output_dir from model.rpn.bbox_transform import clip_boxes from model.nms.nms_wrapper import nms from model.rpn.bbox_transform import bbox_transform_inv from model.utils.net_utils import save_net, load_net, vis_detections from model.utils.blob import im_list_to_blob from model.faster_rcnn.vgg16 import vgg16 from model.faster_rcnn.resnet import resnet import pdb try: xrange # Python 2 except NameError: xrange = range # Python 3 def parse_args(): """ Parse input arguments """ parser = argparse.ArgumentParser(description='Train a Fast R-CNN network') parser.add_argument('--dataset', dest='dataset', help='training dataset', default='pascal_voc', type=str) parser.add_argument('--cfg', dest='cfg_file', help='optional config file', default='cfgs/vgg16.yml', type=str) parser.add_argument('--net', dest='net', help='vgg16, res50, res101, res152', default='res101', type=str) parser.add_argument('--set', dest='set_cfgs', help='set config keys', default=None, nargs=argparse.REMAINDER) parser.add_argument('--load_dir', dest='load_dir', help='directory to load models', default="/srv/share/jyang375/models") parser.add_argument('--image_dir', dest='image_dir', help='directory to load images for demo', default="images") parser.add_argument('--cuda', dest='cuda', help='whether use CUDA', action='store_true') parser.add_argument('--mGPUs', dest='mGPUs', help='whether use multiple GPUs', action='store_true') parser.add_argument('--cag', dest='class_agnostic', help='whether perform class_agnostic bbox regression', action='store_true') parser.add_argument('--parallel_type', dest='parallel_type', help='which part of model to parallel, 0: all, 1: model before roi pooling', default=0, type=int) parser.add_argument('--checksession', dest='checksession', help='checksession to load model', default=1, type=int) parser.add_argument('--checkepoch', dest='checkepoch', help='checkepoch to load network', default=1, type=int) parser.add_argument('--checkpoint', dest='checkpoint', help='checkpoint to load network', default=10021, type=int) parser.add_argument('--bs', dest='batch_size', help='batch_size', default=1, type=int) parser.add_argument('--vis', dest='vis', help='visualization mode', action='store_true') parser.add_argument('--webcam_num', dest='webcam_num', help='webcam ID number', default=-1, type=int) args = parser.parse_args() return args lr = cfg.TRAIN.LEARNING_RATE momentum = cfg.TRAIN.MOMENTUM weight_decay = cfg.TRAIN.WEIGHT_DECAY def _get_image_blob(im): """Converts an image into a network input. Arguments: im (ndarray): a color image in BGR order Returns: blob (ndarray): a data blob holding an image pyramid im_scale_factors (list): list of image scales (relative to im) used in the image pyramid """ im_orig = im.astype(np.float32, copy=True) im_orig -= cfg.PIXEL_MEANS im_shape = im_orig.shape im_size_min = np.min(im_shape[0:2]) im_size_max = np.max(im_shape[0:2]) processed_ims = [] im_scale_factors = [] for target_size in cfg.TEST.SCALES: im_scale = float(target_size) / float(im_size_min) # Prevent the biggest axis from being more than MAX_SIZE if np.round(im_scale * im_size_max) > cfg.TEST.MAX_SIZE: im_scale = float(cfg.TEST.MAX_SIZE) / float(im_size_max) im = cv2.resize(im_orig, None, None, fx=im_scale, fy=im_scale, interpolation=cv2.INTER_LINEAR) im_scale_factors.append(im_scale) processed_ims.append(im) # Create a blob to hold the input images blob = im_list_to_blob(processed_ims) return blob, np.array(im_scale_factors) if __name__ == '__main__': args = parse_args() print('Called with args:') print(args) if args.cfg_file is not None: cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) cfg.USE_GPU_NMS = args.cuda print('Using config:') pprint.pprint(cfg) np.random.seed(cfg.RNG_SEED) # train set # -- Note: Use validation set and disable the flipped to enable faster loading. input_dir = args.load_dir + "/" + args.net + "/" + args.dataset if not os.path.exists(input_dir): raise Exception('There is no input directory for loading network from ' + input_dir) load_name = os.path.join(input_dir, 'faster_rcnn_{}_{}_{}.pth'.format(args.checksession, args.checkepoch, args.checkpoint)) pascal_classes = np.asarray(['__background__', 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor']) # initilize the network here. if args.net == 'vgg16': fasterRCNN = vgg16(pascal_classes, pretrained=False, class_agnostic=args.class_agnostic) elif args.net == 'res101': fasterRCNN = resnet(pascal_classes, 101, pretrained=False, class_agnostic=args.class_agnostic) elif args.net == 'res50': fasterRCNN = resnet(pascal_classes, 50, pretrained=False, class_agnostic=args.class_agnostic) elif args.net == 'res152': fasterRCNN = resnet(pascal_classes, 152, pretrained=False, class_agnostic=args.class_agnostic) else: print("network is not defined") pdb.set_trace() fasterRCNN.create_architecture() print("load checkpoint %s" % (load_name)) if args.cuda > 0: checkpoint = torch.load(load_name) else: checkpoint = torch.load(load_name, map_location=(lambda storage, loc: storage)) fasterRCNN.load_state_dict(checkpoint['model']) if 'pooling_mode' in checkpoint.keys(): cfg.POOLING_MODE = checkpoint['pooling_mode'] print('load model successfully!') # pdb.set_trace() print("load checkpoint %s" % (load_name)) # initilize the tensor holder here. im_data = torch.FloatTensor(1) im_info = torch.FloatTensor(1) num_boxes = torch.LongTensor(1) gt_boxes = torch.FloatTensor(1) # ship to cuda if args.cuda > 0: im_data = im_data.cuda() im_info = im_info.cuda() num_boxes = num_boxes.cuda() gt_boxes = gt_boxes.cuda() # make variable with torch.no_grad(): im_data = Variable(im_data) im_info = Variable(im_info) num_boxes = Variable(num_boxes) gt_boxes = Variable(gt_boxes) if args.cuda > 0: cfg.CUDA = True if args.cuda > 0: fasterRCNN.cuda() fasterRCNN.eval() start = time.time() max_per_image = 100 thresh = 0.05 vis = True webcam_num = args.webcam_num # Set up webcam or get image directories if webcam_num >= 0 : cap = cv2.VideoCapture(webcam_num) num_images = 0 else: imglist = os.listdir(args.image_dir) num_images = len(imglist) print('Loaded Photo: {} images.'.format(num_images)) while (num_images >= 0): total_tic = time.time() if webcam_num == -1: num_images -= 1 # Get image from the webcam if webcam_num >= 0: if not cap.isOpened(): raise RuntimeError("Webcam could not open. Please check connection.") ret, frame = cap.read() im_in = np.array(frame) # Load the demo image else: im_file = os.path.join(args.image_dir, imglist[num_images]) # im = cv2.imread(im_file) im_in = np.array(imread(im_file)) if len(im_in.shape) == 2: im_in = im_in[:,:,np.newaxis] im_in = np.concatenate((im_in,im_in,im_in), axis=2) # rgb -> bgr im_in = im_in[:,:,::-1] im = im_in blobs, im_scales = _get_image_blob(im) assert len(im_scales) == 1, "Only single-image batch implemented" im_blob = blobs im_info_np = np.array([[im_blob.shape[1], im_blob.shape[2], im_scales[0]]], dtype=np.float32) im_data_pt = torch.from_numpy(im_blob) im_data_pt = im_data_pt.permute(0, 3, 1, 2) im_info_pt = torch.from_numpy(im_info_np) im_data.data.resize_(im_data_pt.size()).copy_(im_data_pt) im_info.data.resize_(im_info_pt.size()).copy_(im_info_pt) gt_boxes.data.resize_(1, 1, 5).zero_() num_boxes.data.resize_(1).zero_() # pdb.set_trace() det_tic = time.time() rois, cls_prob, bbox_pred, \ rpn_loss_cls, rpn_loss_box, \ RCNN_loss_cls, RCNN_loss_bbox, \ rois_label = fasterRCNN(im_data, im_info, gt_boxes, num_boxes) scores = cls_prob.data boxes = rois.data[:, :, 1:5] if cfg.TEST.BBOX_REG: # Apply bounding-box regression deltas box_deltas = bbox_pred.data if cfg.TRAIN.BBOX_NORMALIZE_TARGETS_PRECOMPUTED: # Optionally normalize targets by a precomputed mean and stdev if args.class_agnostic: if args.cuda > 0: box_deltas = box_deltas.view(-1, 4) * torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_STDS).cuda() \ + torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_MEANS).cuda() else: box_deltas = box_deltas.view(-1, 4) * torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_STDS) \ + torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_MEANS) box_deltas = box_deltas.view(1, -1, 4) else: if args.cuda > 0: box_deltas = box_deltas.view(-1, 4) * torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_STDS).cuda() \ + torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_MEANS).cuda() else: box_deltas = box_deltas.view(-1, 4) * torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_STDS) \ + torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_MEANS) box_deltas = box_deltas.view(1, -1, 4 * len(pascal_classes)) pred_boxes = bbox_transform_inv(boxes, box_deltas, 1) pred_boxes = clip_boxes(pred_boxes, im_info.data, 1) else: # Simply repeat the boxes, once for each class _ = torch.from_numpy(np.tile(boxes, (1, scores.shape[1]))) pred_boxes = _.cuda() if args.cuda > 0 else _ pred_boxes /= im_scales[0] scores = scores.squeeze() pred_boxes = pred_boxes.squeeze() det_toc = time.time() detect_time = det_toc - det_tic misc_tic = time.time() if vis: im2show = np.copy(im) for j in xrange(1, len(pascal_classes)): inds = torch.nonzero(scores[:,j]>thresh).view(-1) # if there is det if inds.numel() > 0: cls_scores = scores[:,j][inds] _, order = torch.sort(cls_scores, 0, True) if args.class_agnostic: cls_boxes = pred_boxes[inds, :] else: cls_boxes = pred_boxes[inds][:, j * 4:(j + 1) * 4] cls_dets = torch.cat((cls_boxes, cls_scores.unsqueeze(1)), 1) # cls_dets = torch.cat((cls_boxes, cls_scores), 1) cls_dets = cls_dets[order] keep = nms(cls_dets, cfg.TEST.NMS, force_cpu=not cfg.USE_GPU_NMS) cls_dets = cls_dets[keep.view(-1).long()] if vis: im2show = vis_detections(im2show, pascal_classes[j], cls_dets.cpu().numpy(), 0.5) misc_toc = time.time() nms_time = misc_toc - misc_tic if webcam_num == -1: sys.stdout.write('im_detect: {:d}/{:d} {:.3f}s {:.3f}s \r' \ .format(num_images + 1, len(imglist), detect_time, nms_time)) sys.stdout.flush() if vis and webcam_num == -1: # cv2.imshow('test', im2show) # cv2.waitKey(0) result_path = os.path.join(args.image_dir, imglist[num_images][:-4] + "_det.jpg") cv2.imwrite(result_path, im2show) else: cv2.imshow("frame", im2show) total_toc = time.time() total_time = total_toc - total_tic frame_rate = 1 / total_time print('Frame rate:', frame_rate) if cv2.waitKey(1) & 0xFF == ord('q'): break if webcam_num >= 0: cap.release() cv2.destroyAllWindows() ================================================ FILE: lib/datasets/VOCdevkit-matlab-wrapper/get_voc_opts.m ================================================ function VOCopts = get_voc_opts(path) tmp = pwd; cd(path); try addpath('VOCcode'); VOCinit; catch rmpath('VOCcode'); cd(tmp); error(sprintf('VOCcode directory not found under %s', path)); end rmpath('VOCcode'); cd(tmp); ================================================ FILE: lib/datasets/VOCdevkit-matlab-wrapper/voc_eval.m ================================================ function res = voc_eval(path, comp_id, test_set, output_dir) VOCopts = get_voc_opts(path); VOCopts.testset = test_set; for i = 1:length(VOCopts.classes) cls = VOCopts.classes{i}; res(i) = voc_eval_cls(cls, VOCopts, comp_id, output_dir); end fprintf('\n~~~~~~~~~~~~~~~~~~~~\n'); fprintf('Results:\n'); aps = [res(:).ap]'; fprintf('%.1f\n', aps * 100); fprintf('%.1f\n', mean(aps) * 100); fprintf('~~~~~~~~~~~~~~~~~~~~\n'); function res = voc_eval_cls(cls, VOCopts, comp_id, output_dir) test_set = VOCopts.testset; year = VOCopts.dataset(4:end); addpath(fullfile(VOCopts.datadir, 'VOCcode')); res_fn = sprintf(VOCopts.detrespath, comp_id, cls); recall = []; prec = []; ap = 0; ap_auc = 0; do_eval = (str2num(year) <= 2007) | ~strcmp(test_set, 'test'); if do_eval % Bug in VOCevaldet requires that tic has been called first tic; [recall, prec, ap] = VOCevaldet(VOCopts, comp_id, cls, true); ap_auc = xVOCap(recall, prec); % force plot limits ylim([0 1]); xlim([0 1]); print(gcf, '-djpeg', '-r0', ... [output_dir '/' cls '_pr.jpg']); end fprintf('!!! %s : %.4f %.4f\n', cls, ap, ap_auc); res.recall = recall; res.prec = prec; res.ap = ap; res.ap_auc = ap_auc; save([output_dir '/' cls '_pr.mat'], ... 'res', 'recall', 'prec', 'ap', 'ap_auc'); rmpath(fullfile(VOCopts.datadir, 'VOCcode')); ================================================ FILE: lib/datasets/VOCdevkit-matlab-wrapper/xVOCap.m ================================================ function ap = xVOCap(rec,prec) % From the PASCAL VOC 2011 devkit mrec=[0 ; rec ; 1]; mpre=[0 ; prec ; 0]; for i=numel(mpre)-1:-1:1 mpre(i)=max(mpre(i),mpre(i+1)); end i=find(mrec(2:end)~=mrec(1:end-1))+1; ap=sum((mrec(i)-mrec(i-1)).*mpre(i)); ================================================ FILE: lib/datasets/__init__.py ================================================ # -------------------------------------------------------- # Fast R-CNN # Copyright (c) 2015 Microsoft # Licensed under The MIT License [see LICENSE for details] # Written by Ross Girshick # -------------------------------------------------------- ================================================ FILE: lib/datasets/coco.py ================================================ # -------------------------------------------------------- # Fast/er R-CNN # Licensed under The MIT License [see LICENSE for details] # Written by Ross Girshick and Xinlei Chen # -------------------------------------------------------- from __future__ import absolute_import from __future__ import division from __future__ import print_function from datasets.imdb import imdb import datasets.ds_utils as ds_utils from model.utils.config import cfg import os.path as osp import sys import os import numpy as np import scipy.sparse import scipy.io as sio import pickle import json import uuid # COCO API from pycocotools.coco import COCO from pycocotools.cocoeval import COCOeval from pycocotools import mask as COCOmask class coco(imdb): def __init__(self, image_set, year): imdb.__init__(self, 'coco_' + year + '_' + image_set) # COCO specific config options self.config = {'use_salt': True, 'cleanup': True} # name, paths self._year = year self._image_set = image_set self._data_path = osp.join(cfg.DATA_DIR, 'coco') # load COCO API, classes, class <-> id mappings self._COCO = COCO(self._get_ann_file()) cats = self._COCO.loadCats(self._COCO.getCatIds()) self._classes = tuple(['__background__'] + [c['name'] for c in cats]) self._class_to_ind = dict(list(zip(self.classes, list(range(self.num_classes))))) self._class_to_coco_cat_id = dict(list(zip([c['name'] for c in cats], self._COCO.getCatIds()))) self._image_index = self._load_image_set_index() # Default to roidb handler self.set_proposal_method('gt') self.competition_mode(False) # Some image sets are "views" (i.e. subsets) into others. # For example, minival2014 is a random 5000 image subset of val2014. # This mapping tells us where the view's images and proposals come from. self._view_map = { 'minival2014': 'val2014', # 5k val2014 subset 'valminusminival2014': 'val2014', # val2014 \setminus minival2014 'test-dev2015': 'test2015', 'valminuscapval2014': 'val2014', 'capval2014': 'val2014', 'captest2014': 'val2014' } coco_name = image_set + year # e.g., "val2014" self._data_name = (self._view_map[coco_name] if coco_name in self._view_map else coco_name) # Dataset splits that have ground-truth annotations (test splits # do not have gt annotations) self._gt_splits = ('train', 'val', 'minival') def _get_ann_file(self): prefix = 'instances' if self._image_set.find('test') == -1 \ else 'image_info' return osp.join(self._data_path, 'annotations', prefix + '_' + self._image_set + self._year + '.json') def _load_image_set_index(self): """ Load image ids. """ image_ids = self._COCO.getImgIds() return image_ids def _get_widths(self): anns = self._COCO.loadImgs(self._image_index) widths = [ann['width'] for ann in anns] return widths def image_path_at(self, i): """ Return the absolute path to image i in the image sequence. """ return self.image_path_from_index(self._image_index[i]) def image_id_at(self, i): """ Return the absolute path to image i in the image sequence. """ return self._image_index[i] def image_path_from_index(self, index): """ Construct an image path from the image's "index" identifier. """ # Example image path for index=119993: # images/train2014/COCO_train2014_000000119993.jpg file_name = ('COCO_' + self._data_name + '_' + str(index).zfill(12) + '.jpg') image_path = osp.join(self._data_path, 'images', self._data_name, file_name) assert osp.exists(image_path), \ 'Path does not exist: {}'.format(image_path) return image_path def gt_roidb(self): """ Return the database of ground-truth regions of interest. This function loads/saves from/to a cache file to speed up future calls. """ cache_file = osp.join(self.cache_path, self.name + '_gt_roidb.pkl') if osp.exists(cache_file): with open(cache_file, 'rb') as fid: roidb = pickle.load(fid) print('{} gt roidb loaded from {}'.format(self.name, cache_file)) return roidb gt_roidb = [self._load_coco_annotation(index) for index in self._image_index] with open(cache_file, 'wb') as fid: pickle.dump(gt_roidb, fid, pickle.HIGHEST_PROTOCOL) print('wrote gt roidb to {}'.format(cache_file)) return gt_roidb def _load_coco_annotation(self, index): """ Loads COCO bounding-box instance annotations. Crowd instances are handled by marking their overlaps (with all categories) to -1. This overlap value means that crowd "instances" are excluded from training. """ im_ann = self._COCO.loadImgs(index)[0] width = im_ann['width'] height = im_ann['height'] annIds = self._COCO.getAnnIds(imgIds=index, iscrowd=None) objs = self._COCO.loadAnns(annIds) # Sanitize bboxes -- some are invalid valid_objs = [] for obj in objs: x1 = np.max((0, obj['bbox'][0])) y1 = np.max((0, obj['bbox'][1])) x2 = np.min((width - 1, x1 + np.max((0, obj['bbox'][2] - 1)))) y2 = np.min((height - 1, y1 + np.max((0, obj['bbox'][3] - 1)))) if obj['area'] > 0 and x2 >= x1 and y2 >= y1: obj['clean_bbox'] = [x1, y1, x2, y2] valid_objs.append(obj) objs = valid_objs num_objs = len(objs) boxes = np.zeros((num_objs, 4), dtype=np.uint16) gt_classes = np.zeros((num_objs), dtype=np.int32) overlaps = np.zeros((num_objs, self.num_classes), dtype=np.float32) seg_areas = np.zeros((num_objs), dtype=np.float32) # Lookup table to map from COCO category ids to our internal class # indices coco_cat_id_to_class_ind = dict([(self._class_to_coco_cat_id[cls], self._class_to_ind[cls]) for cls in self._classes[1:]]) for ix, obj in enumerate(objs): cls = coco_cat_id_to_class_ind[obj['category_id']] boxes[ix, :] = obj['clean_bbox'] gt_classes[ix] = cls seg_areas[ix] = obj['area'] if obj['iscrowd']: # Set overlap to -1 for all classes for crowd objects # so they will be excluded during training overlaps[ix, :] = -1.0 else: overlaps[ix, cls] = 1.0 ds_utils.validate_boxes(boxes, width=width, height=height) overlaps = scipy.sparse.csr_matrix(overlaps) return {'width': width, 'height': height, 'boxes': boxes, 'gt_classes': gt_classes, 'gt_overlaps': overlaps, 'flipped': False, 'seg_areas': seg_areas} def _get_widths(self): return [r['width'] for r in self.roidb] def append_flipped_images(self): num_images = self.num_images widths = self._get_widths() for i in range(num_images): boxes = self.roidb[i]['boxes'].copy() oldx1 = boxes[:, 0].copy() oldx2 = boxes[:, 2].copy() boxes[:, 0] = widths[i] - oldx2 - 1 boxes[:, 2] = widths[i] - oldx1 - 1 assert (boxes[:, 2] >= boxes[:, 0]).all() entry = {'width': widths[i], 'height': self.roidb[i]['height'], 'boxes': boxes, 'gt_classes': self.roidb[i]['gt_classes'], 'gt_overlaps': self.roidb[i]['gt_overlaps'], 'flipped': True, 'seg_areas': self.roidb[i]['seg_areas']} self.roidb.append(entry) self._image_index = self._image_index * 2 def _get_box_file(self, index): # first 14 chars / first 22 chars / all chars + .mat # COCO_val2014_0/COCO_val2014_000000447/COCO_val2014_000000447991.mat file_name = ('COCO_' + self._data_name + '_' + str(index).zfill(12) + '.mat') return osp.join(file_name[:14], file_name[:22], file_name) def _print_detection_eval_metrics(self, coco_eval): IoU_lo_thresh = 0.5 IoU_hi_thresh = 0.95 def _get_thr_ind(coco_eval, thr): ind = np.where((coco_eval.params.iouThrs > thr - 1e-5) & (coco_eval.params.iouThrs < thr + 1e-5))[0][0] iou_thr = coco_eval.params.iouThrs[ind] assert np.isclose(iou_thr, thr) return ind ind_lo = _get_thr_ind(coco_eval, IoU_lo_thresh) ind_hi = _get_thr_ind(coco_eval, IoU_hi_thresh) # precision has dims (iou, recall, cls, area range, max dets) # area range index 0: all area ranges # max dets index 2: 100 per image precision = \ coco_eval.eval['precision'][ind_lo:(ind_hi + 1), :, :, 0, 2] ap_default = np.mean(precision[precision > -1]) print(('~~~~ Mean and per-category AP @ IoU=[{:.2f},{:.2f}] ' '~~~~').format(IoU_lo_thresh, IoU_hi_thresh)) print('{:.1f}'.format(100 * ap_default)) for cls_ind, cls in enumerate(self.classes): if cls == '__background__': continue # minus 1 because of __background__ precision = coco_eval.eval['precision'][ind_lo:(ind_hi + 1), :, cls_ind - 1, 0, 2] ap = np.mean(precision[precision > -1]) print('{:.1f}'.format(100 * ap)) print('~~~~ Summary metrics ~~~~') coco_eval.summarize() def _do_detection_eval(self, res_file, output_dir): ann_type = 'bbox' coco_dt = self._COCO.loadRes(res_file) coco_eval = COCOeval(self._COCO, coco_dt) coco_eval.params.useSegm = (ann_type == 'segm') coco_eval.evaluate() coco_eval.accumulate() self._print_detection_eval_metrics(coco_eval) eval_file = osp.join(output_dir, 'detection_results.pkl') with open(eval_file, 'wb') as fid: pickle.dump(coco_eval, fid, pickle.HIGHEST_PROTOCOL) print('Wrote COCO eval results to: {}'.format(eval_file)) def _coco_results_one_category(self, boxes, cat_id): results = [] for im_ind, index in enumerate(self.image_index): dets = boxes[im_ind].astype(np.float) if dets == []: continue scores = dets[:, -1] xs = dets[:, 0] ys = dets[:, 1] ws = dets[:, 2] - xs + 1 hs = dets[:, 3] - ys + 1 results.extend( [{'image_id': index, 'category_id': cat_id, 'bbox': [xs[k], ys[k], ws[k], hs[k]], 'score': scores[k]} for k in range(dets.shape[0])]) return results def _write_coco_results_file(self, all_boxes, res_file): # [{"image_id": 42, # "category_id": 18, # "bbox": [258.15,41.29,348.26,243.78], # "score": 0.236}, ...] results = [] for cls_ind, cls in enumerate(self.classes): if cls == '__background__': continue print('Collecting {} results ({:d}/{:d})'.format(cls, cls_ind, self.num_classes - 1)) coco_cat_id = self._class_to_coco_cat_id[cls] results.extend(self._coco_results_one_category(all_boxes[cls_ind], coco_cat_id)) print('Writing results json to {}'.format(res_file)) with open(res_file, 'w') as fid: json.dump(results, fid) def evaluate_detections(self, all_boxes, output_dir): res_file = osp.join(output_dir, ('detections_' + self._image_set + self._year + '_results')) if self.config['use_salt']: res_file += '_{}'.format(str(uuid.uuid4())) res_file += '.json' self._write_coco_results_file(all_boxes, res_file) # Only do evaluation on non-test sets if self._image_set.find('test') == -1: self._do_detection_eval(res_file, output_dir) # Optionally cleanup results json file if self.config['cleanup']: os.remove(res_file) def competition_mode(self, on): if on: self.config['use_salt'] = False self.config['cleanup'] = False else: self.config['use_salt'] = True self.config['cleanup'] = True ================================================ FILE: lib/datasets/ds_utils.py ================================================ # -------------------------------------------------------- # Fast/er R-CNN # Licensed under The MIT License [see LICENSE for details] # Written by Ross Girshick # -------------------------------------------------------- from __future__ import absolute_import from __future__ import division from __future__ import print_function import numpy as np def unique_boxes(boxes, scale=1.0): """Return indices of unique boxes.""" v = np.array([1, 1e3, 1e6, 1e9]) hashes = np.round(boxes * scale).dot(v) _, index = np.unique(hashes, return_index=True) return np.sort(index) def xywh_to_xyxy(boxes): """Convert [x y w h] box format to [x1 y1 x2 y2] format.""" return np.hstack((boxes[:, 0:2], boxes[:, 0:2] + boxes[:, 2:4] - 1)) def xyxy_to_xywh(boxes): """Convert [x1 y1 x2 y2] box format to [x y w h] format.""" return np.hstack((boxes[:, 0:2], boxes[:, 2:4] - boxes[:, 0:2] + 1)) def validate_boxes(boxes, width=0, height=0): """Check that a set of boxes are valid.""" x1 = boxes[:, 0] y1 = boxes[:, 1] x2 = boxes[:, 2] y2 = boxes[:, 3] assert (x1 >= 0).all() assert (y1 >= 0).all() assert (x2 >= x1).all() assert (y2 >= y1).all() assert (x2 < width).all() assert (y2 < height).all() def filter_small_boxes(boxes, min_size): w = boxes[:, 2] - boxes[:, 0] h = boxes[:, 3] - boxes[:, 1] keep = np.where((w >= min_size) & (h > min_size))[0] return keep ================================================ FILE: lib/datasets/factory.py ================================================ # -------------------------------------------------------- # Fast R-CNN # Copyright (c) 2015 Microsoft # Licensed under The MIT License [see LICENSE for details] # Written by Ross Girshick # -------------------------------------------------------- """Factory method for easily getting imdbs by name.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function __sets = {} from datasets.pascal_voc import pascal_voc from datasets.coco import coco from datasets.imagenet import imagenet from datasets.vg import vg import numpy as np # Set up voc__ for year in ['2007', '2012']: for split in ['train', 'val', 'trainval', 'test']: name = 'voc_{}_{}'.format(year, split) __sets[name] = (lambda split=split, year=year: pascal_voc(split, year)) # Set up coco_2014_ for year in ['2014']: for split in ['train', 'val', 'minival', 'valminusminival', 'trainval']: name = 'coco_{}_{}'.format(year, split) __sets[name] = (lambda split=split, year=year: coco(split, year)) # Set up coco_2014_cap_ for year in ['2014']: for split in ['train', 'val', 'capval', 'valminuscapval', 'trainval']: name = 'coco_{}_{}'.format(year, split) __sets[name] = (lambda split=split, year=year: coco(split, year)) # Set up coco_2015_ for year in ['2015']: for split in ['test', 'test-dev']: name = 'coco_{}_{}'.format(year, split) __sets[name] = (lambda split=split, year=year: coco(split, year)) # Set up vg_ # for version in ['1600-400-20']: # for split in ['minitrain', 'train', 'minival', 'val', 'test']: # name = 'vg_{}_{}'.format(version,split) # __sets[name] = (lambda split=split, version=version: vg(version, split)) for version in ['150-50-20', '150-50-50', '500-150-80', '750-250-150', '1750-700-450', '1600-400-20']: for split in ['minitrain', 'smalltrain', 'train', 'minival', 'smallval', 'val', 'test']: name = 'vg_{}_{}'.format(version,split) __sets[name] = (lambda split=split, version=version: vg(version, split)) # set up image net. for split in ['train', 'val', 'val1', 'val2', 'test']: name = 'imagenet_{}'.format(split) devkit_path = 'data/imagenet/ILSVRC/devkit' data_path = 'data/imagenet/ILSVRC' __sets[name] = (lambda split=split, devkit_path=devkit_path, data_path=data_path: imagenet(split,devkit_path,data_path)) def get_imdb(name): """Get an imdb (image database) by name.""" if name not in __sets: raise KeyError('Unknown dataset: {}'.format(name)) return __sets[name]() def list_imdbs(): """List all registered imdbs.""" return list(__sets.keys()) ================================================ FILE: lib/datasets/imagenet.py ================================================ from __future__ import print_function # -------------------------------------------------------- # Fast R-CNN # Copyright (c) 2015 Microsoft # Licensed under The MIT License [see LICENSE for details] # Written by Ross Girshick # -------------------------------------------------------- import datasets import datasets.imagenet import os, sys from datasets.imdb import imdb import xml.dom.minidom as minidom import numpy as np import scipy.sparse import scipy.io as sio import subprocess import pdb import pickle try: xrange # Python 2 except NameError: xrange = range # Python 3 class imagenet(imdb): def __init__(self, image_set, devkit_path, data_path): imdb.__init__(self, image_set) self._image_set = image_set self._devkit_path = devkit_path self._data_path = data_path synsets_image = sio.loadmat(os.path.join(self._devkit_path, 'data', 'meta_det.mat')) synsets_video = sio.loadmat(os.path.join(self._devkit_path, 'data', 'meta_vid.mat')) self._classes_image = ('__background__',) self._wnid_image = (0,) self._classes = ('__background__',) self._wnid = (0,) for i in xrange(200): self._classes_image = self._classes_image + (synsets_image['synsets'][0][i][2][0],) self._wnid_image = self._wnid_image + (synsets_image['synsets'][0][i][1][0],) for i in xrange(30): self._classes = self._classes + (synsets_video['synsets'][0][i][2][0],) self._wnid = self._wnid + (synsets_video['synsets'][0][i][1][0],) self._wnid_to_ind_image = dict(zip(self._wnid_image, xrange(201))) self._class_to_ind_image = dict(zip(self._classes_image, xrange(201))) self._wnid_to_ind = dict(zip(self._wnid, xrange(31))) self._class_to_ind = dict(zip(self._classes, xrange(31))) #check for valid intersection between video and image classes self._valid_image_flag = [0]*201 for i in range(1,201): if self._wnid_image[i] in self._wnid_to_ind: self._valid_image_flag[i] = 1 self._image_ext = ['.JPEG'] self._image_index = self._load_image_set_index() # Default to roidb handler self._roidb_handler = self.gt_roidb # Specific config options self.config = {'cleanup' : True, 'use_salt' : True, 'top_k' : 2000} assert os.path.exists(self._devkit_path), 'Devkit path does not exist: {}'.format(self._devkit_path) assert os.path.exists(self._data_path), 'Path does not exist: {}'.format(self._data_path) def image_path_at(self, i): """ Return the absolute path to image i in the image sequence. """ return self.image_path_from_index(self._image_index[i]) def image_path_from_index(self, index): """ Construct an image path from the image's "index" identifier. """ image_path = os.path.join(self._data_path, 'Data', self._image_set, index + self._image_ext[0]) assert os.path.exists(image_path), 'path does not exist: {}'.format(image_path) return image_path def _load_image_set_index(self): """ Load the indexes listed in this dataset's image set file. """ # Example path to image set file: # self._data_path + /ImageSets/val.txt if self._image_set == 'train': image_set_file = os.path.join(self._data_path, 'ImageSets', 'trainr.txt') image_index = [] if os.path.exists(image_set_file): f = open(image_set_file, 'r') data = f.read().split() for lines in data: if lines != '': image_index.append(lines) f.close() return image_index for i in range(1,200): print(i) image_set_file = os.path.join(self._data_path, 'ImageSets', 'DET', 'train_' + str(i) + '.txt') with open(image_set_file) as f: tmp_index = [x.strip() for x in f.readlines()] vtmp_index = [] for line in tmp_index: line = line.split(' ') image_list = os.popen('ls ' + self._data_path + '/Data/DET/train/' + line[0] + '/*.JPEG').read().split() tmp_list = [] for imgs in image_list: tmp_list.append(imgs[:-5]) vtmp_index = vtmp_index + tmp_list num_lines = len(vtmp_index) ids = np.random.permutation(num_lines) count = 0 while count < 2000: image_index.append(vtmp_index[ids[count % num_lines]]) count = count + 1 for i in range(1,201): if self._valid_image_flag[i] == 1: image_set_file = os.path.join(self._data_path, 'ImageSets', 'train_pos_' + str(i) + '.txt') with open(image_set_file) as f: tmp_index = [x.strip() for x in f.readlines()] num_lines = len(tmp_index) ids = np.random.permutation(num_lines) count = 0 while count < 2000: image_index.append(tmp_index[ids[count % num_lines]]) count = count + 1 image_set_file = os.path.join(self._data_path, 'ImageSets', 'trainr.txt') f = open(image_set_file, 'w') for lines in image_index: f.write(lines + '\n') f.close() else: image_set_file = os.path.join(self._data_path, 'ImageSets', 'val.txt') with open(image_set_file) as f: image_index = [x.strip() for x in f.readlines()] return image_index def gt_roidb(self): """ Return the database of ground-truth regions of interest. This function loads/saves from/to a cache file to speed up future calls. """ cache_file = os.path.join(self.cache_path, self.name + '_gt_roidb.pkl') if os.path.exists(cache_file): with open(cache_file, 'rb') as fid: roidb = pickle.load(fid) print('{} gt roidb loaded from {}'.format(self.name, cache_file)) return roidb gt_roidb = [self._load_imagenet_annotation(index) for index in self.image_index] with open(cache_file, 'wb') as fid: pickle.dump(gt_roidb, fid, pickle.HIGHEST_PROTOCOL) print('wrote gt roidb to {}'.format(cache_file)) return gt_roidb def _load_imagenet_annotation(self, index): """ Load image and bounding boxes info from txt files of imagenet. """ filename = os.path.join(self._data_path, 'Annotations', self._image_set, index + '.xml') # print 'Loading: {}'.format(filename) def get_data_from_tag(node, tag): return node.getElementsByTagName(tag)[0].childNodes[0].data with open(filename) as f: data = minidom.parseString(f.read()) objs = data.getElementsByTagName('object') num_objs = len(objs) boxes = np.zeros((num_objs, 4), dtype=np.uint16) gt_classes = np.zeros((num_objs), dtype=np.int32) overlaps = np.zeros((num_objs, self.num_classes), dtype=np.float32) # Load object bounding boxes into a data frame. for ix, obj in enumerate(objs): x1 = float(get_data_from_tag(obj, 'xmin')) y1 = float(get_data_from_tag(obj, 'ymin')) x2 = float(get_data_from_tag(obj, 'xmax')) y2 = float(get_data_from_tag(obj, 'ymax')) cls = self._wnid_to_ind[ str(get_data_from_tag(obj, "name")).lower().strip()] boxes[ix, :] = [x1, y1, x2, y2] gt_classes[ix] = cls overlaps[ix, cls] = 1.0 overlaps = scipy.sparse.csr_matrix(overlaps) return {'boxes' : boxes, 'gt_classes': gt_classes, 'gt_overlaps' : overlaps, 'flipped' : False} if __name__ == '__main__': d = datasets.imagenet('val', '') res = d.roidb from IPython import embed; embed() ================================================ FILE: lib/datasets/imdb.py ================================================ # -------------------------------------------------------- # Fast R-CNN # Copyright (c) 2015 Microsoft # Licensed under The MIT License [see LICENSE for details] # Written by Ross Girshick and Xinlei Chen # -------------------------------------------------------- from __future__ import absolute_import from __future__ import division from __future__ import print_function import os import os.path as osp import PIL from model.utils.cython_bbox import bbox_overlaps import numpy as np import scipy.sparse from model.utils.config import cfg import pdb ROOT_DIR = osp.join(osp.dirname(__file__), '..', '..') class imdb(object): """Image database.""" def __init__(self, name, classes=None): self._name = name self._num_classes = 0 if not classes: self._classes = [] else: self._classes = classes self._image_index = [] self._obj_proposer = 'gt' self._roidb = None self._roidb_handler = self.default_roidb # Use this dict for storing dataset specific config options self.config = {} @property def name(self): return self._name @property def num_classes(self): return len(self._classes) @property def classes(self): return self._classes @property def image_index(self): return self._image_index @property def roidb_handler(self): return self._roidb_handler @roidb_handler.setter def roidb_handler(self, val): self._roidb_handler = val def set_proposal_method(self, method): method = eval('self.' + method + '_roidb') self.roidb_handler = method @property def roidb(self): # A roidb is a list of dictionaries, each with the following keys: # boxes # gt_overlaps # gt_classes # flipped if self._roidb is not None: return self._roidb self._roidb = self.roidb_handler() return self._roidb @property def cache_path(self): cache_path = osp.abspath(osp.join(cfg.DATA_DIR, 'cache')) if not os.path.exists(cache_path): os.makedirs(cache_path) return cache_path @property def num_images(self): return len(self.image_index) def image_path_at(self, i): raise NotImplementedError def image_id_at(self, i): raise NotImplementedError def default_roidb(self): raise NotImplementedError def evaluate_detections(self, all_boxes, output_dir=None): """ all_boxes is a list of length number-of-classes. Each list element is a list of length number-of-images. Each of those list elements is either an empty list [] or a numpy array of detection. all_boxes[class][image] = [] or np.array of shape #dets x 5 """ raise NotImplementedError def _get_widths(self): return [PIL.Image.open(self.image_path_at(i)).size[0] for i in range(self.num_images)] def append_flipped_images(self): num_images = self.num_images widths = self._get_widths() for i in range(num_images): boxes = self.roidb[i]['boxes'].copy() oldx1 = boxes[:, 0].copy() oldx2 = boxes[:, 2].copy() boxes[:, 0] = widths[i] - oldx2 - 1 boxes[:, 2] = widths[i] - oldx1 - 1 assert (boxes[:, 2] >= boxes[:, 0]).all() entry = {'boxes': boxes, 'gt_overlaps': self.roidb[i]['gt_overlaps'], 'gt_classes': self.roidb[i]['gt_classes'], 'flipped': True} self.roidb.append(entry) self._image_index = self._image_index * 2 def evaluate_recall(self, candidate_boxes=None, thresholds=None, area='all', limit=None): """Evaluate detection proposal recall metrics. Returns: results: dictionary of results with keys 'ar': average recall 'recalls': vector recalls at each IoU overlap threshold 'thresholds': vector of IoU overlap thresholds 'gt_overlaps': vector of all ground-truth overlaps """ # Record max overlap value for each gt box # Return vector of overlap values areas = {'all': 0, 'small': 1, 'medium': 2, 'large': 3, '96-128': 4, '128-256': 5, '256-512': 6, '512-inf': 7} area_ranges = [[0 ** 2, 1e5 ** 2], # all [0 ** 2, 32 ** 2], # small [32 ** 2, 96 ** 2], # medium [96 ** 2, 1e5 ** 2], # large [96 ** 2, 128 ** 2], # 96-128 [128 ** 2, 256 ** 2], # 128-256 [256 ** 2, 512 ** 2], # 256-512 [512 ** 2, 1e5 ** 2], # 512-inf ] assert area in areas, 'unknown area range: {}'.format(area) area_range = area_ranges[areas[area]] gt_overlaps = np.zeros(0) num_pos = 0 for i in range(self.num_images): # Checking for max_overlaps == 1 avoids including crowd annotations # (...pretty hacking :/) max_gt_overlaps = self.roidb[i]['gt_overlaps'].toarray().max(axis=1) gt_inds = np.where((self.roidb[i]['gt_classes'] > 0) & (max_gt_overlaps == 1))[0] gt_boxes = self.roidb[i]['boxes'][gt_inds, :] gt_areas = self.roidb[i]['seg_areas'][gt_inds] valid_gt_inds = np.where((gt_areas >= area_range[0]) & (gt_areas <= area_range[1]))[0] gt_boxes = gt_boxes[valid_gt_inds, :] num_pos += len(valid_gt_inds) if candidate_boxes is None: # If candidate_boxes is not supplied, the default is to use the # non-ground-truth boxes from this roidb non_gt_inds = np.where(self.roidb[i]['gt_classes'] == 0)[0] boxes = self.roidb[i]['boxes'][non_gt_inds, :] else: boxes = candidate_boxes[i] if boxes.shape[0] == 0: continue if limit is not None and boxes.shape[0] > limit: boxes = boxes[:limit, :] overlaps = bbox_overlaps(boxes.astype(np.float), gt_boxes.astype(np.float)) _gt_overlaps = np.zeros((gt_boxes.shape[0])) for j in range(gt_boxes.shape[0]): # find which proposal box maximally covers each gt box argmax_overlaps = overlaps.argmax(axis=0) # and get the iou amount of coverage for each gt box max_overlaps = overlaps.max(axis=0) # find which gt box is 'best' covered (i.e. 'best' = most iou) gt_ind = max_overlaps.argmax() gt_ovr = max_overlaps.max() assert (gt_ovr >= 0) # find the proposal box that covers the best covered gt box box_ind = argmax_overlaps[gt_ind] # record the iou coverage of this gt box _gt_overlaps[j] = overlaps[box_ind, gt_ind] assert (_gt_overlaps[j] == gt_ovr) # mark the proposal box and the gt box as used overlaps[box_ind, :] = -1 overlaps[:, gt_ind] = -1 # append recorded iou coverage level gt_overlaps = np.hstack((gt_overlaps, _gt_overlaps)) gt_overlaps = np.sort(gt_overlaps) if thresholds is None: step = 0.05 thresholds = np.arange(0.5, 0.95 + 1e-5, step) recalls = np.zeros_like(thresholds) # compute recall for each iou threshold for i, t in enumerate(thresholds): recalls[i] = (gt_overlaps >= t).sum() / float(num_pos) # ar = 2 * np.trapz(recalls, thresholds) ar = recalls.mean() return {'ar': ar, 'recalls': recalls, 'thresholds': thresholds, 'gt_overlaps': gt_overlaps} def create_roidb_from_box_list(self, box_list, gt_roidb): assert len(box_list) == self.num_images, \ 'Number of boxes must match number of ground-truth images' roidb = [] for i in range(self.num_images): boxes = box_list[i] num_boxes = boxes.shape[0] overlaps = np.zeros((num_boxes, self.num_classes), dtype=np.float32) if gt_roidb is not None and gt_roidb[i]['boxes'].size > 0: gt_boxes = gt_roidb[i]['boxes'] gt_classes = gt_roidb[i]['gt_classes'] gt_overlaps = bbox_overlaps(boxes.astype(np.float), gt_boxes.astype(np.float)) argmaxes = gt_overlaps.argmax(axis=1) maxes = gt_overlaps.max(axis=1) I = np.where(maxes > 0)[0] overlaps[I, gt_classes[argmaxes[I]]] = maxes[I] overlaps = scipy.sparse.csr_matrix(overlaps) roidb.append({ 'boxes': boxes, 'gt_classes': np.zeros((num_boxes,), dtype=np.int32), 'gt_overlaps': overlaps, 'flipped': False, 'seg_areas': np.zeros((num_boxes,), dtype=np.float32), }) return roidb @staticmethod def merge_roidbs(a, b): assert len(a) == len(b) for i in range(len(a)): a[i]['boxes'] = np.vstack((a[i]['boxes'], b[i]['boxes'])) a[i]['gt_classes'] = np.hstack((a[i]['gt_classes'], b[i]['gt_classes'])) a[i]['gt_overlaps'] = scipy.sparse.vstack([a[i]['gt_overlaps'], b[i]['gt_overlaps']]) a[i]['seg_areas'] = np.hstack((a[i]['seg_areas'], b[i]['seg_areas'])) return a def competition_mode(self, on): """Turn competition mode on or off.""" pass ================================================ FILE: lib/datasets/pascal_voc.py ================================================ from __future__ import print_function from __future__ import absolute_import # -------------------------------------------------------- # Fast R-CNN # Copyright (c) 2015 Microsoft # Licensed under The MIT License [see LICENSE for details] # Written by Ross Girshick # -------------------------------------------------------- import xml.dom.minidom as minidom import os # import PIL import numpy as np import scipy.sparse import subprocess import math import glob import uuid import scipy.io as sio import xml.etree.ElementTree as ET import pickle from .imdb import imdb from .imdb import ROOT_DIR from . import ds_utils from .voc_eval import voc_eval # TODO: make fast_rcnn irrelevant # >>>> obsolete, because it depends on sth outside of this project from model.utils.config import cfg try: xrange # Python 2 except NameError: xrange = range # Python 3 # <<<< obsolete class pascal_voc(imdb): def __init__(self, image_set, year, devkit_path=None): imdb.__init__(self, 'voc_' + year + '_' + image_set) self._year = year self._image_set = image_set self._devkit_path = self._get_default_path() if devkit_path is None \ else devkit_path self._data_path = os.path.join(self._devkit_path, 'VOC' + self._year) self._classes = ('__background__', # always index 0 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor') self._class_to_ind = dict(zip(self.classes, xrange(self.num_classes))) self._image_ext = '.jpg' self._image_index = self._load_image_set_index() # Default to roidb handler # self._roidb_handler = self.selective_search_roidb self._roidb_handler = self.gt_roidb self._salt = str(uuid.uuid4()) self._comp_id = 'comp4' # PASCAL specific config options self.config = {'cleanup': True, 'use_salt': True, 'use_diff': False, 'matlab_eval': False, 'rpn_file': None, 'min_size': 2} assert os.path.exists(self._devkit_path), \ 'VOCdevkit path does not exist: {}'.format(self._devkit_path) assert os.path.exists(self._data_path), \ 'Path does not exist: {}'.format(self._data_path) def image_path_at(self, i): """ Return the absolute path to image i in the image sequence. """ return self.image_path_from_index(self._image_index[i]) def image_id_at(self, i): """ Return the absolute path to image i in the image sequence. """ return i def image_path_from_index(self, index): """ Construct an image path from the image's "index" identifier. """ image_path = os.path.join(self._data_path, 'JPEGImages', index + self._image_ext) assert os.path.exists(image_path), \ 'Path does not exist: {}'.format(image_path) return image_path def _load_image_set_index(self): """ Load the indexes listed in this dataset's image set file. """ # Example path to image set file: # self._devkit_path + /VOCdevkit2007/VOC2007/ImageSets/Main/val.txt image_set_file = os.path.join(self._data_path, 'ImageSets', 'Main', self._image_set + '.txt') assert os.path.exists(image_set_file), \ 'Path does not exist: {}'.format(image_set_file) with open(image_set_file) as f: image_index = [x.strip() for x in f.readlines()] return image_index def _get_default_path(self): """ Return the default path where PASCAL VOC is expected to be installed. """ return os.path.join(cfg.DATA_DIR, 'VOCdevkit' + self._year) def gt_roidb(self): """ Return the database of ground-truth regions of interest. This function loads/saves from/to a cache file to speed up future calls. """ cache_file = os.path.join(self.cache_path, self.name + '_gt_roidb.pkl') if os.path.exists(cache_file): with open(cache_file, 'rb') as fid: roidb = pickle.load(fid) print('{} gt roidb loaded from {}'.format(self.name, cache_file)) return roidb gt_roidb = [self._load_pascal_annotation(index) for index in self.image_index] with open(cache_file, 'wb') as fid: pickle.dump(gt_roidb, fid, pickle.HIGHEST_PROTOCOL) print('wrote gt roidb to {}'.format(cache_file)) return gt_roidb def selective_search_roidb(self): """ Return the database of selective search regions of interest. Ground-truth ROIs are also included. This function loads/saves from/to a cache file to speed up future calls. """ cache_file = os.path.join(self.cache_path, self.name + '_selective_search_roidb.pkl') if os.path.exists(cache_file): with open(cache_file, 'rb') as fid: roidb = pickle.load(fid) print('{} ss roidb loaded from {}'.format(self.name, cache_file)) return roidb if int(self._year) == 2007 or self._image_set != 'test': gt_roidb = self.gt_roidb() ss_roidb = self._load_selective_search_roidb(gt_roidb) roidb = imdb.merge_roidbs(gt_roidb, ss_roidb) else: roidb = self._load_selective_search_roidb(None) with open(cache_file, 'wb') as fid: pickle.dump(roidb, fid, pickle.HIGHEST_PROTOCOL) print('wrote ss roidb to {}'.format(cache_file)) return roidb def rpn_roidb(self): if int(self._year) == 2007 or self._image_set != 'test': gt_roidb = self.gt_roidb() rpn_roidb = self._load_rpn_roidb(gt_roidb) roidb = imdb.merge_roidbs(gt_roidb, rpn_roidb) else: roidb = self._load_rpn_roidb(None) return roidb def _load_rpn_roidb(self, gt_roidb): filename = self.config['rpn_file'] print('loading {}'.format(filename)) assert os.path.exists(filename), \ 'rpn data not found at: {}'.format(filename) with open(filename, 'rb') as f: box_list = pickle.load(f) return self.create_roidb_from_box_list(box_list, gt_roidb) def _load_selective_search_roidb(self, gt_roidb): filename = os.path.abspath(os.path.join(cfg.DATA_DIR, 'selective_search_data', self.name + '.mat')) assert os.path.exists(filename), \ 'Selective search data not found at: {}'.format(filename) raw_data = sio.loadmat(filename)['boxes'].ravel() box_list = [] for i in xrange(raw_data.shape[0]): boxes = raw_data[i][:, (1, 0, 3, 2)] - 1 keep = ds_utils.unique_boxes(boxes) boxes = boxes[keep, :] keep = ds_utils.filter_small_boxes(boxes, self.config['min_size']) boxes = boxes[keep, :] box_list.append(boxes) return self.create_roidb_from_box_list(box_list, gt_roidb) def _load_pascal_annotation(self, index): """ Load image and bounding boxes info from XML file in the PASCAL VOC format. """ filename = os.path.join(self._data_path, 'Annotations', index + '.xml') tree = ET.parse(filename) objs = tree.findall('object') # if not self.config['use_diff']: # # Exclude the samples labeled as difficult # non_diff_objs = [ # obj for obj in objs if int(obj.find('difficult').text) == 0] # # if len(non_diff_objs) != len(objs): # # print 'Removed {} difficult objects'.format( # # len(objs) - len(non_diff_objs)) # objs = non_diff_objs num_objs = len(objs) boxes = np.zeros((num_objs, 4), dtype=np.uint16) gt_classes = np.zeros((num_objs), dtype=np.int32) overlaps = np.zeros((num_objs, self.num_classes), dtype=np.float32) # "Seg" area for pascal is just the box area seg_areas = np.zeros((num_objs), dtype=np.float32) ishards = np.zeros((num_objs), dtype=np.int32) # Load object bounding boxes into a data frame. for ix, obj in enumerate(objs): bbox = obj.find('bndbox') # Make pixel indexes 0-based x1 = float(bbox.find('xmin').text) - 1 y1 = float(bbox.find('ymin').text) - 1 x2 = float(bbox.find('xmax').text) - 1 y2 = float(bbox.find('ymax').text) - 1 diffc = obj.find('difficult') difficult = 0 if diffc == None else int(diffc.text) ishards[ix] = difficult cls = self._class_to_ind[obj.find('name').text.lower().strip()] boxes[ix, :] = [x1, y1, x2, y2] gt_classes[ix] = cls overlaps[ix, cls] = 1.0 seg_areas[ix] = (x2 - x1 + 1) * (y2 - y1 + 1) overlaps = scipy.sparse.csr_matrix(overlaps) return {'boxes': boxes, 'gt_classes': gt_classes, 'gt_ishard': ishards, 'gt_overlaps': overlaps, 'flipped': False, 'seg_areas': seg_areas} def _get_comp_id(self): comp_id = (self._comp_id + '_' + self._salt if self.config['use_salt'] else self._comp_id) return comp_id def _get_voc_results_file_template(self): # VOCdevkit/results/VOC2007/Main/_det_test_aeroplane.txt filename = self._get_comp_id() + '_det_' + self._image_set + '_{:s}.txt' filedir = os.path.join(self._devkit_path, 'results', 'VOC' + self._year, 'Main') if not os.path.exists(filedir): os.makedirs(filedir) path = os.path.join(filedir, filename) return path def _write_voc_results_file(self, all_boxes): for cls_ind, cls in enumerate(self.classes): if cls == '__background__': continue print('Writing {} VOC results file'.format(cls)) filename = self._get_voc_results_file_template().format(cls) with open(filename, 'wt') as f: for im_ind, index in enumerate(self.image_index): dets = all_boxes[cls_ind][im_ind] if dets == []: continue # the VOCdevkit expects 1-based indices for k in xrange(dets.shape[0]): f.write('{:s} {:.3f} {:.1f} {:.1f} {:.1f} {:.1f}\n'. format(index, dets[k, -1], dets[k, 0] + 1, dets[k, 1] + 1, dets[k, 2] + 1, dets[k, 3] + 1)) def _do_python_eval(self, output_dir='output'): annopath = os.path.join( self._devkit_path, 'VOC' + self._year, 'Annotations', '{:s}.xml') imagesetfile = os.path.join( self._devkit_path, 'VOC' + self._year, 'ImageSets', 'Main', self._image_set + '.txt') cachedir = os.path.join(self._devkit_path, 'annotations_cache') aps = [] # The PASCAL VOC metric changed in 2010 use_07_metric = True if int(self._year) < 2010 else False print('VOC07 metric? ' + ('Yes' if use_07_metric else 'No')) if not os.path.isdir(output_dir): os.mkdir(output_dir) for i, cls in enumerate(self._classes): if cls == '__background__': continue filename = self._get_voc_results_file_template().format(cls) rec, prec, ap = voc_eval( filename, annopath, imagesetfile, cls, cachedir, ovthresh=0.5, use_07_metric=use_07_metric) aps += [ap] print('AP for {} = {:.4f}'.format(cls, ap)) with open(os.path.join(output_dir, cls + '_pr.pkl'), 'wb') as f: pickle.dump({'rec': rec, 'prec': prec, 'ap': ap}, f) print('Mean AP = {:.4f}'.format(np.mean(aps))) print('~~~~~~~~') print('Results:') for ap in aps: print('{:.3f}'.format(ap)) print('{:.3f}'.format(np.mean(aps))) print('~~~~~~~~') print('') print('--------------------------------------------------------------') print('Results computed with the **unofficial** Python eval code.') print('Results should be very close to the official MATLAB eval code.') print('Recompute with `./tools/reval.py --matlab ...` for your paper.') print('-- Thanks, The Management') print('--------------------------------------------------------------') def _do_matlab_eval(self, output_dir='output'): print('-----------------------------------------------------') print('Computing results with the official MATLAB eval code.') print('-----------------------------------------------------') path = os.path.join(cfg.ROOT_DIR, 'lib', 'datasets', 'VOCdevkit-matlab-wrapper') cmd = 'cd {} && '.format(path) cmd += '{:s} -nodisplay -nodesktop '.format(cfg.MATLAB) cmd += '-r "dbstop if error; ' cmd += 'voc_eval(\'{:s}\',\'{:s}\',\'{:s}\',\'{:s}\'); quit;"' \ .format(self._devkit_path, self._get_comp_id(), self._image_set, output_dir) print('Running:\n{}'.format(cmd)) status = subprocess.call(cmd, shell=True) def evaluate_detections(self, all_boxes, output_dir): self._write_voc_results_file(all_boxes) self._do_python_eval(output_dir) if self.config['matlab_eval']: self._do_matlab_eval(output_dir) if self.config['cleanup']: for cls in self._classes: if cls == '__background__': continue filename = self._get_voc_results_file_template().format(cls) os.remove(filename) def competition_mode(self, on): if on: self.config['use_salt'] = False self.config['cleanup'] = False else: self.config['use_salt'] = True self.config['cleanup'] = True if __name__ == '__main__': d = pascal_voc('trainval', '2007') res = d.roidb from IPython import embed; embed() ================================================ FILE: lib/datasets/pascal_voc_rbg.py ================================================ # -------------------------------------------------------- # Fast R-CNN # Copyright (c) 2015 Microsoft # Licensed under The MIT License [see LICENSE for details] # Written by Ross Girshick and Xinlei Chen # -------------------------------------------------------- from __future__ import absolute_import from __future__ import division from __future__ import print_function import os from datasets.imdb import imdb import datasets.ds_utils as ds_utils import xml.etree.ElementTree as ET import numpy as np import scipy.sparse import scipy.io as sio import model.utils.cython_bbox import pickle import subprocess import uuid from .voc_eval import voc_eval from model.utils.config import cfg import pdb class pascal_voc(imdb): def __init__(self, image_set, year, devkit_path=None): imdb.__init__(self, 'voc_' + year + '_' + image_set) self._year = year self._image_set = image_set self._devkit_path = self._get_default_path() if devkit_path is None \ else devkit_path self._data_path = os.path.join(self._devkit_path, 'VOC' + self._year) self._classes = ('__background__', # always index 0 'aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog', 'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor') self._class_to_ind = dict(list(zip(self.classes, list(range(self.num_classes))))) self._image_ext = '.jpg' self._image_index = self._load_image_set_index() # Default to roidb handler self._roidb_handler = self.gt_roidb self._salt = str(uuid.uuid4()) self._comp_id = 'comp4' # PASCAL specific config options self.config = {'cleanup': True, 'use_salt': True, 'use_diff': False, 'matlab_eval': False, 'rpn_file': None} assert os.path.exists(self._devkit_path), \ 'VOCdevkit path does not exist: {}'.format(self._devkit_path) assert os.path.exists(self._data_path), \ 'Path does not exist: {}'.format(self._data_path) def image_path_at(self, i): """ Return the absolute path to image i in the image sequence. """ return self.image_path_from_index(self._image_index[i]) def image_path_from_index(self, index): """ Construct an image path from the image's "index" identifier. """ image_path = os.path.join(self._data_path, 'JPEGImages', index + self._image_ext) assert os.path.exists(image_path), \ 'Path does not exist: {}'.format(image_path) return image_path def _load_image_set_index(self): """ Load the indexes listed in this dataset's image set file. """ # Example path to image set file: # self._devkit_path + /VOCdevkit2007/VOC2007/ImageSets/Main/val.txt image_set_file = os.path.join(self._data_path, 'ImageSets', 'Main', self._image_set + '.txt') assert os.path.exists(image_set_file), \ 'Path does not exist: {}'.format(image_set_file) with open(image_set_file) as f: image_index = [x.strip() for x in f.readlines()] return image_index def _get_default_path(self): """ Return the default path where PASCAL VOC is expected to be installed. """ return os.path.join(cfg.DATA_DIR, 'VOCdevkit' + self._year) def gt_roidb(self): """ Return the database of ground-truth regions of interest. This function loads/saves from/to a cache file to speed up future calls. """ cache_file = os.path.join(self.cache_path, self.name + '_gt_roidb.pkl') if os.path.exists(cache_file): with open(cache_file, 'rb') as fid: try: roidb = pickle.load(fid) except: roidb = pickle.load(fid, encoding='bytes') print('{} gt roidb loaded from {}'.format(self.name, cache_file)) return roidb gt_roidb = [self._load_pascal_annotation(index) for index in self.image_index] with open(cache_file, 'wb') as fid: pickle.dump(gt_roidb, fid, pickle.HIGHEST_PROTOCOL) print('wrote gt roidb to {}'.format(cache_file)) return gt_roidb def rpn_roidb(self): if int(self._year) == 2007 or self._image_set != 'test': gt_roidb = self.gt_roidb() rpn_roidb = self._load_rpn_roidb(gt_roidb) roidb = imdb.merge_roidbs(gt_roidb, rpn_roidb) else: roidb = self._load_rpn_roidb(None) return roidb def _load_rpn_roidb(self, gt_roidb): filename = self.config['rpn_file'] print('loading {}'.format(filename)) assert os.path.exists(filename), \ 'rpn data not found at: {}'.format(filename) with open(filename, 'rb') as f: box_list = pickle.load(f) return self.create_roidb_from_box_list(box_list, gt_roidb) def _load_pascal_annotation(self, index): """ Load image and bounding boxes info from XML file in the PASCAL VOC format. """ filename = os.path.join(self._data_path, 'Annotations', index + '.xml') tree = ET.parse(filename) objs = tree.findall('object') if not self.config['use_diff']: # Exclude the samples labeled as difficult non_diff_objs = [ obj for obj in objs if int(obj.find('difficult').text) == 0] # if len(non_diff_objs) != len(objs): # print 'Removed {} difficult objects'.format( # len(objs) - len(non_diff_objs)) objs = non_diff_objs num_objs = len(objs) boxes = np.zeros((num_objs, 4), dtype=np.uint16) gt_classes = np.zeros((num_objs), dtype=np.int32) overlaps = np.zeros((num_objs, self.num_classes), dtype=np.float32) # "Seg" area for pascal is just the box area seg_areas = np.zeros((num_objs), dtype=np.float32) # Load object bounding boxes into a data frame. for ix, obj in enumerate(objs): bbox = obj.find('bndbox') # Make pixel indexes 0-based x1 = float(bbox.find('xmin').text) - 1 y1 = float(bbox.find('ymin').text) - 1 x2 = float(bbox.find('xmax').text) - 1 y2 = float(bbox.find('ymax').text) - 1 cls = self._class_to_ind[obj.find('name').text.lower().strip()] boxes[ix, :] = [x1, y1, x2, y2] gt_classes[ix] = cls overlaps[ix, cls] = 1.0 seg_areas[ix] = (x2 - x1 + 1) * (y2 - y1 + 1) overlaps = scipy.sparse.csr_matrix(overlaps) return {'boxes': boxes, 'gt_classes': gt_classes, 'gt_overlaps': overlaps, 'flipped': False, 'seg_areas': seg_areas} def _get_comp_id(self): comp_id = (self._comp_id + '_' + self._salt if self.config['use_salt'] else self._comp_id) return comp_id def _get_voc_results_file_template(self): # VOCdevkit/results/VOC2007/Main/_det_test_aeroplane.txt filename = self._get_comp_id() + '_det_' + self._image_set + '_{:s}.txt' path = os.path.join( self._devkit_path, 'results', 'VOC' + self._year, 'Main', filename) return path def _write_voc_results_file(self, all_boxes): for cls_ind, cls in enumerate(self.classes): if cls == '__background__': continue print('Writing {} VOC results file'.format(cls)) filename = self._get_voc_results_file_template().format(cls) with open(filename, 'wt') as f: for im_ind, index in enumerate(self.image_index): dets = all_boxes[cls_ind][im_ind] if dets == []: continue # the VOCdevkit expects 1-based indices for k in range(dets.shape[0]): f.write('{:s} {:.3f} {:.1f} {:.1f} {:.1f} {:.1f}\n'. format(index, dets[k, -1], dets[k, 0] + 1, dets[k, 1] + 1, dets[k, 2] + 1, dets[k, 3] + 1)) def _do_python_eval(self, output_dir='output'): annopath = os.path.join( self._devkit_path, 'VOC' + self._year, 'Annotations', '{:s}.xml') imagesetfile = os.path.join( self._devkit_path, 'VOC' + self._year, 'ImageSets', 'Main', self._image_set + '.txt') cachedir = os.path.join(self._devkit_path, 'annotations_cache') aps = [] # The PASCAL VOC metric changed in 2010 use_07_metric = True if int(self._year) < 2010 else False print('VOC07 metric? ' + ('Yes' if use_07_metric else 'No')) if not os.path.isdir(output_dir): os.mkdir(output_dir) for i, cls in enumerate(self._classes): if cls == '__background__': continue filename = self._get_voc_results_file_template().format(cls) rec, prec, ap = voc_eval( filename, annopath, imagesetfile, cls, cachedir, ovthresh=0.5, use_07_metric=use_07_metric) aps += [ap] print(('AP for {} = {:.4f}'.format(cls, ap))) with open(os.path.join(output_dir, cls + '_pr.pkl'), 'wb') as f: pickle.dump({'rec': rec, 'prec': prec, 'ap': ap}, f) print(('Mean AP = {:.4f}'.format(np.mean(aps)))) print('~~~~~~~~') print('Results:') for ap in aps: print(('{:.3f}'.format(ap))) print(('{:.3f}'.format(np.mean(aps)))) print('~~~~~~~~') print('') print('--------------------------------------------------------------') print('Results computed with the **unofficial** Python eval code.') print('Results should be very close to the official MATLAB eval code.') print('Recompute with `./tools/reval.py --matlab ...` for your paper.') print('-- Thanks, The Management') print('--------------------------------------------------------------') def _do_matlab_eval(self, output_dir='output'): print('-----------------------------------------------------') print('Computing results with the official MATLAB eval code.') print('-----------------------------------------------------') path = os.path.join(cfg.ROOT_DIR, 'lib', 'datasets', 'VOCdevkit-matlab-wrapper') cmd = 'cd {} && '.format(path) cmd += '{:s} -nodisplay -nodesktop '.format(cfg.MATLAB) cmd += '-r "dbstop if error; ' cmd += 'voc_eval(\'{:s}\',\'{:s}\',\'{:s}\',\'{:s}\'); quit;"' \ .format(self._devkit_path, self._get_comp_id(), self._image_set, output_dir) print(('Running:\n{}'.format(cmd))) status = subprocess.call(cmd, shell=True) def evaluate_detections(self, all_boxes, output_dir): pdb.set_trace() self._write_voc_results_file(all_boxes) self._do_python_eval(output_dir) if self.config['matlab_eval']: self._do_matlab_eval(output_dir) if self.config['cleanup']: for cls in self._classes: if cls == '__background__': continue filename = self._get_voc_results_file_template().format(cls) os.remove(filename) def competition_mode(self, on): if on: self.config['use_salt'] = False self.config['cleanup'] = False else: self.config['use_salt'] = True self.config['cleanup'] = True if __name__ == '__main__': from datasets.pascal_voc import pascal_voc d = pascal_voc('trainval', '2007') res = d.roidb from IPython import embed; embed() ================================================ FILE: lib/datasets/tools/mcg_munge.py ================================================ from __future__ import print_function import os import sys """Hacky tool to convert file system layout of MCG boxes downloaded from http://www.eecs.berkeley.edu/Research/Projects/CS/vision/grouping/mcg/ so that it's consistent with those computed by Jan Hosang (see: http://www.mpi-inf.mpg.de/departments/computer-vision-and-multimodal- computing/research/object-recognition-and-scene-understanding/how- good-are-detection-proposals-really/) NB: Boxes from the MCG website are in (y1, x1, y2, x2) order. Boxes from Hosang et al. are in (x1, y1, x2, y2) order. """ def munge(src_dir): # stored as: ./MCG-COCO-val2014-boxes/COCO_val2014_000000193401.mat # want: ./MCG/mat/COCO_val2014_0/COCO_val2014_000000141/COCO_val2014_000000141334.mat files = os.listdir(src_dir) for fn in files: base, ext = os.path.splitext(fn) # first 14 chars / first 22 chars / all chars + .mat # COCO_val2014_0/COCO_val2014_000000447/COCO_val2014_000000447991.mat first = base[:14] second = base[:22] dst_dir = os.path.join('MCG', 'mat', first, second) if not os.path.exists(dst_dir): os.makedirs(dst_dir) src = os.path.join(src_dir, fn) dst = os.path.join(dst_dir, fn) print('MV: {} -> {}'.format(src, dst)) os.rename(src, dst) if __name__ == '__main__': # src_dir should look something like: # src_dir = 'MCG-COCO-val2014-boxes' src_dir = sys.argv[1] munge(src_dir) ================================================ FILE: lib/datasets/vg.py ================================================ from __future__ import print_function from __future__ import absolute_import # -------------------------------------------------------- # Fast R-CNN # Copyright (c) 2015 Microsoft # Licensed under The MIT License [see LICENSE for details] # Written by Ross Girshick # -------------------------------------------------------- import os from datasets.imdb import imdb import datasets.ds_utils as ds_utils import xml.etree.ElementTree as ET import numpy as np import scipy.sparse import gzip import PIL import json from .vg_eval import vg_eval from model.utils.config import cfg import pickle import pdb try: xrange # Python 2 except NameError: xrange = range # Python 3 class vg(imdb): def __init__(self, version, image_set, ): imdb.__init__(self, 'vg_' + version + '_' + image_set) self._version = version self._image_set = image_set self._data_path = os.path.join(cfg.DATA_DIR, 'genome') self._img_path = os.path.join(cfg.DATA_DIR, 'vg') # VG specific config options self.config = {'cleanup' : False} # Load classes self._classes = ['__background__'] self._class_to_ind = {} self._class_to_ind[self._classes[0]] = 0 with open(os.path.join(self._data_path, self._version, 'objects_vocab.txt')) as f: count = 1 for object in f.readlines(): names = [n.lower().strip() for n in object.split(',')] self._classes.append(names[0]) for n in names: self._class_to_ind[n] = count count += 1 # Load attributes self._attributes = ['__no_attribute__'] self._attribute_to_ind = {} self._attribute_to_ind[self._attributes[0]] = 0 with open(os.path.join(self._data_path, self._version, 'attributes_vocab.txt')) as f: count = 1 for att in f.readlines(): names = [n.lower().strip() for n in att.split(',')] self._attributes.append(names[0]) for n in names: self._attribute_to_ind[n] = count count += 1 # Load relations self._relations = ['__no_relation__'] self._relation_to_ind = {} self._relation_to_ind[self._relations[0]] = 0 with open(os.path.join(self._data_path, self._version, 'relations_vocab.txt')) as f: count = 1 for rel in f.readlines(): names = [n.lower().strip() for n in rel.split(',')] self._relations.append(names[0]) for n in names: self._relation_to_ind[n] = count count += 1 self._image_ext = '.jpg' load_index_from_file = False if os.path.exists(os.path.join(self._data_path, "vg_image_index_{}.p".format(self._image_set))): with open(os.path.join(self._data_path, "vg_image_index_{}.p".format(self._image_set)), 'rb') as fp: self._image_index = pickle.load(fp) load_index_from_file = True load_id_from_file = False if os.path.exists(os.path.join(self._data_path, "vg_id_to_dir_{}.p".format(self._image_set))): with open(os.path.join(self._data_path, "vg_id_to_dir_{}.p".format(self._image_set)), 'rb') as fp: self._id_to_dir = pickle.load(fp) load_id_from_file = True if not load_index_from_file or not load_id_from_file: self._image_index, self._id_to_dir = self._load_image_set_index() with open(os.path.join(self._data_path, "vg_image_index_{}.p".format(self._image_set)), 'wb') as fp: pickle.dump(self._image_index, fp) with open(os.path.join(self._data_path, "vg_id_to_dir_{}.p".format(self._image_set)), 'wb') as fp: pickle.dump(self._id_to_dir, fp) self._roidb_handler = self.gt_roidb def image_path_at(self, i): """ Return the absolute path to image i in the image sequence. """ return self.image_path_from_index(self._image_index[i]) def image_id_at(self, i): """ Return the absolute path to image i in the image sequence. """ return i # return self._image_index[i] def image_path_from_index(self, index): """ Construct an image path from the image's "index" identifier. """ folder = self._id_to_dir[index] image_path = os.path.join(self._img_path, folder, str(index) + self._image_ext) assert os.path.exists(image_path), \ 'Path does not exist: {}'.format(image_path) return image_path def _image_split_path(self): if self._image_set == "minitrain": return os.path.join(self._data_path, 'train.txt') if self._image_set == "smalltrain": return os.path.join(self._data_path, 'train.txt') if self._image_set == "minival": return os.path.join(self._data_path, 'val.txt') if self._image_set == "smallval": return os.path.join(self._data_path, 'val.txt') else: return os.path.join(self._data_path, self._image_set+'.txt') def _load_image_set_index(self): """ Load the indexes listed in this dataset's image set file. """ training_split_file = self._image_split_path() assert os.path.exists(training_split_file), \ 'Path does not exist: {}'.format(training_split_file) with open(training_split_file) as f: metadata = f.readlines() if self._image_set == "minitrain": metadata = metadata[:1000] elif self._image_set == "smalltrain": metadata = metadata[:20000] elif self._image_set == "minival": metadata = metadata[:100] elif self._image_set == "smallval": metadata = metadata[:2000] image_index = [] id_to_dir = {} for line in metadata: im_file,ann_file = line.split() image_id = int(ann_file.split('/')[-1].split('.')[0]) filename = self._annotation_path(image_id) if os.path.exists(filename): # Some images have no bboxes after object filtering, so there # is no xml annotation for these. tree = ET.parse(filename) for obj in tree.findall('object'): obj_name = obj.find('name').text.lower().strip() if obj_name in self._class_to_ind: # We have to actually load and check these to make sure they have # at least one object actually in vocab image_index.append(image_id) id_to_dir[image_id] = im_file.split('/')[0] break return image_index, id_to_dir def gt_roidb(self): """ Return the database of ground-truth regions of interest. This function loads/saves from/to a cache file to speed up future calls. """ cache_file = os.path.join(self.cache_path, self.name + '_gt_roidb.pkl') if os.path.exists(cache_file): fid = gzip.open(cache_file,'rb') roidb = pickle.load(fid) fid.close() print('{} gt roidb loaded from {}'.format(self.name, cache_file)) return roidb gt_roidb = [self._load_vg_annotation(index) for index in self.image_index] fid = gzip.open(cache_file,'wb') pickle.dump(gt_roidb, fid, pickle.HIGHEST_PROTOCOL) fid.close() print('wrote gt roidb to {}'.format(cache_file)) return gt_roidb def _get_size(self, index): return PIL.Image.open(self.image_path_from_index(index)).size def _annotation_path(self, index): return os.path.join(self._data_path, 'xml', str(index) + '.xml') def _load_vg_annotation(self, index): """ Load image and bounding boxes info from XML file in the PASCAL VOC format. """ width, height = self._get_size(index) filename = self._annotation_path(index) tree = ET.parse(filename) objs = tree.findall('object') num_objs = len(objs) boxes = np.zeros((num_objs, 4), dtype=np.uint16) gt_classes = np.zeros((num_objs), dtype=np.int32) # Max of 16 attributes are observed in the data gt_attributes = np.zeros((num_objs, 16), dtype=np.int32) overlaps = np.zeros((num_objs, self.num_classes), dtype=np.float32) # "Seg" area for pascal is just the box area seg_areas = np.zeros((num_objs), dtype=np.float32) # Load object bounding boxes into a data frame. obj_dict = {} ix = 0 for obj in objs: obj_name = obj.find('name').text.lower().strip() if obj_name in self._class_to_ind: bbox = obj.find('bndbox') x1 = max(0,float(bbox.find('xmin').text)) y1 = max(0,float(bbox.find('ymin').text)) x2 = min(width-1,float(bbox.find('xmax').text)) y2 = min(height-1,float(bbox.find('ymax').text)) # If bboxes are not positive, just give whole image coords (there are a few examples) if x2 < x1 or y2 < y1: print('Failed bbox in %s, object %s' % (filename, obj_name)) x1 = 0 y1 = 0 x2 = width-1 y2 = width-1 cls = self._class_to_ind[obj_name] obj_dict[obj.find('object_id').text] = ix atts = obj.findall('attribute') n = 0 for att in atts: att = att.text.lower().strip() if att in self._attribute_to_ind: gt_attributes[ix, n] = self._attribute_to_ind[att] n += 1 if n >= 16: break boxes[ix, :] = [x1, y1, x2, y2] gt_classes[ix] = cls overlaps[ix, cls] = 1.0 seg_areas[ix] = (x2 - x1 + 1) * (y2 - y1 + 1) ix += 1 # clip gt_classes and gt_relations gt_classes = gt_classes[:ix] gt_attributes = gt_attributes[:ix, :] overlaps = scipy.sparse.csr_matrix(overlaps) gt_attributes = scipy.sparse.csr_matrix(gt_attributes) rels = tree.findall('relation') num_rels = len(rels) gt_relations = set() # Avoid duplicates for rel in rels: pred = rel.find('predicate').text if pred: # One is empty pred = pred.lower().strip() if pred in self._relation_to_ind: try: triple = [] triple.append(obj_dict[rel.find('subject_id').text]) triple.append(self._relation_to_ind[pred]) triple.append(obj_dict[rel.find('object_id').text]) gt_relations.add(tuple(triple)) except: pass # Object not in dictionary gt_relations = np.array(list(gt_relations), dtype=np.int32) return {'boxes' : boxes, 'gt_classes': gt_classes, 'gt_attributes' : gt_attributes, 'gt_relations' : gt_relations, 'gt_overlaps' : overlaps, 'width' : width, 'height': height, 'flipped' : False, 'seg_areas' : seg_areas} def evaluate_detections(self, all_boxes, output_dir): self._write_voc_results_file(self.classes, all_boxes, output_dir) self._do_python_eval(output_dir) if self.config['cleanup']: for cls in self._classes: if cls == '__background__': continue filename = self._get_vg_results_file_template(output_dir).format(cls) os.remove(filename) def evaluate_attributes(self, all_boxes, output_dir): self._write_voc_results_file(self.attributes, all_boxes, output_dir) self._do_python_eval(output_dir, eval_attributes = True) if self.config['cleanup']: for cls in self._attributes: if cls == '__no_attribute__': continue filename = self._get_vg_results_file_template(output_dir).format(cls) os.remove(filename) def _get_vg_results_file_template(self, output_dir): filename = 'detections_' + self._image_set + '_{:s}.txt' path = os.path.join(output_dir, filename) return path def _write_voc_results_file(self, classes, all_boxes, output_dir): for cls_ind, cls in enumerate(classes): if cls == '__background__': continue print('Writing "{}" vg results file'.format(cls)) filename = self._get_vg_results_file_template(output_dir).format(cls) with open(filename, 'wt') as f: for im_ind, index in enumerate(self.image_index): dets = all_boxes[cls_ind][im_ind] if dets == []: continue # the VOCdevkit expects 1-based indices for k in xrange(dets.shape[0]): f.write('{:s} {:.3f} {:.1f} {:.1f} {:.1f} {:.1f}\n'. format(str(index), dets[k, -1], dets[k, 0] + 1, dets[k, 1] + 1, dets[k, 2] + 1, dets[k, 3] + 1)) def _do_python_eval(self, output_dir, pickle=True, eval_attributes = False): # We re-use parts of the pascal voc python code for visual genome aps = [] nposs = [] thresh = [] # The PASCAL VOC metric changed in 2010 use_07_metric = False print('VOC07 metric? ' + ('Yes' if use_07_metric else 'No')) if not os.path.isdir(output_dir): os.mkdir(output_dir) # Load ground truth gt_roidb = self.gt_roidb() if eval_attributes: classes = self._attributes else: classes = self._classes for i, cls in enumerate(classes): if cls == '__background__' or cls == '__no_attribute__': continue filename = self._get_vg_results_file_template(output_dir).format(cls) rec, prec, ap, scores, npos = vg_eval( filename, gt_roidb, self.image_index, i, ovthresh=0.5, use_07_metric=use_07_metric, eval_attributes=eval_attributes) # Determine per class detection thresholds that maximise f score if npos > 1: f = np.nan_to_num((prec*rec)/(prec+rec)) thresh += [scores[np.argmax(f)]] else: thresh += [0] aps += [ap] nposs += [float(npos)] print('AP for {} = {:.4f} (npos={:,})'.format(cls, ap, npos)) if pickle: with open(os.path.join(output_dir, cls + '_pr.pkl'), 'wb') as f: pickle.dump({'rec': rec, 'prec': prec, 'ap': ap, 'scores': scores, 'npos':npos}, f) # Set thresh to mean for classes with poor results thresh = np.array(thresh) avg_thresh = np.mean(thresh[thresh!=0]) thresh[thresh==0] = avg_thresh if eval_attributes: filename = 'attribute_thresholds_' + self._image_set + '.txt' else: filename = 'object_thresholds_' + self._image_set + '.txt' path = os.path.join(output_dir, filename) with open(path, 'wt') as f: for i, cls in enumerate(classes[1:]): f.write('{:s} {:.3f}\n'.format(cls, thresh[i])) weights = np.array(nposs) weights /= weights.sum() print('Mean AP = {:.4f}'.format(np.mean(aps))) print('Weighted Mean AP = {:.4f}'.format(np.average(aps, weights=weights))) print('Mean Detection Threshold = {:.3f}'.format(avg_thresh)) print('~~~~~~~~') print('Results:') for ap,npos in zip(aps,nposs): print('{:.3f}\t{:.3f}'.format(ap,npos)) print('{:.3f}'.format(np.mean(aps))) print('~~~~~~~~') print('') print('--------------------------------------------------------------') print('Results computed with the **unofficial** PASCAL VOC Python eval code.') print('--------------------------------------------------------------') if __name__ == '__main__': d = vg('val') res = d.roidb from IPython import embed; embed() ================================================ FILE: lib/datasets/vg_eval.py ================================================ from __future__ import absolute_import # -------------------------------------------------------- # Fast/er R-CNN # Licensed under The MIT License [see LICENSE for details] # Written by Bharath Hariharan # -------------------------------------------------------- import xml.etree.ElementTree as ET import os import numpy as np from .voc_eval import voc_ap def vg_eval( detpath, gt_roidb, image_index, classindex, ovthresh=0.5, use_07_metric=False, eval_attributes=False): """rec, prec, ap, sorted_scores, npos = voc_eval( detpath, gt_roidb, image_index, classindex, [ovthresh], [use_07_metric]) Top level function that does the Visual Genome evaluation. detpath: Path to detections gt_roidb: List of ground truth structs. image_index: List of image ids. classindex: Category index [ovthresh]: Overlap threshold (default = 0.5) [use_07_metric]: Whether to use VOC07's 11 point AP computation (default False) """ # extract gt objects for this class class_recs = {} npos = 0 for item,imagename in zip(gt_roidb,image_index): if eval_attributes: bbox = item['boxes'][np.where(np.any(item['gt_attributes'].toarray() == classindex, axis=1))[0], :] else: bbox = item['boxes'][np.where(item['gt_classes'] == classindex)[0], :] difficult = np.zeros((bbox.shape[0],)).astype(np.bool) det = [False] * bbox.shape[0] npos = npos + sum(~difficult) class_recs[str(imagename)] = {'bbox': bbox, 'difficult': difficult, 'det': det} if npos == 0: # No ground truth examples return 0,0,0,0,npos # read dets with open(detpath, 'r') as f: lines = f.readlines() if len(lines) == 0: # No detection examples return 0,0,0,0,npos splitlines = [x.strip().split(' ') for x in lines] image_ids = [x[0] for x in splitlines] confidence = np.array([float(x[1]) for x in splitlines]) BB = np.array([[float(z) for z in x[2:]] for x in splitlines]) # sort by confidence sorted_ind = np.argsort(-confidence) sorted_scores = -np.sort(-confidence) BB = BB[sorted_ind, :] image_ids = [image_ids[x] for x in sorted_ind] # go down dets and mark TPs and FPs nd = len(image_ids) tp = np.zeros(nd) fp = np.zeros(nd) for d in range(nd): R = class_recs[image_ids[d]] bb = BB[d, :].astype(float) ovmax = -np.inf BBGT = R['bbox'].astype(float) if BBGT.size > 0: # compute overlaps # intersection ixmin = np.maximum(BBGT[:, 0], bb[0]) iymin = np.maximum(BBGT[:, 1], bb[1]) ixmax = np.minimum(BBGT[:, 2], bb[2]) iymax = np.minimum(BBGT[:, 3], bb[3]) iw = np.maximum(ixmax - ixmin + 1., 0.) ih = np.maximum(iymax - iymin + 1., 0.) inters = iw * ih # union uni = ((bb[2] - bb[0] + 1.) * (bb[3] - bb[1] + 1.) + (BBGT[:, 2] - BBGT[:, 0] + 1.) * (BBGT[:, 3] - BBGT[:, 1] + 1.) - inters) overlaps = inters / uni ovmax = np.max(overlaps) jmax = np.argmax(overlaps) if ovmax > ovthresh: if not R['difficult'][jmax]: if not R['det'][jmax]: tp[d] = 1. R['det'][jmax] = 1 else: fp[d] = 1. else: fp[d] = 1. # compute precision recall fp = np.cumsum(fp) tp = np.cumsum(tp) rec = tp / float(npos) # avoid divide by zero in case the first detection matches a difficult # ground truth prec = tp / np.maximum(tp + fp, np.finfo(np.float64).eps) ap = voc_ap(rec, prec, use_07_metric) return rec, prec, ap, sorted_scores, npos ================================================ FILE: lib/datasets/voc_eval.py ================================================ # -------------------------------------------------------- # Fast/er R-CNN # Licensed under The MIT License [see LICENSE for details] # Written by Bharath Hariharan # -------------------------------------------------------- from __future__ import absolute_import from __future__ import division from __future__ import print_function import xml.etree.ElementTree as ET import os import pickle import numpy as np def parse_rec(filename): """ Parse a PASCAL VOC xml file """ tree = ET.parse(filename) objects = [] for obj in tree.findall('object'): obj_struct = {} obj_struct['name'] = obj.find('name').text obj_struct['pose'] = obj.find('pose').text obj_struct['truncated'] = int(obj.find('truncated').text) obj_struct['difficult'] = int(obj.find('difficult').text) bbox = obj.find('bndbox') obj_struct['bbox'] = [int(bbox.find('xmin').text), int(bbox.find('ymin').text), int(bbox.find('xmax').text), int(bbox.find('ymax').text)] objects.append(obj_struct) return objects def voc_ap(rec, prec, use_07_metric=False): """ ap = voc_ap(rec, prec, [use_07_metric]) Compute VOC AP given precision and recall. If use_07_metric is true, uses the VOC 07 11 point method (default:False). """ if use_07_metric: # 11 point metric ap = 0. for t in np.arange(0., 1.1, 0.1): if np.sum(rec >= t) == 0: p = 0 else: p = np.max(prec[rec >= t]) ap = ap + p / 11. else: # correct AP calculation # first append sentinel values at the end mrec = np.concatenate(([0.], rec, [1.])) mpre = np.concatenate(([0.], prec, [0.])) # compute the precision envelope for i in range(mpre.size - 1, 0, -1): mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i]) # to calculate area under PR curve, look for points # where X axis (recall) changes value i = np.where(mrec[1:] != mrec[:-1])[0] # and sum (\Delta recall) * prec ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1]) return ap def voc_eval(detpath, annopath, imagesetfile, classname, cachedir, ovthresh=0.5, use_07_metric=False): """rec, prec, ap = voc_eval(detpath, annopath, imagesetfile, classname, [ovthresh], [use_07_metric]) Top level function that does the PASCAL VOC evaluation. detpath: Path to detections detpath.format(classname) should produce the detection results file. annopath: Path to annotations annopath.format(imagename) should be the xml annotations file. imagesetfile: Text file containing the list of images, one image per line. classname: Category name (duh) cachedir: Directory for caching the annotations [ovthresh]: Overlap threshold (default = 0.5) [use_07_metric]: Whether to use VOC07's 11 point AP computation (default False) """ # assumes detections are in detpath.format(classname) # assumes annotations are in annopath.format(imagename) # assumes imagesetfile is a text file with each line an image name # cachedir caches the annotations in a pickle file # first load gt if not os.path.isdir(cachedir): os.mkdir(cachedir) cachefile = os.path.join(cachedir, '%s_annots.pkl' % imagesetfile) # read list of images with open(imagesetfile, 'r') as f: lines = f.readlines() imagenames = [x.strip() for x in lines] if not os.path.isfile(cachefile): # load annotations recs = {} for i, imagename in enumerate(imagenames): recs[imagename] = parse_rec(annopath.format(imagename)) if i % 100 == 0: print('Reading annotation for {:d}/{:d}'.format( i + 1, len(imagenames))) # save print('Saving cached annotations to {:s}'.format(cachefile)) with open(cachefile, 'wb') as f: pickle.dump(recs, f) else: # load with open(cachefile, 'rb') as f: try: recs = pickle.load(f) except: recs = pickle.load(f, encoding='bytes') # extract gt objects for this class class_recs = {} npos = 0 for imagename in imagenames: R = [obj for obj in recs[imagename] if obj['name'] == classname] bbox = np.array([x['bbox'] for x in R]) difficult = np.array([x['difficult'] for x in R]).astype(np.bool) det = [False] * len(R) npos = npos + sum(~difficult) class_recs[imagename] = {'bbox': bbox, 'difficult': difficult, 'det': det} # read dets detfile = detpath.format(classname) with open(detfile, 'r') as f: lines = f.readlines() splitlines = [x.strip().split(' ') for x in lines] image_ids = [x[0] for x in splitlines] confidence = np.array([float(x[1]) for x in splitlines]) BB = np.array([[float(z) for z in x[2:]] for x in splitlines]) nd = len(image_ids) tp = np.zeros(nd) fp = np.zeros(nd) if BB.shape[0] > 0: # sort by confidence sorted_ind = np.argsort(-confidence) sorted_scores = np.sort(-confidence) BB = BB[sorted_ind, :] image_ids = [image_ids[x] for x in sorted_ind] # go down dets and mark TPs and FPs for d in range(nd): R = class_recs[image_ids[d]] bb = BB[d, :].astype(float) ovmax = -np.inf BBGT = R['bbox'].astype(float) if BBGT.size > 0: # compute overlaps # intersection ixmin = np.maximum(BBGT[:, 0], bb[0]) iymin = np.maximum(BBGT[:, 1], bb[1]) ixmax = np.minimum(BBGT[:, 2], bb[2]) iymax = np.minimum(BBGT[:, 3], bb[3]) iw = np.maximum(ixmax - ixmin + 1., 0.) ih = np.maximum(iymax - iymin + 1., 0.) inters = iw * ih # union uni = ((bb[2] - bb[0] + 1.) * (bb[3] - bb[1] + 1.) + (BBGT[:, 2] - BBGT[:, 0] + 1.) * (BBGT[:, 3] - BBGT[:, 1] + 1.) - inters) overlaps = inters / uni ovmax = np.max(overlaps) jmax = np.argmax(overlaps) if ovmax > ovthresh: if not R['difficult'][jmax]: if not R['det'][jmax]: tp[d] = 1. R['det'][jmax] = 1 else: fp[d] = 1. else: fp[d] = 1. # compute precision recall fp = np.cumsum(fp) tp = np.cumsum(tp) rec = tp / float(npos) # avoid divide by zero in case the first detection matches a difficult # ground truth prec = tp / np.maximum(tp + fp, np.finfo(np.float64).eps) ap = voc_ap(rec, prec, use_07_metric) return rec, prec, ap ================================================ FILE: lib/make.sh ================================================ #!/usr/bin/env bash # CUDA_PATH=/usr/local/cuda/ export CUDA_PATH=/usr/local/cuda/ #You may also want to ad the following #export C_INCLUDE_PATH=/opt/cuda/include export CXXFLAGS="-std=c++11" export CFLAGS="-std=c99" python setup.py build_ext --inplace rm -rf build CUDA_ARCH="-gencode arch=compute_30,code=sm_30 \ -gencode arch=compute_35,code=sm_35 \ -gencode arch=compute_50,code=sm_50 \ -gencode arch=compute_52,code=sm_52 \ -gencode arch=compute_60,code=sm_60 \ -gencode arch=compute_61,code=sm_61 " # compile NMS cd model/nms/src echo "Compiling nms kernels by nvcc..." nvcc -c -o nms_cuda_kernel.cu.o nms_cuda_kernel.cu \ -D GOOGLE_CUDA=1 -x cu -Xcompiler -fPIC $CUDA_ARCH cd ../ python build.py # compile roi_pooling cd ../../ cd model/roi_pooling/src echo "Compiling roi pooling kernels by nvcc..." nvcc -c -o roi_pooling.cu.o roi_pooling_kernel.cu \ -D GOOGLE_CUDA=1 -x cu -Xcompiler -fPIC $CUDA_ARCH cd ../ python build.py # compile roi_align cd ../../ cd model/roi_align/src echo "Compiling roi align kernels by nvcc..." nvcc -c -o roi_align_kernel.cu.o roi_align_kernel.cu \ -D GOOGLE_CUDA=1 -x cu -Xcompiler -fPIC $CUDA_ARCH cd ../ python build.py # compile roi_crop cd ../../ cd model/roi_crop/src echo "Compiling roi crop kernels by nvcc..." nvcc -c -o roi_crop_cuda_kernel.cu.o roi_crop_cuda_kernel.cu \ -D GOOGLE_CUDA=1 -x cu -Xcompiler -fPIC $CUDA_ARCH cd ../ python build.py ================================================ FILE: lib/model/__init__.py ================================================ ================================================ FILE: lib/model/faster_rcnn/__init__.py ================================================ ================================================ FILE: lib/model/faster_rcnn/faster_rcnn.py ================================================ import random import torch import torch.nn as nn import torch.nn.functional as F from torch.autograd import Variable import torchvision.models as models from torch.autograd import Variable import numpy as np from model.utils.config import cfg from model.rpn.rpn import _RPN from model.roi_pooling.modules.roi_pool import _RoIPooling from model.roi_crop.modules.roi_crop import _RoICrop from model.roi_align.modules.roi_align import RoIAlignAvg from model.rpn.proposal_target_layer_cascade import _ProposalTargetLayer import time import pdb from model.utils.net_utils import _smooth_l1_loss, _crop_pool_layer, _affine_grid_gen, _affine_theta class _fasterRCNN(nn.Module): """ faster RCNN """ def __init__(self, classes, class_agnostic): super(_fasterRCNN, self).__init__() self.classes = classes self.n_classes = len(classes) self.class_agnostic = class_agnostic # loss self.RCNN_loss_cls = 0 self.RCNN_loss_bbox = 0 # define rpn self.RCNN_rpn = _RPN(self.dout_base_model) self.RCNN_proposal_target = _ProposalTargetLayer(self.n_classes) self.RCNN_roi_pool = _RoIPooling(cfg.POOLING_SIZE, cfg.POOLING_SIZE, 1.0/16.0) self.RCNN_roi_align = RoIAlignAvg(cfg.POOLING_SIZE, cfg.POOLING_SIZE, 1.0/16.0) self.grid_size = cfg.POOLING_SIZE * 2 if cfg.CROP_RESIZE_WITH_MAX_POOL else cfg.POOLING_SIZE self.RCNN_roi_crop = _RoICrop() def forward(self, im_data, im_info, gt_boxes, num_boxes): batch_size = im_data.size(0) im_info = im_info.data gt_boxes = gt_boxes.data num_boxes = num_boxes.data # feed image data to base model to obtain base feature map base_feat = self.RCNN_base(im_data) # feed base feature map to RPN to obtain rois rois, rpn_loss_cls, rpn_loss_bbox = self.RCNN_rpn(base_feat, im_info, gt_boxes, num_boxes) # if it is training phase, then use ground truth bboxes for refining if self.training: roi_data = self.RCNN_proposal_target(rois, gt_boxes, num_boxes) rois, rois_label, rois_target, rois_inside_ws, rois_outside_ws = roi_data rois_label = Variable(rois_label.view(-1).long()) rois_target = Variable(rois_target.view(-1, rois_target.size(2))) rois_inside_ws = Variable(rois_inside_ws.view(-1, rois_inside_ws.size(2))) rois_outside_ws = Variable(rois_outside_ws.view(-1, rois_outside_ws.size(2))) else: rois_label = None rois_target = None rois_inside_ws = None rois_outside_ws = None rpn_loss_cls = 0 rpn_loss_bbox = 0 rois = Variable(rois) # do roi pooling based on predicted rois if cfg.POOLING_MODE == 'crop': # pdb.set_trace() # pooled_feat_anchor = _crop_pool_layer(base_feat, rois.view(-1, 5)) grid_xy = _affine_grid_gen(rois.view(-1, 5), base_feat.size()[2:], self.grid_size) grid_yx = torch.stack([grid_xy.data[:,:,:,1], grid_xy.data[:,:,:,0]], 3).contiguous() pooled_feat = self.RCNN_roi_crop(base_feat, Variable(grid_yx).detach()) if cfg.CROP_RESIZE_WITH_MAX_POOL: pooled_feat = F.max_pool2d(pooled_feat, 2, 2) elif cfg.POOLING_MODE == 'align': pooled_feat = self.RCNN_roi_align(base_feat, rois.view(-1, 5)) elif cfg.POOLING_MODE == 'pool': pooled_feat = self.RCNN_roi_pool(base_feat, rois.view(-1,5)) # feed pooled features to top model pooled_feat = self._head_to_tail(pooled_feat) # compute bbox offset bbox_pred = self.RCNN_bbox_pred(pooled_feat) if self.training and not self.class_agnostic: # select the corresponding columns according to roi labels bbox_pred_view = bbox_pred.view(bbox_pred.size(0), int(bbox_pred.size(1) / 4), 4) bbox_pred_select = torch.gather(bbox_pred_view, 1, rois_label.view(rois_label.size(0), 1, 1).expand(rois_label.size(0), 1, 4)) bbox_pred = bbox_pred_select.squeeze(1) # compute object classification probability cls_score = self.RCNN_cls_score(pooled_feat) cls_prob = F.softmax(cls_score, 1) RCNN_loss_cls = 0 RCNN_loss_bbox = 0 if self.training: # classification loss RCNN_loss_cls = F.cross_entropy(cls_score, rois_label) # bounding box regression L1 loss RCNN_loss_bbox = _smooth_l1_loss(bbox_pred, rois_target, rois_inside_ws, rois_outside_ws) cls_prob = cls_prob.view(batch_size, rois.size(1), -1) bbox_pred = bbox_pred.view(batch_size, rois.size(1), -1) return rois, cls_prob, bbox_pred, rpn_loss_cls, rpn_loss_bbox, RCNN_loss_cls, RCNN_loss_bbox, rois_label def _init_weights(self): def normal_init(m, mean, stddev, truncated=False): """ weight initalizer: truncated normal and random normal. """ # x is a parameter if truncated: m.weight.data.normal_().fmod_(2).mul_(stddev).add_(mean) # not a perfect approximation else: m.weight.data.normal_(mean, stddev) m.bias.data.zero_() normal_init(self.RCNN_rpn.RPN_Conv, 0, 0.01, cfg.TRAIN.TRUNCATED) normal_init(self.RCNN_rpn.RPN_cls_score, 0, 0.01, cfg.TRAIN.TRUNCATED) normal_init(self.RCNN_rpn.RPN_bbox_pred, 0, 0.01, cfg.TRAIN.TRUNCATED) normal_init(self.RCNN_cls_score, 0, 0.01, cfg.TRAIN.TRUNCATED) normal_init(self.RCNN_bbox_pred, 0, 0.001, cfg.TRAIN.TRUNCATED) def create_architecture(self): self._init_modules() self._init_weights() ================================================ FILE: lib/model/faster_rcnn/resnet.py ================================================ from __future__ import absolute_import from __future__ import division from __future__ import print_function from model.utils.config import cfg from model.faster_rcnn.faster_rcnn import _fasterRCNN import torch import torch.nn as nn import torch.nn.functional as F from torch.autograd import Variable import math import torch.utils.model_zoo as model_zoo import pdb __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152'] model_urls = { 'resnet18': 'https://s3.amazonaws.com/pytorch/models/resnet18-5c106cde.pth', 'resnet34': 'https://s3.amazonaws.com/pytorch/models/resnet34-333f7ec4.pth', 'resnet50': 'https://s3.amazonaws.com/pytorch/models/resnet50-19c8e357.pth', 'resnet101': 'https://s3.amazonaws.com/pytorch/models/resnet101-5d3b4d8f.pth', 'resnet152': 'https://s3.amazonaws.com/pytorch/models/resnet152-b121ed2d.pth', } def conv3x3(in_planes, out_planes, stride=1): "3x3 convolution with padding" return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False) class BasicBlock(nn.Module): expansion = 1 def __init__(self, inplanes, planes, stride=1, downsample=None): super(BasicBlock, self).__init__() self.conv1 = conv3x3(inplanes, planes, stride) self.bn1 = nn.BatchNorm2d(planes) self.relu = nn.ReLU(inplace=True) self.conv2 = conv3x3(planes, planes) self.bn2 = nn.BatchNorm2d(planes) self.downsample = downsample self.stride = stride def forward(self, x): residual = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) out = self.conv2(out) out = self.bn2(out) if self.downsample is not None: residual = self.downsample(x) out += residual out = self.relu(out) return out class Bottleneck(nn.Module): expansion = 4 def __init__(self, inplanes, planes, stride=1, downsample=None): super(Bottleneck, self).__init__() self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, stride=stride, bias=False) # change self.bn1 = nn.BatchNorm2d(planes) self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, # change padding=1, bias=False) self.bn2 = nn.BatchNorm2d(planes) self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) self.bn3 = nn.BatchNorm2d(planes * 4) self.relu = nn.ReLU(inplace=True) self.downsample = downsample self.stride = stride def forward(self, x): residual = x out = self.conv1(x) out = self.bn1(out) out = self.relu(out) out = self.conv2(out) out = self.bn2(out) out = self.relu(out) out = self.conv3(out) out = self.bn3(out) if self.downsample is not None: residual = self.downsample(x) out += residual out = self.relu(out) return out class ResNet(nn.Module): def __init__(self, block, layers, num_classes=1000): self.inplanes = 64 super(ResNet, self).__init__() self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False) self.bn1 = nn.BatchNorm2d(64) self.relu = nn.ReLU(inplace=True) self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=0, ceil_mode=True) # change self.layer1 = self._make_layer(block, 64, layers[0]) self.layer2 = self._make_layer(block, 128, layers[1], stride=2) self.layer3 = self._make_layer(block, 256, layers[2], stride=2) self.layer4 = self._make_layer(block, 512, layers[3], stride=2) # it is slightly better whereas slower to set stride = 1 # self.layer4 = self._make_layer(block, 512, layers[3], stride=1) self.avgpool = nn.AvgPool2d(7) self.fc = nn.Linear(512 * block.expansion, num_classes) for m in self.modules(): if isinstance(m, nn.Conv2d): n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels m.weight.data.normal_(0, math.sqrt(2. / n)) elif isinstance(m, nn.BatchNorm2d): m.weight.data.fill_(1) m.bias.data.zero_() def _make_layer(self, block, planes, blocks, stride=1): downsample = None if stride != 1 or self.inplanes != planes * block.expansion: downsample = nn.Sequential( nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False), nn.BatchNorm2d(planes * block.expansion), ) layers = [] layers.append(block(self.inplanes, planes, stride, downsample)) self.inplanes = planes * block.expansion for i in range(1, blocks): layers.append(block(self.inplanes, planes)) return nn.Sequential(*layers) def forward(self, x): x = self.conv1(x) x = self.bn1(x) x = self.relu(x) x = self.maxpool(x) x = self.layer1(x) x = self.layer2(x) x = self.layer3(x) x = self.layer4(x) x = self.avgpool(x) x = x.view(x.size(0), -1) x = self.fc(x) return x def resnet18(pretrained=False): """Constructs a ResNet-18 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet """ model = ResNet(BasicBlock, [2, 2, 2, 2]) if pretrained: model.load_state_dict(model_zoo.load_url(model_urls['resnet18'])) return model def resnet34(pretrained=False): """Constructs a ResNet-34 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet """ model = ResNet(BasicBlock, [3, 4, 6, 3]) if pretrained: model.load_state_dict(model_zoo.load_url(model_urls['resnet34'])) return model def resnet50(pretrained=False): """Constructs a ResNet-50 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet """ model = ResNet(Bottleneck, [3, 4, 6, 3]) if pretrained: model.load_state_dict(model_zoo.load_url(model_urls['resnet50'])) return model def resnet101(pretrained=False): """Constructs a ResNet-101 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet """ model = ResNet(Bottleneck, [3, 4, 23, 3]) if pretrained: model.load_state_dict(model_zoo.load_url(model_urls['resnet101'])) return model def resnet152(pretrained=False): """Constructs a ResNet-152 model. Args: pretrained (bool): If True, returns a model pre-trained on ImageNet """ model = ResNet(Bottleneck, [3, 8, 36, 3]) if pretrained: model.load_state_dict(model_zoo.load_url(model_urls['resnet152'])) return model class resnet(_fasterRCNN): def __init__(self, classes, num_layers=101, pretrained=False, class_agnostic=False): self.model_path = 'data/pretrained_model/resnet101_caffe.pth' self.dout_base_model = 1024 self.pretrained = pretrained self.class_agnostic = class_agnostic self.num_layers = num_layers _fasterRCNN.__init__(self, classes, class_agnostic) def _init_modules(self): resnet = resnet101() if self.num_layers == 18: resnet = resnet18() if self.num_layers == 34: resnet = resnet34() if self.num_layers == 50: resnet = resnet50() if self.num_layers == 152: resnet = resnet152() if self.pretrained == True: print("Loading pretrained weights from %s" %(self.model_path)) state_dict = torch.load(self.model_path) resnet.load_state_dict({k:v for k,v in state_dict.items() if k in resnet.state_dict()}) # Build resnet. self.RCNN_base = nn.Sequential(resnet.conv1, resnet.bn1,resnet.relu, resnet.maxpool,resnet.layer1,resnet.layer2,resnet.layer3) self.RCNN_top = nn.Sequential(resnet.layer4) self.RCNN_cls_score = nn.Linear(2048, self.n_classes) if self.class_agnostic: self.RCNN_bbox_pred = nn.Linear(2048, 4) else: self.RCNN_bbox_pred = nn.Linear(2048, 4 * self.n_classes) # Fix blocks for p in self.RCNN_base[0].parameters(): p.requires_grad=False for p in self.RCNN_base[1].parameters(): p.requires_grad=False assert (0 <= cfg.RESNET.FIXED_BLOCKS < 4) if cfg.RESNET.FIXED_BLOCKS >= 3: for p in self.RCNN_base[6].parameters(): p.requires_grad=False if cfg.RESNET.FIXED_BLOCKS >= 2: for p in self.RCNN_base[5].parameters(): p.requires_grad=False if cfg.RESNET.FIXED_BLOCKS >= 1: for p in self.RCNN_base[4].parameters(): p.requires_grad=False def set_bn_fix(m): classname = m.__class__.__name__ if classname.find('BatchNorm') != -1: for p in m.parameters(): p.requires_grad=False self.RCNN_base.apply(set_bn_fix) self.RCNN_top.apply(set_bn_fix) def train(self, mode=True): # Override train so that the training mode is set as we want nn.Module.train(self, mode) if mode: # Set fixed blocks to be in eval mode self.RCNN_base.eval() self.RCNN_base[5].train() self.RCNN_base[6].train() def set_bn_eval(m): classname = m.__class__.__name__ if classname.find('BatchNorm') != -1: m.eval() self.RCNN_base.apply(set_bn_eval) self.RCNN_top.apply(set_bn_eval) def _head_to_tail(self, pool5): fc7 = self.RCNN_top(pool5).mean(3).mean(2) return fc7 ================================================ FILE: lib/model/faster_rcnn/vgg16.py ================================================ # -------------------------------------------------------- # Tensorflow Faster R-CNN # Licensed under The MIT License [see LICENSE for details] # Written by Xinlei Chen # -------------------------------------------------------- from __future__ import absolute_import from __future__ import division from __future__ import print_function import torch import torch.nn as nn import torch.nn.functional as F from torch.autograd import Variable import math import torchvision.models as models from model.faster_rcnn.faster_rcnn import _fasterRCNN import pdb class vgg16(_fasterRCNN): def __init__(self, classes, pretrained=False, class_agnostic=False): self.model_path = 'data/pretrained_model/vgg16_caffe.pth' self.dout_base_model = 512 self.pretrained = pretrained self.class_agnostic = class_agnostic _fasterRCNN.__init__(self, classes, class_agnostic) def _init_modules(self): vgg = models.vgg16() if self.pretrained: print("Loading pretrained weights from %s" %(self.model_path)) state_dict = torch.load(self.model_path) vgg.load_state_dict({k:v for k,v in state_dict.items() if k in vgg.state_dict()}) vgg.classifier = nn.Sequential(*list(vgg.classifier._modules.values())[:-1]) # not using the last maxpool layer self.RCNN_base = nn.Sequential(*list(vgg.features._modules.values())[:-1]) # Fix the layers before conv3: for layer in range(10): for p in self.RCNN_base[layer].parameters(): p.requires_grad = False # self.RCNN_base = _RCNN_base(vgg.features, self.classes, self.dout_base_model) self.RCNN_top = vgg.classifier # not using the last maxpool layer self.RCNN_cls_score = nn.Linear(4096, self.n_classes) if self.class_agnostic: self.RCNN_bbox_pred = nn.Linear(4096, 4) else: self.RCNN_bbox_pred = nn.Linear(4096, 4 * self.n_classes) def _head_to_tail(self, pool5): pool5_flat = pool5.view(pool5.size(0), -1) fc7 = self.RCNN_top(pool5_flat) return fc7 ================================================ FILE: lib/model/nms/.gitignore ================================================ *.c *.cpp *.so ================================================ FILE: lib/model/nms/__init__.py ================================================ ================================================ FILE: lib/model/nms/_ext/__init__.py ================================================ ================================================ FILE: lib/model/nms/_ext/nms/__init__.py ================================================ from torch.utils.ffi import _wrap_function from ._nms import lib as _lib, ffi as _ffi __all__ = [] def _import_symbols(locals): for symbol in dir(_lib): fn = getattr(_lib, symbol) if callable(fn): locals[symbol] = _wrap_function(fn, _ffi) else: locals[symbol] = fn __all__.append(symbol) _import_symbols(locals()) ================================================ FILE: lib/model/nms/build.py ================================================ from __future__ import print_function import os import torch from torch.utils.ffi import create_extension #this_file = os.path.dirname(__file__) sources = [] headers = [] defines = [] with_cuda = False if torch.cuda.is_available(): print('Including CUDA code.') sources += ['src/nms_cuda.c'] headers += ['src/nms_cuda.h'] defines += [('WITH_CUDA', None)] with_cuda = True this_file = os.path.dirname(os.path.realpath(__file__)) print(this_file) extra_objects = ['src/nms_cuda_kernel.cu.o'] extra_objects = [os.path.join(this_file, fname) for fname in extra_objects] print(extra_objects) ffi = create_extension( '_ext.nms', headers=headers, sources=sources, define_macros=defines, relative_to=__file__, with_cuda=with_cuda, extra_objects=extra_objects ) if __name__ == '__main__': ffi.build() ================================================ FILE: lib/model/nms/make.sh ================================================ #!/usr/bin/env bash # CUDA_PATH=/usr/local/cuda/ cd src echo "Compiling stnm kernels by nvcc..." nvcc -c -o nms_cuda_kernel.cu.o nms_cuda_kernel.cu -x cu -Xcompiler -fPIC -arch=sm_52 cd ../ python build.py ================================================ FILE: lib/model/nms/nms_cpu.py ================================================ from __future__ import absolute_import import numpy as np import torch def nms_cpu(dets, thresh): dets = dets.numpy() x1 = dets[:, 0] y1 = dets[:, 1] x2 = dets[:, 2] y2 = dets[:, 3] scores = dets[:, 4] areas = (x2 - x1 + 1) * (y2 - y1 + 1) order = scores.argsort()[::-1] keep = [] while order.size > 0: i = order.item(0) keep.append(i) xx1 = np.maximum(x1[i], x1[order[1:]]) yy1 = np.maximum(y1[i], y1[order[1:]]) xx2 = np.minimum(x2[i], x2[order[1:]]) yy2 = np.minimum(y2[i], y2[order[1:]]) w = np.maximum(0.0, xx2 - xx1 + 1) h = np.maximum(0.0, yy2 - yy1 + 1) inter = w * h ovr = inter / (areas[i] + areas[order[1:]] - inter) inds = np.where(ovr <= thresh)[0] order = order[inds + 1] return torch.IntTensor(keep) ================================================ FILE: lib/model/nms/nms_gpu.py ================================================ from __future__ import absolute_import import torch import numpy as np from ._ext import nms import pdb def nms_gpu(dets, thresh): keep = dets.new(dets.size(0), 1).zero_().int() num_out = dets.new(1).zero_().int() nms.nms_cuda(keep, dets, num_out, thresh) keep = keep[:num_out[0]] return keep ================================================ FILE: lib/model/nms/nms_kernel.cu ================================================ // ------------------------------------------------------------------ // Faster R-CNN // Copyright (c) 2015 Microsoft // Licensed under The MIT License [see fast-rcnn/LICENSE for details] // Written by Shaoqing Ren // ------------------------------------------------------------------ #include "gpu_nms.hpp" #include #include #define CUDA_CHECK(condition) \ /* Code block avoids redefinition of cudaError_t error */ \ do { \ cudaError_t error = condition; \ if (error != cudaSuccess) { \ std::cout << cudaGetErrorString(error) << std::endl; \ } \ } while (0) #define DIVUP(m,n) ((m) / (n) + ((m) % (n) > 0)) int const threadsPerBlock = sizeof(unsigned long long) * 8; __device__ inline float devIoU(float const * const a, float const * const b) { float left = max(a[0], b[0]), right = min(a[2], b[2]); float top = max(a[1], b[1]), bottom = min(a[3], b[3]); float width = max(right - left + 1, 0.f), height = max(bottom - top + 1, 0.f); float interS = width * height; float Sa = (a[2] - a[0] + 1) * (a[3] - a[1] + 1); float Sb = (b[2] - b[0] + 1) * (b[3] - b[1] + 1); return interS / (Sa + Sb - interS); } __global__ void nms_kernel(const int n_boxes, const float nms_overlap_thresh, const float *dev_boxes, unsigned long long *dev_mask) { const int row_start = blockIdx.y; const int col_start = blockIdx.x; // if (row_start > col_start) return; const int row_size = min(n_boxes - row_start * threadsPerBlock, threadsPerBlock); const int col_size = min(n_boxes - col_start * threadsPerBlock, threadsPerBlock); __shared__ float block_boxes[threadsPerBlock * 5]; if (threadIdx.x < col_size) { block_boxes[threadIdx.x * 5 + 0] = dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 0]; block_boxes[threadIdx.x * 5 + 1] = dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 1]; block_boxes[threadIdx.x * 5 + 2] = dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 2]; block_boxes[threadIdx.x * 5 + 3] = dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 3]; block_boxes[threadIdx.x * 5 + 4] = dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 4]; } __syncthreads(); if (threadIdx.x < row_size) { const int cur_box_idx = threadsPerBlock * row_start + threadIdx.x; const float *cur_box = dev_boxes + cur_box_idx * 5; int i = 0; unsigned long long t = 0; int start = 0; if (row_start == col_start) { start = threadIdx.x + 1; } for (i = start; i < col_size; i++) { if (devIoU(cur_box, block_boxes + i * 5) > nms_overlap_thresh) { t |= 1ULL << i; } } const int col_blocks = DIVUP(n_boxes, threadsPerBlock); dev_mask[cur_box_idx * col_blocks + col_start] = t; } } void _set_device(int device_id) { int current_device; CUDA_CHECK(cudaGetDevice(¤t_device)); if (current_device == device_id) { return; } // The call to cudaSetDevice must come before any calls to Get, which // may perform initialization using the GPU. CUDA_CHECK(cudaSetDevice(device_id)); } void _nms(int* keep_out, int* num_out, const float* boxes_host, int boxes_num, int boxes_dim, float nms_overlap_thresh, int device_id) { _set_device(device_id); float* boxes_dev = NULL; unsigned long long* mask_dev = NULL; const int col_blocks = DIVUP(boxes_num, threadsPerBlock); CUDA_CHECK(cudaMalloc(&boxes_dev, boxes_num * boxes_dim * sizeof(float))); CUDA_CHECK(cudaMemcpy(boxes_dev, boxes_host, boxes_num * boxes_dim * sizeof(float), cudaMemcpyHostToDevice)); CUDA_CHECK(cudaMalloc(&mask_dev, boxes_num * col_blocks * sizeof(unsigned long long))); dim3 blocks(DIVUP(boxes_num, threadsPerBlock), DIVUP(boxes_num, threadsPerBlock)); dim3 threads(threadsPerBlock); nms_kernel<<>>(boxes_num, nms_overlap_thresh, boxes_dev, mask_dev); std::vector mask_host(boxes_num * col_blocks); CUDA_CHECK(cudaMemcpy(&mask_host[0], mask_dev, sizeof(unsigned long long) * boxes_num * col_blocks, cudaMemcpyDeviceToHost)); std::vector remv(col_blocks); memset(&remv[0], 0, sizeof(unsigned long long) * col_blocks); int num_to_keep = 0; for (int i = 0; i < boxes_num; i++) { int nblock = i / threadsPerBlock; int inblock = i % threadsPerBlock; if (!(remv[nblock] & (1ULL << inblock))) { keep_out[num_to_keep++] = i; unsigned long long *p = &mask_host[0] + i * col_blocks; for (int j = nblock; j < col_blocks; j++) { remv[j] |= p[j]; } } } *num_out = num_to_keep; CUDA_CHECK(cudaFree(boxes_dev)); CUDA_CHECK(cudaFree(mask_dev)); } ================================================ FILE: lib/model/nms/nms_wrapper.py ================================================ # -------------------------------------------------------- # Fast R-CNN # Copyright (c) 2015 Microsoft # Licensed under The MIT License [see LICENSE for details] # Written by Ross Girshick # -------------------------------------------------------- import torch from model.utils.config import cfg if torch.cuda.is_available(): from model.nms.nms_gpu import nms_gpu from model.nms.nms_cpu import nms_cpu def nms(dets, thresh, force_cpu=False): """Dispatch to either CPU or GPU NMS implementations.""" if dets.shape[0] == 0: return [] # ---numpy version--- # original: return gpu_nms(dets, thresh, device_id=cfg.GPU_ID) # ---pytorch version--- return nms_gpu(dets, thresh) if force_cpu == False else nms_cpu(dets, thresh) ================================================ FILE: lib/model/nms/src/nms_cuda.h ================================================ // int nms_cuda(THCudaTensor *keep_out, THCudaTensor *num_out, // THCudaTensor *boxes_host, THCudaTensor *nms_overlap_thresh); int nms_cuda(THCudaIntTensor *keep_out, THCudaTensor *boxes_host, THCudaIntTensor *num_out, float nms_overlap_thresh); ================================================ FILE: lib/model/nms/src/nms_cuda_kernel.cu ================================================ // ------------------------------------------------------------------ // Faster R-CNN // Copyright (c) 2015 Microsoft // Licensed under The MIT License [see fast-rcnn/LICENSE for details] // Written by Shaoqing Ren // ------------------------------------------------------------------ #include #include #include #include #include "nms_cuda_kernel.h" #define CUDA_WARN(XXX) \ do { if (XXX != cudaSuccess) std::cout << "CUDA Error: " << \ cudaGetErrorString(XXX) << ", at line " << __LINE__ \ << std::endl; cudaDeviceSynchronize(); } while (0) #define CUDA_CHECK(condition) \ /* Code block avoids redefinition of cudaError_t error */ \ do { \ cudaError_t error = condition; \ if (error != cudaSuccess) { \ std::cout << cudaGetErrorString(error) << std::endl; \ } \ } while (0) #define DIVUP(m,n) ((m) / (n) + ((m) % (n) > 0)) int const threadsPerBlock = sizeof(unsigned long long) * 8; __device__ inline float devIoU(float const * const a, float const * const b) { float left = max(a[0], b[0]), right = min(a[2], b[2]); float top = max(a[1], b[1]), bottom = min(a[3], b[3]); float width = max(right - left + 1, 0.f), height = max(bottom - top + 1, 0.f); float interS = width * height; float Sa = (a[2] - a[0] + 1) * (a[3] - a[1] + 1); float Sb = (b[2] - b[0] + 1) * (b[3] - b[1] + 1); return interS / (Sa + Sb - interS); } __global__ void nms_kernel(int n_boxes, float nms_overlap_thresh, float *dev_boxes, unsigned long long *dev_mask) { const int row_start = blockIdx.y; const int col_start = blockIdx.x; // if (row_start > col_start) return; const int row_size = min(n_boxes - row_start * threadsPerBlock, threadsPerBlock); const int col_size = min(n_boxes - col_start * threadsPerBlock, threadsPerBlock); __shared__ float block_boxes[threadsPerBlock * 5]; if (threadIdx.x < col_size) { block_boxes[threadIdx.x * 5 + 0] = dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 0]; block_boxes[threadIdx.x * 5 + 1] = dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 1]; block_boxes[threadIdx.x * 5 + 2] = dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 2]; block_boxes[threadIdx.x * 5 + 3] = dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 3]; block_boxes[threadIdx.x * 5 + 4] = dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 4]; } __syncthreads(); if (threadIdx.x < row_size) { const int cur_box_idx = threadsPerBlock * row_start + threadIdx.x; const float *cur_box = dev_boxes + cur_box_idx * 5; int i = 0; unsigned long long t = 0; int start = 0; if (row_start == col_start) { start = threadIdx.x + 1; } for (i = start; i < col_size; i++) { if (devIoU(cur_box, block_boxes + i * 5) > nms_overlap_thresh) { t |= 1ULL << i; } } const int col_blocks = DIVUP(n_boxes, threadsPerBlock); dev_mask[cur_box_idx * col_blocks + col_start] = t; } } void nms_cuda_compute(int* keep_out, int *num_out, float* boxes_host, int boxes_num, int boxes_dim, float nms_overlap_thresh) { float* boxes_dev = NULL; unsigned long long* mask_dev = NULL; const int col_blocks = DIVUP(boxes_num, threadsPerBlock); CUDA_CHECK(cudaMalloc(&boxes_dev, boxes_num * boxes_dim * sizeof(float))); CUDA_CHECK(cudaMemcpy(boxes_dev, boxes_host, boxes_num * boxes_dim * sizeof(float), cudaMemcpyHostToDevice)); CUDA_CHECK(cudaMalloc(&mask_dev, boxes_num * col_blocks * sizeof(unsigned long long))); dim3 blocks(DIVUP(boxes_num, threadsPerBlock), DIVUP(boxes_num, threadsPerBlock)); dim3 threads(threadsPerBlock); // printf("i am at line %d\n", boxes_num); // printf("i am at line %d\n", boxes_dim); nms_kernel<<>>(boxes_num, nms_overlap_thresh, boxes_dev, mask_dev); std::vector mask_host(boxes_num * col_blocks); CUDA_CHECK(cudaMemcpy(&mask_host[0], mask_dev, sizeof(unsigned long long) * boxes_num * col_blocks, cudaMemcpyDeviceToHost)); std::vector remv(col_blocks); memset(&remv[0], 0, sizeof(unsigned long long) * col_blocks); // we need to create a memory for keep_out on cpu // otherwise, the following code cannot run int* keep_out_cpu = new int[boxes_num]; int num_to_keep = 0; for (int i = 0; i < boxes_num; i++) { int nblock = i / threadsPerBlock; int inblock = i % threadsPerBlock; if (!(remv[nblock] & (1ULL << inblock))) { // orignal: keep_out[num_to_keep++] = i; keep_out_cpu[num_to_keep++] = i; unsigned long long *p = &mask_host[0] + i * col_blocks; for (int j = nblock; j < col_blocks; j++) { remv[j] |= p[j]; } } } // copy keep_out_cpu to keep_out on gpu CUDA_WARN(cudaMemcpy(keep_out, keep_out_cpu, boxes_num * sizeof(int),cudaMemcpyHostToDevice)); // *num_out = num_to_keep; // original: *num_out = num_to_keep; // copy num_to_keep to num_out on gpu CUDA_WARN(cudaMemcpy(num_out, &num_to_keep, 1 * sizeof(int),cudaMemcpyHostToDevice)); // release cuda memory CUDA_CHECK(cudaFree(boxes_dev)); CUDA_CHECK(cudaFree(mask_dev)); // release cpu memory delete []keep_out_cpu; } ================================================ FILE: lib/model/nms/src/nms_cuda_kernel.h ================================================ #ifdef __cplusplus extern "C" { #endif void nms_cuda_compute(int* keep_out, int *num_out, float* boxes_host, int boxes_num, int boxes_dim, float nms_overlap_thresh); #ifdef __cplusplus } #endif ================================================ FILE: lib/model/roi_align/__init__.py ================================================ ================================================ FILE: lib/model/roi_align/_ext/__init__.py ================================================ ================================================ FILE: lib/model/roi_align/_ext/roi_align/__init__.py ================================================ from torch.utils.ffi import _wrap_function from ._roi_align import lib as _lib, ffi as _ffi __all__ = [] def _import_symbols(locals): for symbol in dir(_lib): fn = getattr(_lib, symbol) if callable(fn): locals[symbol] = _wrap_function(fn, _ffi) else: locals[symbol] = fn __all__.append(symbol) _import_symbols(locals()) ================================================ FILE: lib/model/roi_align/build.py ================================================ from __future__ import print_function import os import torch from torch.utils.ffi import create_extension sources = ['src/roi_align.c'] headers = ['src/roi_align.h'] extra_objects = [] #sources = [] #headers = [] defines = [] with_cuda = False this_file = os.path.dirname(os.path.realpath(__file__)) print(this_file) if torch.cuda.is_available(): print('Including CUDA code.') sources += ['src/roi_align_cuda.c'] headers += ['src/roi_align_cuda.h'] defines += [('WITH_CUDA', None)] with_cuda = True extra_objects = ['src/roi_align_kernel.cu.o'] extra_objects = [os.path.join(this_file, fname) for fname in extra_objects] ffi = create_extension( '_ext.roi_align', headers=headers, sources=sources, define_macros=defines, relative_to=__file__, with_cuda=with_cuda, extra_objects=extra_objects ) if __name__ == '__main__': ffi.build() ================================================ FILE: lib/model/roi_align/functions/__init__.py ================================================ ================================================ FILE: lib/model/roi_align/functions/roi_align.py ================================================ import torch from torch.autograd import Function from .._ext import roi_align # TODO use save_for_backward instead class RoIAlignFunction(Function): def __init__(self, aligned_height, aligned_width, spatial_scale): self.aligned_width = int(aligned_width) self.aligned_height = int(aligned_height) self.spatial_scale = float(spatial_scale) self.rois = None self.feature_size = None def forward(self, features, rois): self.rois = rois self.feature_size = features.size() batch_size, num_channels, data_height, data_width = features.size() num_rois = rois.size(0) output = features.new(num_rois, num_channels, self.aligned_height, self.aligned_width).zero_() if features.is_cuda: roi_align.roi_align_forward_cuda(self.aligned_height, self.aligned_width, self.spatial_scale, features, rois, output) else: roi_align.roi_align_forward(self.aligned_height, self.aligned_width, self.spatial_scale, features, rois, output) # raise NotImplementedError return output def backward(self, grad_output): assert(self.feature_size is not None and grad_output.is_cuda) batch_size, num_channels, data_height, data_width = self.feature_size grad_input = self.rois.new(batch_size, num_channels, data_height, data_width).zero_() roi_align.roi_align_backward_cuda(self.aligned_height, self.aligned_width, self.spatial_scale, grad_output, self.rois, grad_input) # print grad_input return grad_input, None ================================================ FILE: lib/model/roi_align/make.sh ================================================ #!/usr/bin/env bash CUDA_PATH=/usr/local/cuda/ cd src echo "Compiling my_lib kernels by nvcc..." nvcc -c -o roi_align_kernel.cu.o roi_align_kernel.cu -x cu -Xcompiler -fPIC -arch=sm_52 cd ../ python build.py ================================================ FILE: lib/model/roi_align/modules/__init__.py ================================================ ================================================ FILE: lib/model/roi_align/modules/roi_align.py ================================================ from torch.nn.modules.module import Module from torch.nn.functional import avg_pool2d, max_pool2d from ..functions.roi_align import RoIAlignFunction class RoIAlign(Module): def __init__(self, aligned_height, aligned_width, spatial_scale): super(RoIAlign, self).__init__() self.aligned_width = int(aligned_width) self.aligned_height = int(aligned_height) self.spatial_scale = float(spatial_scale) def forward(self, features, rois): return RoIAlignFunction(self.aligned_height, self.aligned_width, self.spatial_scale)(features, rois) class RoIAlignAvg(Module): def __init__(self, aligned_height, aligned_width, spatial_scale): super(RoIAlignAvg, self).__init__() self.aligned_width = int(aligned_width) self.aligned_height = int(aligned_height) self.spatial_scale = float(spatial_scale) def forward(self, features, rois): x = RoIAlignFunction(self.aligned_height+1, self.aligned_width+1, self.spatial_scale)(features, rois) return avg_pool2d(x, kernel_size=2, stride=1) class RoIAlignMax(Module): def __init__(self, aligned_height, aligned_width, spatial_scale): super(RoIAlignMax, self).__init__() self.aligned_width = int(aligned_width) self.aligned_height = int(aligned_height) self.spatial_scale = float(spatial_scale) def forward(self, features, rois): x = RoIAlignFunction(self.aligned_height+1, self.aligned_width+1, self.spatial_scale)(features, rois) return max_pool2d(x, kernel_size=2, stride=1) ================================================ FILE: lib/model/roi_align/src/roi_align.c ================================================ #include #include #include void ROIAlignForwardCpu(const float* bottom_data, const float spatial_scale, const int num_rois, const int height, const int width, const int channels, const int aligned_height, const int aligned_width, const float * bottom_rois, float* top_data); void ROIAlignBackwardCpu(const float* top_diff, const float spatial_scale, const int num_rois, const int height, const int width, const int channels, const int aligned_height, const int aligned_width, const float * bottom_rois, float* top_data); int roi_align_forward(int aligned_height, int aligned_width, float spatial_scale, THFloatTensor * features, THFloatTensor * rois, THFloatTensor * output) { //Grab the input tensor float * data_flat = THFloatTensor_data(features); float * rois_flat = THFloatTensor_data(rois); float * output_flat = THFloatTensor_data(output); // Number of ROIs int num_rois = THFloatTensor_size(rois, 0); int size_rois = THFloatTensor_size(rois, 1); if (size_rois != 5) { return 0; } // data height int data_height = THFloatTensor_size(features, 2); // data width int data_width = THFloatTensor_size(features, 3); // Number of channels int num_channels = THFloatTensor_size(features, 1); // do ROIAlignForward ROIAlignForwardCpu(data_flat, spatial_scale, num_rois, data_height, data_width, num_channels, aligned_height, aligned_width, rois_flat, output_flat); return 1; } int roi_align_backward(int aligned_height, int aligned_width, float spatial_scale, THFloatTensor * top_grad, THFloatTensor * rois, THFloatTensor * bottom_grad) { //Grab the input tensor float * top_grad_flat = THFloatTensor_data(top_grad); float * rois_flat = THFloatTensor_data(rois); float * bottom_grad_flat = THFloatTensor_data(bottom_grad); // Number of ROIs int num_rois = THFloatTensor_size(rois, 0); int size_rois = THFloatTensor_size(rois, 1); if (size_rois != 5) { return 0; } // batch size // int batch_size = THFloatTensor_size(bottom_grad, 0); // data height int data_height = THFloatTensor_size(bottom_grad, 2); // data width int data_width = THFloatTensor_size(bottom_grad, 3); // Number of channels int num_channels = THFloatTensor_size(bottom_grad, 1); // do ROIAlignBackward ROIAlignBackwardCpu(top_grad_flat, spatial_scale, num_rois, data_height, data_width, num_channels, aligned_height, aligned_width, rois_flat, bottom_grad_flat); return 1; } void ROIAlignForwardCpu(const float* bottom_data, const float spatial_scale, const int num_rois, const int height, const int width, const int channels, const int aligned_height, const int aligned_width, const float * bottom_rois, float* top_data) { const int output_size = num_rois * aligned_height * aligned_width * channels; int idx = 0; for (idx = 0; idx < output_size; ++idx) { // (n, c, ph, pw) is an element in the aligned output int pw = idx % aligned_width; int ph = (idx / aligned_width) % aligned_height; int c = (idx / aligned_width / aligned_height) % channels; int n = idx / aligned_width / aligned_height / channels; float roi_batch_ind = bottom_rois[n * 5 + 0]; float roi_start_w = bottom_rois[n * 5 + 1] * spatial_scale; float roi_start_h = bottom_rois[n * 5 + 2] * spatial_scale; float roi_end_w = bottom_rois[n * 5 + 3] * spatial_scale; float roi_end_h = bottom_rois[n * 5 + 4] * spatial_scale; // Force malformed ROI to be 1x1 float roi_width = fmaxf(roi_end_w - roi_start_w + 1., 0.); float roi_height = fmaxf(roi_end_h - roi_start_h + 1., 0.); float bin_size_h = roi_height / (aligned_height - 1.); float bin_size_w = roi_width / (aligned_width - 1.); float h = (float)(ph) * bin_size_h + roi_start_h; float w = (float)(pw) * bin_size_w + roi_start_w; int hstart = fminf(floor(h), height - 2); int wstart = fminf(floor(w), width - 2); int img_start = roi_batch_ind * channels * height * width; // bilinear interpolation if (h < 0 || h >= height || w < 0 || w >= width) { top_data[idx] = 0.; } else { float h_ratio = h - (float)(hstart); float w_ratio = w - (float)(wstart); int upleft = img_start + (c * height + hstart) * width + wstart; int upright = upleft + 1; int downleft = upleft + width; int downright = downleft + 1; top_data[idx] = bottom_data[upleft] * (1. - h_ratio) * (1. - w_ratio) + bottom_data[upright] * (1. - h_ratio) * w_ratio + bottom_data[downleft] * h_ratio * (1. - w_ratio) + bottom_data[downright] * h_ratio * w_ratio; } } } void ROIAlignBackwardCpu(const float* top_diff, const float spatial_scale, const int num_rois, const int height, const int width, const int channels, const int aligned_height, const int aligned_width, const float * bottom_rois, float* bottom_diff) { const int output_size = num_rois * aligned_height * aligned_width * channels; int idx = 0; for (idx = 0; idx < output_size; ++idx) { // (n, c, ph, pw) is an element in the aligned output int pw = idx % aligned_width; int ph = (idx / aligned_width) % aligned_height; int c = (idx / aligned_width / aligned_height) % channels; int n = idx / aligned_width / aligned_height / channels; float roi_batch_ind = bottom_rois[n * 5 + 0]; float roi_start_w = bottom_rois[n * 5 + 1] * spatial_scale; float roi_start_h = bottom_rois[n * 5 + 2] * spatial_scale; float roi_end_w = bottom_rois[n * 5 + 3] * spatial_scale; float roi_end_h = bottom_rois[n * 5 + 4] * spatial_scale; // Force malformed ROI to be 1x1 float roi_width = fmaxf(roi_end_w - roi_start_w + 1., 0.); float roi_height = fmaxf(roi_end_h - roi_start_h + 1., 0.); float bin_size_h = roi_height / (aligned_height - 1.); float bin_size_w = roi_width / (aligned_width - 1.); float h = (float)(ph) * bin_size_h + roi_start_h; float w = (float)(pw) * bin_size_w + roi_start_w; int hstart = fminf(floor(h), height - 2); int wstart = fminf(floor(w), width - 2); int img_start = roi_batch_ind * channels * height * width; // bilinear interpolation if (h < 0 || h >= height || w < 0 || w >= width) { float h_ratio = h - (float)(hstart); float w_ratio = w - (float)(wstart); int upleft = img_start + (c * height + hstart) * width + wstart; int upright = upleft + 1; int downleft = upleft + width; int downright = downleft + 1; bottom_diff[upleft] += top_diff[idx] * (1. - h_ratio) * (1. - w_ratio); bottom_diff[upright] += top_diff[idx] * (1. - h_ratio) * w_ratio; bottom_diff[downleft] += top_diff[idx] * h_ratio * (1. - w_ratio); bottom_diff[downright] += top_diff[idx] * h_ratio * w_ratio; } } } ================================================ FILE: lib/model/roi_align/src/roi_align.h ================================================ int roi_align_forward(int aligned_height, int aligned_width, float spatial_scale, THFloatTensor * features, THFloatTensor * rois, THFloatTensor * output); int roi_align_backward(int aligned_height, int aligned_width, float spatial_scale, THFloatTensor * top_grad, THFloatTensor * rois, THFloatTensor * bottom_grad); ================================================ FILE: lib/model/roi_align/src/roi_align_cuda.c ================================================ #include #include #include "roi_align_kernel.h" extern THCState *state; int roi_align_forward_cuda(int aligned_height, int aligned_width, float spatial_scale, THCudaTensor * features, THCudaTensor * rois, THCudaTensor * output) { // Grab the input tensor float * data_flat = THCudaTensor_data(state, features); float * rois_flat = THCudaTensor_data(state, rois); float * output_flat = THCudaTensor_data(state, output); // Number of ROIs int num_rois = THCudaTensor_size(state, rois, 0); int size_rois = THCudaTensor_size(state, rois, 1); if (size_rois != 5) { return 0; } // data height int data_height = THCudaTensor_size(state, features, 2); // data width int data_width = THCudaTensor_size(state, features, 3); // Number of channels int num_channels = THCudaTensor_size(state, features, 1); cudaStream_t stream = THCState_getCurrentStream(state); ROIAlignForwardLaucher( data_flat, spatial_scale, num_rois, data_height, data_width, num_channels, aligned_height, aligned_width, rois_flat, output_flat, stream); return 1; } int roi_align_backward_cuda(int aligned_height, int aligned_width, float spatial_scale, THCudaTensor * top_grad, THCudaTensor * rois, THCudaTensor * bottom_grad) { // Grab the input tensor float * top_grad_flat = THCudaTensor_data(state, top_grad); float * rois_flat = THCudaTensor_data(state, rois); float * bottom_grad_flat = THCudaTensor_data(state, bottom_grad); // Number of ROIs int num_rois = THCudaTensor_size(state, rois, 0); int size_rois = THCudaTensor_size(state, rois, 1); if (size_rois != 5) { return 0; } // batch size int batch_size = THCudaTensor_size(state, bottom_grad, 0); // data height int data_height = THCudaTensor_size(state, bottom_grad, 2); // data width int data_width = THCudaTensor_size(state, bottom_grad, 3); // Number of channels int num_channels = THCudaTensor_size(state, bottom_grad, 1); cudaStream_t stream = THCState_getCurrentStream(state); ROIAlignBackwardLaucher( top_grad_flat, spatial_scale, batch_size, num_rois, data_height, data_width, num_channels, aligned_height, aligned_width, rois_flat, bottom_grad_flat, stream); return 1; } ================================================ FILE: lib/model/roi_align/src/roi_align_cuda.h ================================================ int roi_align_forward_cuda(int aligned_height, int aligned_width, float spatial_scale, THCudaTensor * features, THCudaTensor * rois, THCudaTensor * output); int roi_align_backward_cuda(int aligned_height, int aligned_width, float spatial_scale, THCudaTensor * top_grad, THCudaTensor * rois, THCudaTensor * bottom_grad); ================================================ FILE: lib/model/roi_align/src/roi_align_kernel.cu ================================================ #ifdef __cplusplus extern "C" { #endif #include #include #include #include "roi_align_kernel.h" #define CUDA_1D_KERNEL_LOOP(i, n) \ for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; \ i += blockDim.x * gridDim.x) __global__ void ROIAlignForward(const int nthreads, const float* bottom_data, const float spatial_scale, const int height, const int width, const int channels, const int aligned_height, const int aligned_width, const float* bottom_rois, float* top_data) { CUDA_1D_KERNEL_LOOP(index, nthreads) { // (n, c, ph, pw) is an element in the aligned output // int n = index; // int pw = n % aligned_width; // n /= aligned_width; // int ph = n % aligned_height; // n /= aligned_height; // int c = n % channels; // n /= channels; int pw = index % aligned_width; int ph = (index / aligned_width) % aligned_height; int c = (index / aligned_width / aligned_height) % channels; int n = index / aligned_width / aligned_height / channels; // bottom_rois += n * 5; float roi_batch_ind = bottom_rois[n * 5 + 0]; float roi_start_w = bottom_rois[n * 5 + 1] * spatial_scale; float roi_start_h = bottom_rois[n * 5 + 2] * spatial_scale; float roi_end_w = bottom_rois[n * 5 + 3] * spatial_scale; float roi_end_h = bottom_rois[n * 5 + 4] * spatial_scale; // Force malformed ROIs to be 1x1 float roi_width = fmaxf(roi_end_w - roi_start_w + 1., 0.); float roi_height = fmaxf(roi_end_h - roi_start_h + 1., 0.); float bin_size_h = roi_height / (aligned_height - 1.); float bin_size_w = roi_width / (aligned_width - 1.); float h = (float)(ph) * bin_size_h + roi_start_h; float w = (float)(pw) * bin_size_w + roi_start_w; int hstart = fminf(floor(h), height - 2); int wstart = fminf(floor(w), width - 2); int img_start = roi_batch_ind * channels * height * width; // bilinear interpolation if (h < 0 || h >= height || w < 0 || w >= width) { top_data[index] = 0.; } else { float h_ratio = h - (float)(hstart); float w_ratio = w - (float)(wstart); int upleft = img_start + (c * height + hstart) * width + wstart; int upright = upleft + 1; int downleft = upleft + width; int downright = downleft + 1; top_data[index] = bottom_data[upleft] * (1. - h_ratio) * (1. - w_ratio) + bottom_data[upright] * (1. - h_ratio) * w_ratio + bottom_data[downleft] * h_ratio * (1. - w_ratio) + bottom_data[downright] * h_ratio * w_ratio; } } } int ROIAlignForwardLaucher(const float* bottom_data, const float spatial_scale, const int num_rois, const int height, const int width, const int channels, const int aligned_height, const int aligned_width, const float* bottom_rois, float* top_data, cudaStream_t stream) { const int kThreadsPerBlock = 1024; const int output_size = num_rois * aligned_height * aligned_width * channels; cudaError_t err; ROIAlignForward<<<(output_size + kThreadsPerBlock - 1) / kThreadsPerBlock, kThreadsPerBlock, 0, stream>>>( output_size, bottom_data, spatial_scale, height, width, channels, aligned_height, aligned_width, bottom_rois, top_data); err = cudaGetLastError(); if(cudaSuccess != err) { fprintf( stderr, "cudaCheckError() failed : %s\n", cudaGetErrorString( err ) ); exit( -1 ); } return 1; } __global__ void ROIAlignBackward(const int nthreads, const float* top_diff, const float spatial_scale, const int height, const int width, const int channels, const int aligned_height, const int aligned_width, float* bottom_diff, const float* bottom_rois) { CUDA_1D_KERNEL_LOOP(index, nthreads) { // (n, c, ph, pw) is an element in the aligned output int pw = index % aligned_width; int ph = (index / aligned_width) % aligned_height; int c = (index / aligned_width / aligned_height) % channels; int n = index / aligned_width / aligned_height / channels; float roi_batch_ind = bottom_rois[n * 5 + 0]; float roi_start_w = bottom_rois[n * 5 + 1] * spatial_scale; float roi_start_h = bottom_rois[n * 5 + 2] * spatial_scale; float roi_end_w = bottom_rois[n * 5 + 3] * spatial_scale; float roi_end_h = bottom_rois[n * 5 + 4] * spatial_scale; /* int roi_start_w = round(bottom_rois[1] * spatial_scale); */ /* int roi_start_h = round(bottom_rois[2] * spatial_scale); */ /* int roi_end_w = round(bottom_rois[3] * spatial_scale); */ /* int roi_end_h = round(bottom_rois[4] * spatial_scale); */ // Force malformed ROIs to be 1x1 float roi_width = fmaxf(roi_end_w - roi_start_w + 1., 0.); float roi_height = fmaxf(roi_end_h - roi_start_h + 1., 0.); float bin_size_h = roi_height / (aligned_height - 1.); float bin_size_w = roi_width / (aligned_width - 1.); float h = (float)(ph) * bin_size_h + roi_start_h; float w = (float)(pw) * bin_size_w + roi_start_w; int hstart = fminf(floor(h), height - 2); int wstart = fminf(floor(w), width - 2); int img_start = roi_batch_ind * channels * height * width; // bilinear interpolation if (!(h < 0 || h >= height || w < 0 || w >= width)) { float h_ratio = h - (float)(hstart); float w_ratio = w - (float)(wstart); int upleft = img_start + (c * height + hstart) * width + wstart; int upright = upleft + 1; int downleft = upleft + width; int downright = downleft + 1; atomicAdd(bottom_diff + upleft, top_diff[index] * (1. - h_ratio) * (1 - w_ratio)); atomicAdd(bottom_diff + upright, top_diff[index] * (1. - h_ratio) * w_ratio); atomicAdd(bottom_diff + downleft, top_diff[index] * h_ratio * (1 - w_ratio)); atomicAdd(bottom_diff + downright, top_diff[index] * h_ratio * w_ratio); } } } int ROIAlignBackwardLaucher(const float* top_diff, const float spatial_scale, const int batch_size, const int num_rois, const int height, const int width, const int channels, const int aligned_height, const int aligned_width, const float* bottom_rois, float* bottom_diff, cudaStream_t stream) { const int kThreadsPerBlock = 1024; const int output_size = num_rois * aligned_height * aligned_width * channels; cudaError_t err; ROIAlignBackward<<<(output_size + kThreadsPerBlock - 1) / kThreadsPerBlock, kThreadsPerBlock, 0, stream>>>( output_size, top_diff, spatial_scale, height, width, channels, aligned_height, aligned_width, bottom_diff, bottom_rois); err = cudaGetLastError(); if(cudaSuccess != err) { fprintf( stderr, "cudaCheckError() failed : %s\n", cudaGetErrorString( err ) ); exit( -1 ); } return 1; } #ifdef __cplusplus } #endif ================================================ FILE: lib/model/roi_align/src/roi_align_kernel.h ================================================ #ifndef _ROI_ALIGN_KERNEL #define _ROI_ALIGN_KERNEL #ifdef __cplusplus extern "C" { #endif __global__ void ROIAlignForward(const int nthreads, const float* bottom_data, const float spatial_scale, const int height, const int width, const int channels, const int aligned_height, const int aligned_width, const float* bottom_rois, float* top_data); int ROIAlignForwardLaucher( const float* bottom_data, const float spatial_scale, const int num_rois, const int height, const int width, const int channels, const int aligned_height, const int aligned_width, const float* bottom_rois, float* top_data, cudaStream_t stream); __global__ void ROIAlignBackward(const int nthreads, const float* top_diff, const float spatial_scale, const int height, const int width, const int channels, const int aligned_height, const int aligned_width, float* bottom_diff, const float* bottom_rois); int ROIAlignBackwardLaucher(const float* top_diff, const float spatial_scale, const int batch_size, const int num_rois, const int height, const int width, const int channels, const int aligned_height, const int aligned_width, const float* bottom_rois, float* bottom_diff, cudaStream_t stream); #ifdef __cplusplus } #endif #endif ================================================ FILE: lib/model/roi_crop/__init__.py ================================================ ================================================ FILE: lib/model/roi_crop/_ext/__init__.py ================================================ ================================================ FILE: lib/model/roi_crop/_ext/crop_resize/__init__.py ================================================ from torch.utils.ffi import _wrap_function from ._crop_resize import lib as _lib, ffi as _ffi __all__ = [] def _import_symbols(locals): for symbol in dir(_lib): fn = getattr(_lib, symbol) locals[symbol] = _wrap_function(fn, _ffi) __all__.append(symbol) _import_symbols(locals()) ================================================ FILE: lib/model/roi_crop/_ext/roi_crop/__init__.py ================================================ from torch.utils.ffi import _wrap_function from ._roi_crop import lib as _lib, ffi as _ffi __all__ = [] def _import_symbols(locals): for symbol in dir(_lib): fn = getattr(_lib, symbol) if callable(fn): locals[symbol] = _wrap_function(fn, _ffi) else: locals[symbol] = fn __all__.append(symbol) _import_symbols(locals()) ================================================ FILE: lib/model/roi_crop/build.py ================================================ from __future__ import print_function import os import torch from torch.utils.ffi import create_extension #this_file = os.path.dirname(__file__) sources = ['src/roi_crop.c'] headers = ['src/roi_crop.h'] defines = [] with_cuda = False if torch.cuda.is_available(): print('Including CUDA code.') sources += ['src/roi_crop_cuda.c'] headers += ['src/roi_crop_cuda.h'] defines += [('WITH_CUDA', None)] with_cuda = True this_file = os.path.dirname(os.path.realpath(__file__)) print(this_file) extra_objects = ['src/roi_crop_cuda_kernel.cu.o'] extra_objects = [os.path.join(this_file, fname) for fname in extra_objects] ffi = create_extension( '_ext.roi_crop', headers=headers, sources=sources, define_macros=defines, relative_to=__file__, with_cuda=with_cuda, extra_objects=extra_objects ) if __name__ == '__main__': ffi.build() ================================================ FILE: lib/model/roi_crop/functions/__init__.py ================================================ ================================================ FILE: lib/model/roi_crop/functions/crop_resize.py ================================================ # functions/add.py import torch from torch.autograd import Function from .._ext import roi_crop from cffi import FFI ffi = FFI() class RoICropFunction(Function): def forward(self, input1, input2): self.input1 = input1 self.input2 = input2 self.device_c = ffi.new("int *") output = torch.zeros(input2.size()[0], input1.size()[1], input2.size()[1], input2.size()[2]) #print('decice %d' % torch.cuda.current_device()) if input1.is_cuda: self.device = torch.cuda.current_device() else: self.device = -1 self.device_c[0] = self.device if not input1.is_cuda: roi_crop.BilinearSamplerBHWD_updateOutput(input1, input2, output) else: output = output.cuda(self.device) roi_crop.BilinearSamplerBHWD_updateOutput_cuda(input1, input2, output) return output def backward(self, grad_output): grad_input1 = torch.zeros(self.input1.size()) grad_input2 = torch.zeros(self.input2.size()) #print('backward decice %d' % self.device) if not grad_output.is_cuda: roi_crop.BilinearSamplerBHWD_updateGradInput(self.input1, self.input2, grad_input1, grad_input2, grad_output) else: grad_input1 = grad_input1.cuda(self.device) grad_input2 = grad_input2.cuda(self.device) roi_crop.BilinearSamplerBHWD_updateGradInput_cuda(self.input1, self.input2, grad_input1, grad_input2, grad_output) return grad_input1, grad_input2 ================================================ FILE: lib/model/roi_crop/functions/gridgen.py ================================================ # functions/add.py import torch from torch.autograd import Function import numpy as np class AffineGridGenFunction(Function): def __init__(self, height, width,lr=1): super(AffineGridGenFunction, self).__init__() self.lr = lr self.height, self.width = height, width self.grid = np.zeros( [self.height, self.width, 3], dtype=np.float32) self.grid[:,:,0] = np.expand_dims(np.repeat(np.expand_dims(np.arange(-1, 1, 2.0/(self.height)), 0), repeats = self.width, axis = 0).T, 0) self.grid[:,:,1] = np.expand_dims(np.repeat(np.expand_dims(np.arange(-1, 1, 2.0/(self.width)), 0), repeats = self.height, axis = 0), 0) # self.grid[:,:,0] = np.expand_dims(np.repeat(np.expand_dims(np.arange(-1, 1, 2.0/(self.height - 1)), 0), repeats = self.width, axis = 0).T, 0) # self.grid[:,:,1] = np.expand_dims(np.repeat(np.expand_dims(np.arange(-1, 1, 2.0/(self.width - 1)), 0), repeats = self.height, axis = 0), 0) self.grid[:,:,2] = np.ones([self.height, width]) self.grid = torch.from_numpy(self.grid.astype(np.float32)) #print(self.grid) def forward(self, input1): self.input1 = input1 output = input1.new(torch.Size([input1.size(0)]) + self.grid.size()).zero_() self.batchgrid = input1.new(torch.Size([input1.size(0)]) + self.grid.size()).zero_() for i in range(input1.size(0)): self.batchgrid[i] = self.grid.astype(self.batchgrid[i]) # if input1.is_cuda: # self.batchgrid = self.batchgrid.cuda() # output = output.cuda() for i in range(input1.size(0)): output = torch.bmm(self.batchgrid.view(-1, self.height*self.width, 3), torch.transpose(input1, 1, 2)).view(-1, self.height, self.width, 2) return output def backward(self, grad_output): grad_input1 = self.input1.new(self.input1.size()).zero_() # if grad_output.is_cuda: # self.batchgrid = self.batchgrid.cuda() # grad_input1 = grad_input1.cuda() grad_input1 = torch.baddbmm(grad_input1, torch.transpose(grad_output.view(-1, self.height*self.width, 2), 1,2), self.batchgrid.view(-1, self.height*self.width, 3)) return grad_input1 ================================================ FILE: lib/model/roi_crop/functions/roi_crop.py ================================================ # functions/add.py import torch from torch.autograd import Function from .._ext import roi_crop import pdb class RoICropFunction(Function): def forward(self, input1, input2): self.input1 = input1.clone() self.input2 = input2.clone() output = input2.new(input2.size()[0], input1.size()[1], input2.size()[1], input2.size()[2]).zero_() assert output.get_device() == input1.get_device(), "output and input1 must on the same device" assert output.get_device() == input2.get_device(), "output and input2 must on the same device" roi_crop.BilinearSamplerBHWD_updateOutput_cuda(input1, input2, output) return output def backward(self, grad_output): grad_input1 = self.input1.new(self.input1.size()).zero_() grad_input2 = self.input2.new(self.input2.size()).zero_() roi_crop.BilinearSamplerBHWD_updateGradInput_cuda(self.input1, self.input2, grad_input1, grad_input2, grad_output) return grad_input1, grad_input2 ================================================ FILE: lib/model/roi_crop/make.sh ================================================ #!/usr/bin/env bash CUDA_PATH=/usr/local/cuda/ cd src echo "Compiling my_lib kernels by nvcc..." nvcc -c -o roi_crop_cuda_kernel.cu.o roi_crop_cuda_kernel.cu -x cu -Xcompiler -fPIC -arch=sm_52 cd ../ python build.py ================================================ FILE: lib/model/roi_crop/modules/__init__.py ================================================ ================================================ FILE: lib/model/roi_crop/modules/gridgen.py ================================================ from torch.nn.modules.module import Module import torch from torch.autograd import Variable import numpy as np from ..functions.gridgen import AffineGridGenFunction import pyximport pyximport.install(setup_args={"include_dirs":np.get_include()}, reload_support=True) class _AffineGridGen(Module): def __init__(self, height, width, lr = 1, aux_loss = False): super(_AffineGridGen, self).__init__() self.height, self.width = height, width self.aux_loss = aux_loss self.f = AffineGridGenFunction(self.height, self.width, lr=lr) self.lr = lr def forward(self, input): # if not self.aux_loss: return self.f(input) # else: # identity = torch.from_numpy(np.array([[1,0,0], [0,1,0]], dtype=np.float32)) # batch_identity = torch.zeros([input.size(0), 2,3]) # for i in range(input.size(0)): # batch_identity[i] = identity # batch_identity = Variable(batch_identity) # loss = torch.mul(input - batch_identity, input - batch_identity) # loss = torch.sum(loss,1) # loss = torch.sum(loss,2) # return self.f(input), loss.view(-1,1) class CylinderGridGen(Module): def __init__(self, height, width, lr = 1, aux_loss = False): super(CylinderGridGen, self).__init__() self.height, self.width = height, width self.aux_loss = aux_loss self.f = CylinderGridGenFunction(self.height, self.width, lr=lr) self.lr = lr def forward(self, input): if not self.aux_loss: return self.f(input) else: return self.f(input), torch.mul(input, input).view(-1,1) class AffineGridGenV2(Module): def __init__(self, height, width, lr = 1, aux_loss = False): super(AffineGridGenV2, self).__init__() self.height, self.width = height, width self.aux_loss = aux_loss self.lr = lr self.grid = np.zeros( [self.height, self.width, 3], dtype=np.float32) self.grid[:,:,0] = np.expand_dims(np.repeat(np.expand_dims(np.arange(-1, 1, 2.0/self.height), 0), repeats = self.width, axis = 0).T, 0) self.grid[:,:,1] = np.expand_dims(np.repeat(np.expand_dims(np.arange(-1, 1, 2.0/self.width), 0), repeats = self.height, axis = 0), 0) self.grid[:,:,2] = np.ones([self.height, width]) self.grid = torch.from_numpy(self.grid.astype(np.float32)) def forward(self, input1): self.batchgrid = torch.zeros(torch.Size([input1.size(0)]) + self.grid.size()) for i in range(input1.size(0)): self.batchgrid[i] = self.grid self.batchgrid = Variable(self.batchgrid) if input1.is_cuda: self.batchgrid = self.batchgrid.cuda() output = torch.bmm(self.batchgrid.view(-1, self.height*self.width, 3), torch.transpose(input1, 1, 2)).view(-1, self.height, self.width, 2) return output class CylinderGridGenV2(Module): def __init__(self, height, width, lr = 1): super(CylinderGridGenV2, self).__init__() self.height, self.width = height, width self.lr = lr self.grid = np.zeros( [self.height, self.width, 3], dtype=np.float32) self.grid[:,:,0] = np.expand_dims(np.repeat(np.expand_dims(np.arange(-1, 1, 2.0/self.height), 0), repeats = self.width, axis = 0).T, 0) self.grid[:,:,1] = np.expand_dims(np.repeat(np.expand_dims(np.arange(-1, 1, 2.0/self.width), 0), repeats = self.height, axis = 0), 0) self.grid[:,:,2] = np.ones([self.height, width]) self.grid = torch.from_numpy(self.grid.astype(np.float32)) def forward(self, input): self.batchgrid = torch.zeros(torch.Size([input.size(0)]) + self.grid.size() ) #print(self.batchgrid.size()) for i in range(input.size(0)): self.batchgrid[i,:,:,:] = self.grid self.batchgrid = Variable(self.batchgrid) #print(self.batchgrid.size()) input_u = input.view(-1,1,1,1).repeat(1,self.height, self.width,1) #print(input_u.requires_grad, self.batchgrid) output0 = self.batchgrid[:,:,:,0:1] output1 = torch.atan(torch.tan(np.pi/2.0*(self.batchgrid[:,:,:,1:2] + self.batchgrid[:,:,:,2:] * input_u[:,:,:,:]))) /(np.pi/2) #print(output0.size(), output1.size()) output = torch.cat([output0, output1], 3) return output class DenseAffineGridGen(Module): def __init__(self, height, width, lr = 1, aux_loss = False): super(DenseAffineGridGen, self).__init__() self.height, self.width = height, width self.aux_loss = aux_loss self.lr = lr self.grid = np.zeros( [self.height, self.width, 3], dtype=np.float32) self.grid[:,:,0] = np.expand_dims(np.repeat(np.expand_dims(np.arange(-1, 1, 2.0/self.height), 0), repeats = self.width, axis = 0).T, 0) self.grid[:,:,1] = np.expand_dims(np.repeat(np.expand_dims(np.arange(-1, 1, 2.0/self.width), 0), repeats = self.height, axis = 0), 0) self.grid[:,:,2] = np.ones([self.height, width]) self.grid = torch.from_numpy(self.grid.astype(np.float32)) def forward(self, input1): self.batchgrid = torch.zeros(torch.Size([input1.size(0)]) + self.grid.size()) for i in range(input1.size(0)): self.batchgrid[i] = self.grid self.batchgrid = Variable(self.batchgrid) #print self.batchgrid, input1[:,:,:,0:3] #print self.batchgrid, input1[:,:,:,4:6] x = torch.mul(self.batchgrid, input1[:,:,:,0:3]) y = torch.mul(self.batchgrid, input1[:,:,:,3:6]) output = torch.cat([torch.sum(x,3),torch.sum(y,3)], 3) return output class DenseAffine3DGridGen(Module): def __init__(self, height, width, lr = 1, aux_loss = False): super(DenseAffine3DGridGen, self).__init__() self.height, self.width = height, width self.aux_loss = aux_loss self.lr = lr self.grid = np.zeros( [self.height, self.width, 3], dtype=np.float32) self.grid[:,:,0] = np.expand_dims(np.repeat(np.expand_dims(np.arange(-1, 1, 2.0/self.height), 0), repeats = self.width, axis = 0).T, 0) self.grid[:,:,1] = np.expand_dims(np.repeat(np.expand_dims(np.arange(-1, 1, 2.0/self.width), 0), repeats = self.height, axis = 0), 0) self.grid[:,:,2] = np.ones([self.height, width]) self.grid = torch.from_numpy(self.grid.astype(np.float32)) self.theta = self.grid[:,:,0] * np.pi/2 + np.pi/2 self.phi = self.grid[:,:,1] * np.pi self.x = torch.sin(self.theta) * torch.cos(self.phi) self.y = torch.sin(self.theta) * torch.sin(self.phi) self.z = torch.cos(self.theta) self.grid3d = torch.from_numpy(np.zeros( [self.height, self.width, 4], dtype=np.float32)) self.grid3d[:,:,0] = self.x self.grid3d[:,:,1] = self.y self.grid3d[:,:,2] = self.z self.grid3d[:,:,3] = self.grid[:,:,2] def forward(self, input1): self.batchgrid3d = torch.zeros(torch.Size([input1.size(0)]) + self.grid3d.size()) for i in range(input1.size(0)): self.batchgrid3d[i] = self.grid3d self.batchgrid3d = Variable(self.batchgrid3d) #print(self.batchgrid3d) x = torch.sum(torch.mul(self.batchgrid3d, input1[:,:,:,0:4]), 3) y = torch.sum(torch.mul(self.batchgrid3d, input1[:,:,:,4:8]), 3) z = torch.sum(torch.mul(self.batchgrid3d, input1[:,:,:,8:]), 3) #print(x) r = torch.sqrt(x**2 + y**2 + z**2) + 1e-5 #print(r) theta = torch.acos(z/r)/(np.pi/2) - 1 #phi = torch.atan(y/x) phi = torch.atan(y/(x + 1e-5)) + np.pi * x.lt(0).type(torch.FloatTensor) * (y.ge(0).type(torch.FloatTensor) - y.lt(0).type(torch.FloatTensor)) phi = phi/np.pi output = torch.cat([theta,phi], 3) return output class DenseAffine3DGridGen_rotate(Module): def __init__(self, height, width, lr = 1, aux_loss = False): super(DenseAffine3DGridGen_rotate, self).__init__() self.height, self.width = height, width self.aux_loss = aux_loss self.lr = lr self.grid = np.zeros( [self.height, self.width, 3], dtype=np.float32) self.grid[:,:,0] = np.expand_dims(np.repeat(np.expand_dims(np.arange(-1, 1, 2.0/self.height), 0), repeats = self.width, axis = 0).T, 0) self.grid[:,:,1] = np.expand_dims(np.repeat(np.expand_dims(np.arange(-1, 1, 2.0/self.width), 0), repeats = self.height, axis = 0), 0) self.grid[:,:,2] = np.ones([self.height, width]) self.grid = torch.from_numpy(self.grid.astype(np.float32)) self.theta = self.grid[:,:,0] * np.pi/2 + np.pi/2 self.phi = self.grid[:,:,1] * np.pi self.x = torch.sin(self.theta) * torch.cos(self.phi) self.y = torch.sin(self.theta) * torch.sin(self.phi) self.z = torch.cos(self.theta) self.grid3d = torch.from_numpy(np.zeros( [self.height, self.width, 4], dtype=np.float32)) self.grid3d[:,:,0] = self.x self.grid3d[:,:,1] = self.y self.grid3d[:,:,2] = self.z self.grid3d[:,:,3] = self.grid[:,:,2] def forward(self, input1, input2): self.batchgrid3d = torch.zeros(torch.Size([input1.size(0)]) + self.grid3d.size()) for i in range(input1.size(0)): self.batchgrid3d[i] = self.grid3d self.batchgrid3d = Variable(self.batchgrid3d) self.batchgrid = torch.zeros(torch.Size([input1.size(0)]) + self.grid.size()) for i in range(input1.size(0)): self.batchgrid[i] = self.grid self.batchgrid = Variable(self.batchgrid) #print(self.batchgrid3d) x = torch.sum(torch.mul(self.batchgrid3d, input1[:,:,:,0:4]), 3) y = torch.sum(torch.mul(self.batchgrid3d, input1[:,:,:,4:8]), 3) z = torch.sum(torch.mul(self.batchgrid3d, input1[:,:,:,8:]), 3) #print(x) r = torch.sqrt(x**2 + y**2 + z**2) + 1e-5 #print(r) theta = torch.acos(z/r)/(np.pi/2) - 1 #phi = torch.atan(y/x) phi = torch.atan(y/(x + 1e-5)) + np.pi * x.lt(0).type(torch.FloatTensor) * (y.ge(0).type(torch.FloatTensor) - y.lt(0).type(torch.FloatTensor)) phi = phi/np.pi input_u = input2.view(-1,1,1,1).repeat(1,self.height, self.width,1) output = torch.cat([theta,phi], 3) output1 = torch.atan(torch.tan(np.pi/2.0*(output[:,:,:,1:2] + self.batchgrid[:,:,:,2:] * input_u[:,:,:,:]))) /(np.pi/2) output2 = torch.cat([output[:,:,:,0:1], output1], 3) return output2 class Depth3DGridGen(Module): def __init__(self, height, width, lr = 1, aux_loss = False): super(Depth3DGridGen, self).__init__() self.height, self.width = height, width self.aux_loss = aux_loss self.lr = lr self.grid = np.zeros( [self.height, self.width, 3], dtype=np.float32) self.grid[:,:,0] = np.expand_dims(np.repeat(np.expand_dims(np.arange(-1, 1, 2.0/self.height), 0), repeats = self.width, axis = 0).T, 0) self.grid[:,:,1] = np.expand_dims(np.repeat(np.expand_dims(np.arange(-1, 1, 2.0/self.width), 0), repeats = self.height, axis = 0), 0) self.grid[:,:,2] = np.ones([self.height, width]) self.grid = torch.from_numpy(self.grid.astype(np.float32)) self.theta = self.grid[:,:,0] * np.pi/2 + np.pi/2 self.phi = self.grid[:,:,1] * np.pi self.x = torch.sin(self.theta) * torch.cos(self.phi) self.y = torch.sin(self.theta) * torch.sin(self.phi) self.z = torch.cos(self.theta) self.grid3d = torch.from_numpy(np.zeros( [self.height, self.width, 4], dtype=np.float32)) self.grid3d[:,:,0] = self.x self.grid3d[:,:,1] = self.y self.grid3d[:,:,2] = self.z self.grid3d[:,:,3] = self.grid[:,:,2] def forward(self, depth, trans0, trans1, rotate): self.batchgrid3d = torch.zeros(torch.Size([depth.size(0)]) + self.grid3d.size()) for i in range(depth.size(0)): self.batchgrid3d[i] = self.grid3d self.batchgrid3d = Variable(self.batchgrid3d) self.batchgrid = torch.zeros(torch.Size([depth.size(0)]) + self.grid.size()) for i in range(depth.size(0)): self.batchgrid[i] = self.grid self.batchgrid = Variable(self.batchgrid) x = self.batchgrid3d[:,:,:,0:1] * depth + trans0.view(-1,1,1,1).repeat(1, self.height, self.width, 1) y = self.batchgrid3d[:,:,:,1:2] * depth + trans1.view(-1,1,1,1).repeat(1, self.height, self.width, 1) z = self.batchgrid3d[:,:,:,2:3] * depth #print(x.size(), y.size(), z.size()) r = torch.sqrt(x**2 + y**2 + z**2) + 1e-5 #print(r) theta = torch.acos(z/r)/(np.pi/2) - 1 #phi = torch.atan(y/x) phi = torch.atan(y/(x + 1e-5)) + np.pi * x.lt(0).type(torch.FloatTensor) * (y.ge(0).type(torch.FloatTensor) - y.lt(0).type(torch.FloatTensor)) phi = phi/np.pi #print(theta.size(), phi.size()) input_u = rotate.view(-1,1,1,1).repeat(1,self.height, self.width,1) output = torch.cat([theta,phi], 3) #print(output.size()) output1 = torch.atan(torch.tan(np.pi/2.0*(output[:,:,:,1:2] + self.batchgrid[:,:,:,2:] * input_u[:,:,:,:]))) /(np.pi/2) output2 = torch.cat([output[:,:,:,0:1], output1], 3) return output2 class Depth3DGridGen_with_mask(Module): def __init__(self, height, width, lr = 1, aux_loss = False, ray_tracing = False): super(Depth3DGridGen_with_mask, self).__init__() self.height, self.width = height, width self.aux_loss = aux_loss self.lr = lr self.ray_tracing = ray_tracing self.grid = np.zeros( [self.height, self.width, 3], dtype=np.float32) self.grid[:,:,0] = np.expand_dims(np.repeat(np.expand_dims(np.arange(-1, 1, 2.0/self.height), 0), repeats = self.width, axis = 0).T, 0) self.grid[:,:,1] = np.expand_dims(np.repeat(np.expand_dims(np.arange(-1, 1, 2.0/self.width), 0), repeats = self.height, axis = 0), 0) self.grid[:,:,2] = np.ones([self.height, width]) self.grid = torch.from_numpy(self.grid.astype(np.float32)) self.theta = self.grid[:,:,0] * np.pi/2 + np.pi/2 self.phi = self.grid[:,:,1] * np.pi self.x = torch.sin(self.theta) * torch.cos(self.phi) self.y = torch.sin(self.theta) * torch.sin(self.phi) self.z = torch.cos(self.theta) self.grid3d = torch.from_numpy(np.zeros( [self.height, self.width, 4], dtype=np.float32)) self.grid3d[:,:,0] = self.x self.grid3d[:,:,1] = self.y self.grid3d[:,:,2] = self.z self.grid3d[:,:,3] = self.grid[:,:,2] def forward(self, depth, trans0, trans1, rotate): self.batchgrid3d = torch.zeros(torch.Size([depth.size(0)]) + self.grid3d.size()) for i in range(depth.size(0)): self.batchgrid3d[i] = self.grid3d self.batchgrid3d = Variable(self.batchgrid3d) self.batchgrid = torch.zeros(torch.Size([depth.size(0)]) + self.grid.size()) for i in range(depth.size(0)): self.batchgrid[i] = self.grid self.batchgrid = Variable(self.batchgrid) if depth.is_cuda: self.batchgrid = self.batchgrid.cuda() self.batchgrid3d = self.batchgrid3d.cuda() x_ = self.batchgrid3d[:,:,:,0:1] * depth + trans0.view(-1,1,1,1).repeat(1, self.height, self.width, 1) y_ = self.batchgrid3d[:,:,:,1:2] * depth + trans1.view(-1,1,1,1).repeat(1, self.height, self.width, 1) z = self.batchgrid3d[:,:,:,2:3] * depth #print(x.size(), y.size(), z.size()) rotate_z = rotate.view(-1,1,1,1).repeat(1,self.height, self.width,1) * np.pi x = x_ * torch.cos(rotate_z) - y_ * torch.sin(rotate_z) y = x_ * torch.sin(rotate_z) + y_ * torch.cos(rotate_z) r = torch.sqrt(x**2 + y**2 + z**2) + 1e-5 #print(r) theta = torch.acos(z/r)/(np.pi/2) - 1 #phi = torch.atan(y/x) if depth.is_cuda: phi = torch.atan(y/(x + 1e-5)) + np.pi * x.lt(0).type(torch.cuda.FloatTensor) * (y.ge(0).type(torch.cuda.FloatTensor) - y.lt(0).type(torch.cuda.FloatTensor)) else: phi = torch.atan(y/(x + 1e-5)) + np.pi * x.lt(0).type(torch.FloatTensor) * (y.ge(0).type(torch.FloatTensor) - y.lt(0).type(torch.FloatTensor)) phi = phi/np.pi output = torch.cat([theta,phi], 3) return output ================================================ FILE: lib/model/roi_crop/modules/roi_crop.py ================================================ from torch.nn.modules.module import Module from ..functions.roi_crop import RoICropFunction class _RoICrop(Module): def __init__(self, layout = 'BHWD'): super(_RoICrop, self).__init__() def forward(self, input1, input2): return RoICropFunction()(input1, input2) ================================================ FILE: lib/model/roi_crop/src/roi_crop.c ================================================ #include #include #include #define real float int BilinearSamplerBHWD_updateOutput(THFloatTensor *inputImages, THFloatTensor *grids, THFloatTensor *output) { int batchsize = THFloatTensor_size(inputImages, 0); int inputImages_height = THFloatTensor_size(inputImages, 1); int inputImages_width = THFloatTensor_size(inputImages, 2); int output_height = THFloatTensor_size(output, 1); int output_width = THFloatTensor_size(output, 2); int inputImages_channels = THFloatTensor_size(inputImages, 3); int output_strideBatch = THFloatTensor_stride(output, 0); int output_strideHeight = THFloatTensor_stride(output, 1); int output_strideWidth = THFloatTensor_stride(output, 2); int inputImages_strideBatch = THFloatTensor_stride(inputImages, 0); int inputImages_strideHeight = THFloatTensor_stride(inputImages, 1); int inputImages_strideWidth = THFloatTensor_stride(inputImages, 2); int grids_strideBatch = THFloatTensor_stride(grids, 0); int grids_strideHeight = THFloatTensor_stride(grids, 1); int grids_strideWidth = THFloatTensor_stride(grids, 2); real *inputImages_data, *output_data, *grids_data; inputImages_data = THFloatTensor_data(inputImages); output_data = THFloatTensor_data(output); grids_data = THFloatTensor_data(grids); int b, yOut, xOut; for(b=0; b < batchsize; b++) { for(yOut=0; yOut < output_height; yOut++) { for(xOut=0; xOut < output_width; xOut++) { //read the grid real yf = grids_data[b*grids_strideBatch + yOut*grids_strideHeight + xOut*grids_strideWidth]; real xf = grids_data[b*grids_strideBatch + yOut*grids_strideHeight + xOut*grids_strideWidth + 1]; // get the weights for interpolation int yInTopLeft, xInTopLeft; real yWeightTopLeft, xWeightTopLeft; real xcoord = (xf + 1) * (inputImages_width - 1) / 2; xInTopLeft = floor(xcoord); xWeightTopLeft = 1 - (xcoord - xInTopLeft); real ycoord = (yf + 1) * (inputImages_height - 1) / 2; yInTopLeft = floor(ycoord); yWeightTopLeft = 1 - (ycoord - yInTopLeft); const int outAddress = output_strideBatch * b + output_strideHeight * yOut + output_strideWidth * xOut; const int inTopLeftAddress = inputImages_strideBatch * b + inputImages_strideHeight * yInTopLeft + inputImages_strideWidth * xInTopLeft; const int inTopRightAddress = inTopLeftAddress + inputImages_strideWidth; const int inBottomLeftAddress = inTopLeftAddress + inputImages_strideHeight; const int inBottomRightAddress = inBottomLeftAddress + inputImages_strideWidth; real v=0; real inTopLeft=0; real inTopRight=0; real inBottomLeft=0; real inBottomRight=0; // we are careful with the boundaries bool topLeftIsIn = xInTopLeft >= 0 && xInTopLeft <= inputImages_width-1 && yInTopLeft >= 0 && yInTopLeft <= inputImages_height-1; bool topRightIsIn = xInTopLeft+1 >= 0 && xInTopLeft+1 <= inputImages_width-1 && yInTopLeft >= 0 && yInTopLeft <= inputImages_height-1; bool bottomLeftIsIn = xInTopLeft >= 0 && xInTopLeft <= inputImages_width-1 && yInTopLeft+1 >= 0 && yInTopLeft+1 <= inputImages_height-1; bool bottomRightIsIn = xInTopLeft+1 >= 0 && xInTopLeft+1 <= inputImages_width-1 && yInTopLeft+1 >= 0 && yInTopLeft+1 <= inputImages_height-1; int t; // interpolation happens here for(t=0; t= 0 && xInTopLeft <= inputImages_width-1 && yInTopLeft >= 0 && yInTopLeft <= inputImages_height-1; bool topRightIsIn = xInTopLeft+1 >= 0 && xInTopLeft+1 <= inputImages_width-1 && yInTopLeft >= 0 && yInTopLeft <= inputImages_height-1; bool bottomLeftIsIn = xInTopLeft >= 0 && xInTopLeft <= inputImages_width-1 && yInTopLeft+1 >= 0 && yInTopLeft+1 <= inputImages_height-1; bool bottomRightIsIn = xInTopLeft+1 >= 0 && xInTopLeft+1 <= inputImages_width-1 && yInTopLeft+1 >= 0 && yInTopLeft+1 <= inputImages_height-1; int t; for(t=0; t= 0 && xInTopLeft <= inputImages_width-1 && yInTopLeft >= 0 && yInTopLeft <= inputImages_height-1; bool topRightIsIn = xInTopLeft+1 >= 0 && xInTopLeft+1 <= inputImages_width-1 && yInTopLeft >= 0 && yInTopLeft <= inputImages_height-1; bool bottomLeftIsIn = xInTopLeft >= 0 && xInTopLeft <= inputImages_width-1 && yInTopLeft+1 >= 0 && yInTopLeft+1 <= inputImages_height-1; bool bottomRightIsIn = xInTopLeft+1 >= 0 && xInTopLeft+1 <= inputImages_width-1 && yInTopLeft+1 >= 0 && yInTopLeft+1 <= inputImages_height-1; int t; // interpolation happens here for(t=0; t= 0 && xInTopLeft <= inputImages_width-1 && yInTopLeft >= 0 && yInTopLeft <= inputImages_height-1; bool topRightIsIn = xInTopLeft+1 >= 0 && xInTopLeft+1 <= inputImages_width-1 && yInTopLeft >= 0 && yInTopLeft <= inputImages_height-1; bool bottomLeftIsIn = xInTopLeft >= 0 && xInTopLeft <= inputImages_width-1 && yInTopLeft+1 >= 0 && yInTopLeft+1 <= inputImages_height-1; bool bottomRightIsIn = xInTopLeft+1 >= 0 && xInTopLeft+1 <= inputImages_width-1 && yInTopLeft+1 >= 0 && yInTopLeft+1 <= inputImages_height-1; int t; for(t=0; t #include #include #include "roi_crop_cuda_kernel.h" #define real float // this symbol will be resolved automatically from PyTorch libs extern THCState *state; // Bilinear sampling is done in BHWD (coalescing is not obvious in BDHW) // we assume BHWD format in inputImages // we assume BHW(YX) format on grids int BilinearSamplerBHWD_updateOutput_cuda(THCudaTensor *inputImages, THCudaTensor *grids, THCudaTensor *output){ // THCState *state = getCutorchState(L); // THCudaTensor *inputImages = (THCudaTensor *)luaT_checkudata(L, 2, "torch.CudaTensor"); // THCudaTensor *grids = (THCudaTensor *)luaT_checkudata(L, 3, "torch.CudaTensor"); // THCudaTensor *output = (THCudaTensor *)luaT_checkudata(L, 4, "torch.CudaTensor"); int success = 0; success = BilinearSamplerBHWD_updateOutput_cuda_kernel(THCudaTensor_size(state, output, 1), THCudaTensor_size(state, output, 3), THCudaTensor_size(state, output, 2), THCudaTensor_size(state, output, 0), THCudaTensor_size(state, inputImages, 1), THCudaTensor_size(state, inputImages, 2), THCudaTensor_size(state, inputImages, 3), THCudaTensor_size(state, inputImages, 0), THCudaTensor_data(state, inputImages), THCudaTensor_stride(state, inputImages, 0), THCudaTensor_stride(state, inputImages, 1), THCudaTensor_stride(state, inputImages, 2), THCudaTensor_stride(state, inputImages, 3), THCudaTensor_data(state, grids), THCudaTensor_stride(state, grids, 0), THCudaTensor_stride(state, grids, 3), THCudaTensor_stride(state, grids, 1), THCudaTensor_stride(state, grids, 2), THCudaTensor_data(state, output), THCudaTensor_stride(state, output, 0), THCudaTensor_stride(state, output, 1), THCudaTensor_stride(state, output, 2), THCudaTensor_stride(state, output, 3), THCState_getCurrentStream(state)); //check for errors if (!success) { THError("aborting"); } return 1; } int BilinearSamplerBHWD_updateGradInput_cuda(THCudaTensor *inputImages, THCudaTensor *grids, THCudaTensor *gradInputImages, THCudaTensor *gradGrids, THCudaTensor *gradOutput) { // THCState *state = getCutorchState(L); // THCudaTensor *inputImages = (THCudaTensor *)luaT_checkudata(L, 2, "torch.CudaTensor"); // THCudaTensor *grids = (THCudaTensor *)luaT_checkudata(L, 3, "torch.CudaTensor"); // THCudaTensor *gradInputImages = (THCudaTensor *)luaT_checkudata(L, 4, "torch.CudaTensor"); // THCudaTensor *gradGrids = (THCudaTensor *)luaT_checkudata(L, 5, "torch.CudaTensor"); // THCudaTensor *gradOutput = (THCudaTensor *)luaT_checkudata(L, 6, "torch.CudaTensor"); int success = 0; success = BilinearSamplerBHWD_updateGradInput_cuda_kernel(THCudaTensor_size(state, gradOutput, 1), THCudaTensor_size(state, gradOutput, 3), THCudaTensor_size(state, gradOutput, 2), THCudaTensor_size(state, gradOutput, 0), THCudaTensor_size(state, inputImages, 1), THCudaTensor_size(state, inputImages, 2), THCudaTensor_size(state, inputImages, 3), THCudaTensor_size(state, inputImages, 0), THCudaTensor_data(state, inputImages), THCudaTensor_stride(state, inputImages, 0), THCudaTensor_stride(state, inputImages, 1), THCudaTensor_stride(state, inputImages, 2), THCudaTensor_stride(state, inputImages, 3), THCudaTensor_data(state, grids), THCudaTensor_stride(state, grids, 0), THCudaTensor_stride(state, grids, 3), THCudaTensor_stride(state, grids, 1), THCudaTensor_stride(state, grids, 2), THCudaTensor_data(state, gradInputImages), THCudaTensor_stride(state, gradInputImages, 0), THCudaTensor_stride(state, gradInputImages, 1), THCudaTensor_stride(state, gradInputImages, 2), THCudaTensor_stride(state, gradInputImages, 3), THCudaTensor_data(state, gradGrids), THCudaTensor_stride(state, gradGrids, 0), THCudaTensor_stride(state, gradGrids, 3), THCudaTensor_stride(state, gradGrids, 1), THCudaTensor_stride(state, gradGrids, 2), THCudaTensor_data(state, gradOutput), THCudaTensor_stride(state, gradOutput, 0), THCudaTensor_stride(state, gradOutput, 1), THCudaTensor_stride(state, gradOutput, 2), THCudaTensor_stride(state, gradOutput, 3), THCState_getCurrentStream(state)); //check for errors if (!success) { THError("aborting"); } return 1; } ================================================ FILE: lib/model/roi_crop/src/roi_crop_cuda.h ================================================ // Bilinear sampling is done in BHWD (coalescing is not obvious in BDHW) // we assume BHWD format in inputImages // we assume BHW(YX) format on grids int BilinearSamplerBHWD_updateOutput_cuda(THCudaTensor *inputImages, THCudaTensor *grids, THCudaTensor *output); int BilinearSamplerBHWD_updateGradInput_cuda(THCudaTensor *inputImages, THCudaTensor *grids, THCudaTensor *gradInputImages, THCudaTensor *gradGrids, THCudaTensor *gradOutput); ================================================ FILE: lib/model/roi_crop/src/roi_crop_cuda_kernel.cu ================================================ #include #include #include "roi_crop_cuda_kernel.h" #define real float // Bilinear sampling is done in BHWD (coalescing is not obvious in BDHW) // we assume BHWD format in inputImages // we assume BHW(YX) format on grids __device__ void getTopLeft(float x, int width, int& point, float& weight) { /* for interpolation : stores in point and weight : - the x-coordinate of the pixel on the left (or y-coordinate of the upper pixel) - the weight for interpolating */ float xcoord = (x + 1) * (width - 1) / 2; point = floor(xcoord); weight = 1 - (xcoord - point); } __device__ bool between(int value, int lowerBound, int upperBound) { return (value >= lowerBound && value <= upperBound); } __device__ void sumReduceShMem(volatile float s[]) { /* obviously only works for 32 elements */ /* sums up a shared memory array of 32 elements, stores it in s[0] */ /* whole warp can then read first element (broadcasting) */ if(threadIdx.x<16) { s[threadIdx.x] = s[threadIdx.x] + s[threadIdx.x+16]; } if(threadIdx.x<8) { s[threadIdx.x] = s[threadIdx.x] + s[threadIdx.x+8]; } if(threadIdx.x<4) { s[threadIdx.x] = s[threadIdx.x] + s[threadIdx.x+4]; } if(threadIdx.x<2) { s[threadIdx.x] = s[threadIdx.x] + s[threadIdx.x+2]; } if(threadIdx.x<1) { s[threadIdx.x] = s[threadIdx.x] + s[threadIdx.x+1]; } } // CUDA: grid stride looping #define CUDA_KERNEL_LOOP(i, n) \ for (int i = blockIdx.x * blockDim.x + threadIdx.x; \ i < (n); \ i += blockDim.x * gridDim.x) __global__ void bilinearSamplingFromGrid(const int nthreads, float* inputImages_data, int inputImages_strideBatch, int inputImages_strideChannels, int inputImages_strideHeight, int inputImages_strideWidth, float* grids_data, int grids_strideBatch, int grids_strideYX, int grids_strideHeight, int grids_strideWidth, float* output_data, int output_strideBatch, int output_strideChannels, int output_strideHeight, int output_strideWidth, int inputImages_channels, int inputImages_height, int inputImages_width, int output_channels, int output_height, int output_width, int output_batchsize, int roiPerImage) { CUDA_KERNEL_LOOP(index, nthreads) { const int xOut = index % output_width; const int yOut = (index / output_width) % output_height; const int cOut = (index / output_width / output_height) % output_channels; const int b = index / output_width / output_height / output_channels; const int width = inputImages_width; const int height = inputImages_height; const int b_input = b / roiPerImage; float yf = grids_data[b*grids_strideBatch + yOut*grids_strideHeight + xOut*grids_strideWidth]; float xf = grids_data[b*grids_strideBatch + yOut*grids_strideHeight + xOut*grids_strideWidth + 1]; int yInTopLeft, xInTopLeft; float yWeightTopLeft, xWeightTopLeft; getTopLeft(xf, inputImages_width, xInTopLeft, xWeightTopLeft); getTopLeft(yf, inputImages_height, yInTopLeft, yWeightTopLeft); // const int outAddress = output_strideBatch * b + output_strideHeight * yOut + output_strideWidth * xOut; const int outAddress = output_strideBatch * b + output_strideChannels * cOut + output_strideHeight * yOut + xOut; const int inTopLeftAddress = inputImages_strideBatch * b_input + inputImages_strideChannels * cOut + inputImages_strideHeight * yInTopLeft + xInTopLeft; const int inTopRightAddress = inTopLeftAddress + inputImages_strideWidth; const int inBottomLeftAddress = inTopLeftAddress + inputImages_strideHeight; const int inBottomRightAddress = inBottomLeftAddress + inputImages_strideWidth; float v=0; float inTopLeft=0; float inTopRight=0; float inBottomLeft=0; float inBottomRight=0; bool topLeftIsIn = between(xInTopLeft, 0, width-1) && between(yInTopLeft, 0, height-1); bool topRightIsIn = between(xInTopLeft+1, 0, width-1) && between(yInTopLeft, 0, height-1); bool bottomLeftIsIn = between(xInTopLeft, 0, width-1) && between(yInTopLeft+1, 0, height-1); bool bottomRightIsIn = between(xInTopLeft+1, 0, width-1) && between(yInTopLeft+1, 0, height-1); if (!topLeftIsIn && !topRightIsIn && !bottomLeftIsIn && !bottomRightIsIn) continue; if(topLeftIsIn) inTopLeft = inputImages_data[inTopLeftAddress]; if(topRightIsIn) inTopRight = inputImages_data[inTopRightAddress]; if(bottomLeftIsIn) inBottomLeft = inputImages_data[inBottomLeftAddress]; if(bottomRightIsIn) inBottomRight = inputImages_data[inBottomRightAddress]; v = xWeightTopLeft * yWeightTopLeft * inTopLeft + (1 - xWeightTopLeft) * yWeightTopLeft * inTopRight + xWeightTopLeft * (1 - yWeightTopLeft) * inBottomLeft + (1 - xWeightTopLeft) * (1 - yWeightTopLeft) * inBottomRight; output_data[outAddress] = v; } } __global__ void backwardBilinearSampling(const int nthreads, float* inputImages_data, int inputImages_strideBatch, int inputImages_strideChannels, int inputImages_strideHeight, int inputImages_strideWidth, float* gradInputImages_data, int gradInputImages_strideBatch, int gradInputImages_strideChannels, int gradInputImages_strideHeight, int gradInputImages_strideWidth, float* grids_data, int grids_strideBatch, int grids_strideYX, int grids_strideHeight, int grids_strideWidth, float* gradGrids_data, int gradGrids_strideBatch, int gradGrids_strideYX, int gradGrids_strideHeight, int gradGrids_strideWidth, float* gradOutput_data, int gradOutput_strideBatch, int gradOutput_strideChannels, int gradOutput_strideHeight, int gradOutput_strideWidth, int inputImages_channels, int inputImages_height, int inputImages_width, int gradOutput_channels, int gradOutput_height, int gradOutput_width, int gradOutput_batchsize, int roiPerImage) { CUDA_KERNEL_LOOP(index, nthreads) { const int xOut = index % gradOutput_width; const int yOut = (index / gradOutput_width) % gradOutput_height; const int cOut = (index / gradOutput_width / gradOutput_height) % gradOutput_channels; const int b = index / gradOutput_width / gradOutput_height / gradOutput_channels; const int b_input = b / roiPerImage; const int width = inputImages_width; const int height = inputImages_height; float yf = grids_data[b*grids_strideBatch + yOut*grids_strideHeight + xOut*grids_strideWidth]; float xf = grids_data[b*grids_strideBatch + yOut*grids_strideHeight + xOut*grids_strideWidth + 1]; int yInTopLeft, xInTopLeft; float yWeightTopLeft, xWeightTopLeft; getTopLeft(xf, inputImages_width, xInTopLeft, xWeightTopLeft); getTopLeft(yf, inputImages_height, yInTopLeft, yWeightTopLeft); const int inTopLeftAddress = inputImages_strideBatch * b_input + inputImages_strideChannels * cOut + inputImages_strideHeight * yInTopLeft + xInTopLeft; const int inTopRightAddress = inTopLeftAddress + inputImages_strideWidth; const int inBottomLeftAddress = inTopLeftAddress + inputImages_strideHeight; const int inBottomRightAddress = inBottomLeftAddress + inputImages_strideWidth; const int gradInputImagesTopLeftAddress = gradInputImages_strideBatch * b_input + gradInputImages_strideChannels * cOut + gradInputImages_strideHeight * yInTopLeft + xInTopLeft; const int gradInputImagesTopRightAddress = gradInputImagesTopLeftAddress + gradInputImages_strideWidth; const int gradInputImagesBottomLeftAddress = gradInputImagesTopLeftAddress + gradInputImages_strideHeight; const int gradInputImagesBottomRightAddress = gradInputImagesBottomLeftAddress + gradInputImages_strideWidth; const int gradOutputAddress = gradOutput_strideBatch * b + gradOutput_strideChannels * cOut + gradOutput_strideHeight * yOut + xOut; float topLeftDotProduct = 0; float topRightDotProduct = 0; float bottomLeftDotProduct = 0; float bottomRightDotProduct = 0; bool topLeftIsIn = between(xInTopLeft, 0, width-1) && between(yInTopLeft, 0, height-1); bool topRightIsIn = between(xInTopLeft+1, 0, width-1) && between(yInTopLeft, 0, height-1); bool bottomLeftIsIn = between(xInTopLeft, 0, width-1) && between(yInTopLeft+1, 0, height-1); bool bottomRightIsIn = between(xInTopLeft+1, 0, width-1) && between(yInTopLeft+1, 0, height-1); float gradOutValue = gradOutput_data[gradOutputAddress]; // bool between(int value, int lowerBound, int upperBound) if(topLeftIsIn) { float inTopLeft = inputImages_data[inTopLeftAddress]; topLeftDotProduct += inTopLeft * gradOutValue; atomicAdd(&gradInputImages_data[gradInputImagesTopLeftAddress], xWeightTopLeft * yWeightTopLeft * gradOutValue); } if(topRightIsIn) { float inTopRight = inputImages_data[inTopRightAddress]; topRightDotProduct += inTopRight * gradOutValue; atomicAdd(&gradInputImages_data[gradInputImagesTopRightAddress], (1 - xWeightTopLeft) * yWeightTopLeft * gradOutValue); } if(bottomLeftIsIn) { float inBottomLeft = inputImages_data[inBottomLeftAddress]; bottomLeftDotProduct += inBottomLeft * gradOutValue; atomicAdd(&gradInputImages_data[gradInputImagesBottomLeftAddress], xWeightTopLeft * (1 - yWeightTopLeft) * gradOutValue); } if(bottomRightIsIn) { float inBottomRight = inputImages_data[inBottomRightAddress]; bottomRightDotProduct += inBottomRight * gradOutValue; atomicAdd(&gradInputImages_data[gradInputImagesBottomRightAddress], (1 - xWeightTopLeft) * (1 - yWeightTopLeft) * gradOutValue); } } } #ifdef __cplusplus extern "C" { #endif int BilinearSamplerBHWD_updateOutput_cuda_kernel(/*output->size[1]*/int oc, /*output->size[3]*/int ow, /*output->size[2]*/int oh, /*output->size[0]*/int ob, /*THCudaTensor_size(state, inputImages, 1)*/int ic, /*THCudaTensor_size(state, inputImages, 2)*/int ih, /*THCudaTensor_size(state, inputImages, 3)*/int iw, /*THCudaTensor_size(state, inputImages, 0)*/int ib, /*THCudaTensor *inputImages*/float *inputImages, int isb, int isc, int ish, int isw, /*THCudaTensor *grids*/float *grids, int gsb, int gsc, int gsh, int gsw, /*THCudaTensor *output*/float *output, int osb, int osc, int osh, int osw, /*THCState_getCurrentStream(state)*/cudaStream_t stream) { const int kThreadsPerBlock = 1024; int output_size = ob * oh * ow * oc; cudaError_t err; int roiPerImage = ob / ib; // printf("forward pass\n"); bilinearSamplingFromGrid<<<(output_size + kThreadsPerBlock - 1) / kThreadsPerBlock, kThreadsPerBlock, 0, stream>>>( output_size, /*THCudaTensor_data(state, inputImages)*/inputImages, /*THCudaTensor_stride(state, inputImages, 0)*/isb, /*THCudaTensor_stride(state, inputImages, 3)*/isc, /*THCudaTensor_stride(state, inputImages, 1)*/ish, /*THCudaTensor_stride(state, inputImages, 2)*/isw, /*THCudaTensor_data(state, grids)*/grids, /*THCudaTensor_stride(state, grids, 0)*/gsb, /*THCudaTensor_stride(state, grids, 3)*/gsc, /*THCudaTensor_stride(state, grids, 1)*/gsh, /*THCudaTensor_stride(state, grids, 2)*/gsw, /*THCudaTensor_data(state, output)*/output, /*THCudaTensor_stride(state, output, 0)*/osb, /*THCudaTensor_stride(state, output, 3)*/osc, /*THCudaTensor_stride(state, output, 1)*/osh, /*THCudaTensor_stride(state, output, 2)*/osw, /*THCudaTensor_size(state, inputImages, 3)*/ic, /*THCudaTensor_size(state, inputImages, 1)*/ih, /*THCudaTensor_size(state, inputImages, 2)*/iw, /*THCudaTensor_size(state, output, 3)*/oc, /*THCudaTensor_size(state, output, 1)*/oh, /*THCudaTensor_size(state, output, 2)*/ow, /*THCudaTensor_size(state, output, 0)*/ob, /*Number of rois per image*/roiPerImage); // check for errors err = cudaGetLastError(); if (err != cudaSuccess) { printf("error in BilinearSampler.updateOutput: %s\n", cudaGetErrorString(err)); //THError("aborting"); return 0; } return 1; } int BilinearSamplerBHWD_updateGradInput_cuda_kernel(/*gradOutput->size[1]*/int goc, /*gradOutput->size[3]*/int gow, /*gradOutput->size[2]*/int goh, /*gradOutput->size[0]*/int gob, /*THCudaTensor_size(state, inputImages, 1)*/int ic, /*THCudaTensor_size(state, inputImages, 2)*/int ih, /*THCudaTensor_size(state, inputImages, 3)*/int iw, /*THCudaTensor_size(state, inputImages, 0)*/int ib, /*THCudaTensor *inputImages*/float *inputImages, int isb, int isc, int ish, int isw, /*THCudaTensor *grids*/float *grids, int gsb, int gsc, int gsh, int gsw, /*THCudaTensor *gradInputImages*/float *gradInputImages, int gisb, int gisc, int gish, int gisw, /*THCudaTensor *gradGrids*/float *gradGrids, int ggsb, int ggsc, int ggsh, int ggsw, /*THCudaTensor *gradOutput*/float *gradOutput, int gosb, int gosc, int gosh, int gosw, /*THCState_getCurrentStream(state)*/cudaStream_t stream) { const int kThreadsPerBlock = 1024; int output_size = gob * goh * gow * goc; cudaError_t err; int roiPerImage = gob / ib; // printf("%d %d %d %d\n", gob, goh, gow, goc); // printf("%d %d %d %d\n", ib, ih, iw, ic); // printf("backward pass\n"); backwardBilinearSampling<<<(output_size + kThreadsPerBlock - 1) / kThreadsPerBlock, kThreadsPerBlock, 0, stream>>>( output_size, /*THCudaTensor_data(state, inputImages)*/inputImages, /*THCudaTensor_stride(state, inputImages, 0)*/isb, /*THCudaTensor_stride(state, inputImages, 3)*/isc, /*THCudaTensor_stride(state, inputImages, 1)*/ish, /*THCudaTensor_stride(state, inputImages, 2)*/isw, /*THCudaTensor_data(state, gradInputImages)*/gradInputImages, /*THCudaTensor_stride(state, gradInputImages, 0)*/gisb, /*THCudaTensor_stride(state, gradInputImages, 3)*/gisc, /*THCudaTensor_stride(state, gradInputImages, 1)*/gish, /*THCudaTensor_stride(state, gradInputImages, 2)*/gisw, /*THCudaTensor_data(state, grids)*/grids, /*THCudaTensor_stride(state, grids, 0)*/gsb, /*THCudaTensor_stride(state, grids, 3)*/gsc, /*THCudaTensor_stride(state, grids, 1)*/gsh, /*THCudaTensor_stride(state, grids, 2)*/gsw, /*THCudaTensor_data(state, gradGrids)*/gradGrids, /*THCudaTensor_stride(state, gradGrids, 0)*/ggsb, /*THCudaTensor_stride(state, gradGrids, 3)*/ggsc, /*THCudaTensor_stride(state, gradGrids, 1)*/ggsh, /*THCudaTensor_stride(state, gradGrids, 2)*/ggsw, /*THCudaTensor_data(state, gradOutput)*/gradOutput, /*THCudaTensor_stride(state, gradOutput, 0)*/gosb, /*THCudaTensor_stride(state, gradOutput, 3)*/gosc, /*THCudaTensor_stride(state, gradOutput, 1)*/gosh, /*THCudaTensor_stride(state, gradOutput, 2)*/gosw, /*THCudaTensor_size(state, inputImages, 3)*/ic, /*THCudaTensor_size(state, inputImages, 1)*/ih, /*THCudaTensor_size(state, inputImages, 2)*/iw, /*THCudaTensor_size(state, gradOutput, 3)*/goc, /*THCudaTensor_size(state, gradOutput, 1)*/goh, /*THCudaTensor_size(state, gradOutput, 2)*/gow, /*THCudaTensor_size(state, gradOutput, 0)*/gob, /*Number of rois per image*/roiPerImage); // check for errors err = cudaGetLastError(); if (err != cudaSuccess) { printf("error in BilinearSampler.updateGradInput: %s\n", cudaGetErrorString(err)); //THError("aborting"); return 0; } return 1; } #ifdef __cplusplus } #endif ================================================ FILE: lib/model/roi_crop/src/roi_crop_cuda_kernel.h ================================================ #ifdef __cplusplus extern "C" { #endif int BilinearSamplerBHWD_updateOutput_cuda_kernel(/*output->size[3]*/int oc, /*output->size[2]*/int ow, /*output->size[1]*/int oh, /*output->size[0]*/int ob, /*THCudaTensor_size(state, inputImages, 3)*/int ic, /*THCudaTensor_size(state, inputImages, 1)*/int ih, /*THCudaTensor_size(state, inputImages, 2)*/int iw, /*THCudaTensor_size(state, inputImages, 0)*/int ib, /*THCudaTensor *inputImages*/float *inputImages, int isb, int isc, int ish, int isw, /*THCudaTensor *grids*/float *grids, int gsb, int gsc, int gsh, int gsw, /*THCudaTensor *output*/float *output, int osb, int osc, int osh, int osw, /*THCState_getCurrentStream(state)*/cudaStream_t stream); int BilinearSamplerBHWD_updateGradInput_cuda_kernel(/*gradOutput->size[3]*/int goc, /*gradOutput->size[2]*/int gow, /*gradOutput->size[1]*/int goh, /*gradOutput->size[0]*/int gob, /*THCudaTensor_size(state, inputImages, 3)*/int ic, /*THCudaTensor_size(state, inputImages, 1)*/int ih, /*THCudaTensor_size(state, inputImages, 2)*/int iw, /*THCudaTensor_size(state, inputImages, 0)*/int ib, /*THCudaTensor *inputImages*/float *inputImages, int isb, int isc, int ish, int isw, /*THCudaTensor *grids*/float *grids, int gsb, int gsc, int gsh, int gsw, /*THCudaTensor *gradInputImages*/float *gradInputImages, int gisb, int gisc, int gish, int gisw, /*THCudaTensor *gradGrids*/float *gradGrids, int ggsb, int ggsc, int ggsh, int ggsw, /*THCudaTensor *gradOutput*/float *gradOutput, int gosb, int gosc, int gosh, int gosw, /*THCState_getCurrentStream(state)*/cudaStream_t stream); #ifdef __cplusplus } #endif ================================================ FILE: lib/model/roi_pooling/__init__.py ================================================ ================================================ FILE: lib/model/roi_pooling/_ext/__init__.py ================================================ ================================================ FILE: lib/model/roi_pooling/_ext/roi_pooling/__init__.py ================================================ from torch.utils.ffi import _wrap_function from ._roi_pooling import lib as _lib, ffi as _ffi __all__ = [] def _import_symbols(locals): for symbol in dir(_lib): fn = getattr(_lib, symbol) if callable(fn): locals[symbol] = _wrap_function(fn, _ffi) else: locals[symbol] = fn __all__.append(symbol) _import_symbols(locals()) ================================================ FILE: lib/model/roi_pooling/build.py ================================================ from __future__ import print_function import os import torch from torch.utils.ffi import create_extension sources = ['src/roi_pooling.c'] headers = ['src/roi_pooling.h'] extra_objects = [] defines = [] with_cuda = False this_file = os.path.dirname(os.path.realpath(__file__)) print(this_file) if torch.cuda.is_available(): print('Including CUDA code.') sources += ['src/roi_pooling_cuda.c'] headers += ['src/roi_pooling_cuda.h'] defines += [('WITH_CUDA', None)] with_cuda = True extra_objects = ['src/roi_pooling.cu.o'] extra_objects = [os.path.join(this_file, fname) for fname in extra_objects] ffi = create_extension( '_ext.roi_pooling', headers=headers, sources=sources, define_macros=defines, relative_to=__file__, with_cuda=with_cuda, extra_objects=extra_objects ) if __name__ == '__main__': ffi.build() ================================================ FILE: lib/model/roi_pooling/functions/__init__.py ================================================ ================================================ FILE: lib/model/roi_pooling/functions/roi_pool.py ================================================ import torch from torch.autograd import Function from .._ext import roi_pooling import pdb class RoIPoolFunction(Function): def __init__(ctx, pooled_height, pooled_width, spatial_scale): ctx.pooled_width = pooled_width ctx.pooled_height = pooled_height ctx.spatial_scale = spatial_scale ctx.feature_size = None def forward(ctx, features, rois): ctx.feature_size = features.size() batch_size, num_channels, data_height, data_width = ctx.feature_size num_rois = rois.size(0) output = features.new(num_rois, num_channels, ctx.pooled_height, ctx.pooled_width).zero_() ctx.argmax = features.new(num_rois, num_channels, ctx.pooled_height, ctx.pooled_width).zero_().int() ctx.rois = rois if not features.is_cuda: _features = features.permute(0, 2, 3, 1) roi_pooling.roi_pooling_forward(ctx.pooled_height, ctx.pooled_width, ctx.spatial_scale, _features, rois, output) else: roi_pooling.roi_pooling_forward_cuda(ctx.pooled_height, ctx.pooled_width, ctx.spatial_scale, features, rois, output, ctx.argmax) return output def backward(ctx, grad_output): assert(ctx.feature_size is not None and grad_output.is_cuda) batch_size, num_channels, data_height, data_width = ctx.feature_size grad_input = grad_output.new(batch_size, num_channels, data_height, data_width).zero_() roi_pooling.roi_pooling_backward_cuda(ctx.pooled_height, ctx.pooled_width, ctx.spatial_scale, grad_output, ctx.rois, grad_input, ctx.argmax) return grad_input, None ================================================ FILE: lib/model/roi_pooling/modules/__init__.py ================================================ ================================================ FILE: lib/model/roi_pooling/modules/roi_pool.py ================================================ from torch.nn.modules.module import Module from ..functions.roi_pool import RoIPoolFunction class _RoIPooling(Module): def __init__(self, pooled_height, pooled_width, spatial_scale): super(_RoIPooling, self).__init__() self.pooled_width = int(pooled_width) self.pooled_height = int(pooled_height) self.spatial_scale = float(spatial_scale) def forward(self, features, rois): return RoIPoolFunction(self.pooled_height, self.pooled_width, self.spatial_scale)(features, rois) ================================================ FILE: lib/model/roi_pooling/src/roi_pooling.c ================================================ #include #include int roi_pooling_forward(int pooled_height, int pooled_width, float spatial_scale, THFloatTensor * features, THFloatTensor * rois, THFloatTensor * output) { // Grab the input tensor float * data_flat = THFloatTensor_data(features); float * rois_flat = THFloatTensor_data(rois); float * output_flat = THFloatTensor_data(output); // Number of ROIs int num_rois = THFloatTensor_size(rois, 0); int size_rois = THFloatTensor_size(rois, 1); // batch size int batch_size = THFloatTensor_size(features, 0); if(batch_size != 1) { return 0; } // data height int data_height = THFloatTensor_size(features, 1); // data width int data_width = THFloatTensor_size(features, 2); // Number of channels int num_channels = THFloatTensor_size(features, 3); // Set all element of the output tensor to -inf. THFloatStorage_fill(THFloatTensor_storage(output), -1); // For each ROI R = [batch_index x1 y1 x2 y2]: max pool over R int index_roi = 0; int index_output = 0; int n; for (n = 0; n < num_rois; ++n) { int roi_batch_ind = rois_flat[index_roi + 0]; int roi_start_w = round(rois_flat[index_roi + 1] * spatial_scale); int roi_start_h = round(rois_flat[index_roi + 2] * spatial_scale); int roi_end_w = round(rois_flat[index_roi + 3] * spatial_scale); int roi_end_h = round(rois_flat[index_roi + 4] * spatial_scale); // CHECK_GE(roi_batch_ind, 0); // CHECK_LT(roi_batch_ind, batch_size); int roi_height = fmaxf(roi_end_h - roi_start_h + 1, 1); int roi_width = fmaxf(roi_end_w - roi_start_w + 1, 1); float bin_size_h = (float)(roi_height) / (float)(pooled_height); float bin_size_w = (float)(roi_width) / (float)(pooled_width); int index_data = roi_batch_ind * data_height * data_width * num_channels; const int output_area = pooled_width * pooled_height; int c, ph, pw; for (ph = 0; ph < pooled_height; ++ph) { for (pw = 0; pw < pooled_width; ++pw) { int hstart = (floor((float)(ph) * bin_size_h)); int wstart = (floor((float)(pw) * bin_size_w)); int hend = (ceil((float)(ph + 1) * bin_size_h)); int wend = (ceil((float)(pw + 1) * bin_size_w)); hstart = fminf(fmaxf(hstart + roi_start_h, 0), data_height); hend = fminf(fmaxf(hend + roi_start_h, 0), data_height); wstart = fminf(fmaxf(wstart + roi_start_w, 0), data_width); wend = fminf(fmaxf(wend + roi_start_w, 0), data_width); const int pool_index = index_output + (ph * pooled_width + pw); int is_empty = (hend <= hstart) || (wend <= wstart); if (is_empty) { for (c = 0; c < num_channels * output_area; c += output_area) { output_flat[pool_index + c] = 0; } } else { int h, w, c; for (h = hstart; h < hend; ++h) { for (w = wstart; w < wend; ++w) { for (c = 0; c < num_channels; ++c) { const int index = (h * data_width + w) * num_channels + c; if (data_flat[index_data + index] > output_flat[pool_index + c * output_area]) { output_flat[pool_index + c * output_area] = data_flat[index_data + index]; } } } } } } } // Increment ROI index index_roi += size_rois; index_output += pooled_height * pooled_width * num_channels; } return 1; } ================================================ FILE: lib/model/roi_pooling/src/roi_pooling.h ================================================ int roi_pooling_forward(int pooled_height, int pooled_width, float spatial_scale, THFloatTensor * features, THFloatTensor * rois, THFloatTensor * output); ================================================ FILE: lib/model/roi_pooling/src/roi_pooling_cuda.c ================================================ #include #include #include "roi_pooling_kernel.h" extern THCState *state; int roi_pooling_forward_cuda(int pooled_height, int pooled_width, float spatial_scale, THCudaTensor * features, THCudaTensor * rois, THCudaTensor * output, THCudaIntTensor * argmax) { // Grab the input tensor float * data_flat = THCudaTensor_data(state, features); float * rois_flat = THCudaTensor_data(state, rois); float * output_flat = THCudaTensor_data(state, output); int * argmax_flat = THCudaIntTensor_data(state, argmax); // Number of ROIs int num_rois = THCudaTensor_size(state, rois, 0); int size_rois = THCudaTensor_size(state, rois, 1); if (size_rois != 5) { return 0; } // batch size // int batch_size = THCudaTensor_size(state, features, 0); // if (batch_size != 1) // { // return 0; // } // data height int data_height = THCudaTensor_size(state, features, 2); // data width int data_width = THCudaTensor_size(state, features, 3); // Number of channels int num_channels = THCudaTensor_size(state, features, 1); cudaStream_t stream = THCState_getCurrentStream(state); ROIPoolForwardLaucher( data_flat, spatial_scale, num_rois, data_height, data_width, num_channels, pooled_height, pooled_width, rois_flat, output_flat, argmax_flat, stream); return 1; } int roi_pooling_backward_cuda(int pooled_height, int pooled_width, float spatial_scale, THCudaTensor * top_grad, THCudaTensor * rois, THCudaTensor * bottom_grad, THCudaIntTensor * argmax) { // Grab the input tensor float * top_grad_flat = THCudaTensor_data(state, top_grad); float * rois_flat = THCudaTensor_data(state, rois); float * bottom_grad_flat = THCudaTensor_data(state, bottom_grad); int * argmax_flat = THCudaIntTensor_data(state, argmax); // Number of ROIs int num_rois = THCudaTensor_size(state, rois, 0); int size_rois = THCudaTensor_size(state, rois, 1); if (size_rois != 5) { return 0; } // batch size int batch_size = THCudaTensor_size(state, bottom_grad, 0); // if (batch_size != 1) // { // return 0; // } // data height int data_height = THCudaTensor_size(state, bottom_grad, 2); // data width int data_width = THCudaTensor_size(state, bottom_grad, 3); // Number of channels int num_channels = THCudaTensor_size(state, bottom_grad, 1); cudaStream_t stream = THCState_getCurrentStream(state); ROIPoolBackwardLaucher( top_grad_flat, spatial_scale, batch_size, num_rois, data_height, data_width, num_channels, pooled_height, pooled_width, rois_flat, bottom_grad_flat, argmax_flat, stream); return 1; } ================================================ FILE: lib/model/roi_pooling/src/roi_pooling_cuda.h ================================================ int roi_pooling_forward_cuda(int pooled_height, int pooled_width, float spatial_scale, THCudaTensor * features, THCudaTensor * rois, THCudaTensor * output, THCudaIntTensor * argmax); int roi_pooling_backward_cuda(int pooled_height, int pooled_width, float spatial_scale, THCudaTensor * top_grad, THCudaTensor * rois, THCudaTensor * bottom_grad, THCudaIntTensor * argmax); ================================================ FILE: lib/model/roi_pooling/src/roi_pooling_kernel.cu ================================================ // #ifdef __cplusplus // extern "C" { // #endif #include #include #include #include #include "roi_pooling_kernel.h" #define DIVUP(m, n) ((m) / (m) + ((m) % (n) > 0)) #define CUDA_1D_KERNEL_LOOP(i, n) \ for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; \ i += blockDim.x * gridDim.x) // CUDA: grid stride looping #define CUDA_KERNEL_LOOP(i, n) \ for (int i = blockIdx.x * blockDim.x + threadIdx.x; \ i < (n); \ i += blockDim.x * gridDim.x) __global__ void ROIPoolForward(const int nthreads, const float* bottom_data, const float spatial_scale, const int height, const int width, const int channels, const int pooled_height, const int pooled_width, const float* bottom_rois, float* top_data, int* argmax_data) { CUDA_KERNEL_LOOP(index, nthreads) { // (n, c, ph, pw) is an element in the pooled output // int n = index; // int pw = n % pooled_width; // n /= pooled_width; // int ph = n % pooled_height; // n /= pooled_height; // int c = n % channels; // n /= channels; int pw = index % pooled_width; int ph = (index / pooled_width) % pooled_height; int c = (index / pooled_width / pooled_height) % channels; int n = index / pooled_width / pooled_height / channels; // bottom_rois += n * 5; int roi_batch_ind = bottom_rois[n * 5 + 0]; int roi_start_w = round(bottom_rois[n * 5 + 1] * spatial_scale); int roi_start_h = round(bottom_rois[n * 5 + 2] * spatial_scale); int roi_end_w = round(bottom_rois[n * 5 + 3] * spatial_scale); int roi_end_h = round(bottom_rois[n * 5 + 4] * spatial_scale); // Force malformed ROIs to be 1x1 int roi_width = fmaxf(roi_end_w - roi_start_w + 1, 1); int roi_height = fmaxf(roi_end_h - roi_start_h + 1, 1); float bin_size_h = (float)(roi_height) / (float)(pooled_height); float bin_size_w = (float)(roi_width) / (float)(pooled_width); int hstart = (int)(floor((float)(ph) * bin_size_h)); int wstart = (int)(floor((float)(pw) * bin_size_w)); int hend = (int)(ceil((float)(ph + 1) * bin_size_h)); int wend = (int)(ceil((float)(pw + 1) * bin_size_w)); // Add roi offsets and clip to input boundaries hstart = fminf(fmaxf(hstart + roi_start_h, 0), height); hend = fminf(fmaxf(hend + roi_start_h, 0), height); wstart = fminf(fmaxf(wstart + roi_start_w, 0), width); wend = fminf(fmaxf(wend + roi_start_w, 0), width); bool is_empty = (hend <= hstart) || (wend <= wstart); // Define an empty pooling region to be zero float maxval = is_empty ? 0 : -FLT_MAX; // If nothing is pooled, argmax = -1 causes nothing to be backprop'd int maxidx = -1; // bottom_data += roi_batch_ind * channels * height * width; int bottom_data_batch_offset = roi_batch_ind * channels * height * width; int bottom_data_offset = bottom_data_batch_offset + c * height * width; for (int h = hstart; h < hend; ++h) { for (int w = wstart; w < wend; ++w) { // int bottom_index = (h * width + w) * channels + c; // int bottom_index = (c * height + h) * width + w; int bottom_index = h * width + w; if (bottom_data[bottom_data_offset + bottom_index] > maxval) { maxval = bottom_data[bottom_data_offset + bottom_index]; maxidx = bottom_data_offset + bottom_index; } } } top_data[index] = maxval; if (argmax_data != NULL) argmax_data[index] = maxidx; } } int ROIPoolForwardLaucher( const float* bottom_data, const float spatial_scale, const int num_rois, const int height, const int width, const int channels, const int pooled_height, const int pooled_width, const float* bottom_rois, float* top_data, int* argmax_data, cudaStream_t stream) { const int kThreadsPerBlock = 1024; int output_size = num_rois * pooled_height * pooled_width * channels; cudaError_t err; ROIPoolForward<<<(output_size + kThreadsPerBlock - 1) / kThreadsPerBlock, kThreadsPerBlock, 0, stream>>>( output_size, bottom_data, spatial_scale, height, width, channels, pooled_height, pooled_width, bottom_rois, top_data, argmax_data); // dim3 blocks(DIVUP(output_size, kThreadsPerBlock), // DIVUP(output_size, kThreadsPerBlock)); // dim3 threads(kThreadsPerBlock); // // ROIPoolForward<<>>( // output_size, bottom_data, spatial_scale, height, width, channels, pooled_height, // pooled_width, bottom_rois, top_data, argmax_data); err = cudaGetLastError(); if(cudaSuccess != err) { fprintf( stderr, "cudaCheckError() failed : %s\n", cudaGetErrorString( err ) ); exit( -1 ); } return 1; } __global__ void ROIPoolBackward(const int nthreads, const float* top_diff, const int* argmax_data, const int num_rois, const float spatial_scale, const int height, const int width, const int channels, const int pooled_height, const int pooled_width, float* bottom_diff, const float* bottom_rois) { CUDA_1D_KERNEL_LOOP(index, nthreads) { // (n, c, ph, pw) is an element in the pooled output int n = index; int w = n % width; n /= width; int h = n % height; n /= height; int c = n % channels; n /= channels; float gradient = 0; // Accumulate gradient over all ROIs that pooled this element for (int roi_n = 0; roi_n < num_rois; ++roi_n) { const float* offset_bottom_rois = bottom_rois + roi_n * 5; int roi_batch_ind = offset_bottom_rois[0]; // Skip if ROI's batch index doesn't match n if (n != roi_batch_ind) { continue; } int roi_start_w = round(offset_bottom_rois[1] * spatial_scale); int roi_start_h = round(offset_bottom_rois[2] * spatial_scale); int roi_end_w = round(offset_bottom_rois[3] * spatial_scale); int roi_end_h = round(offset_bottom_rois[4] * spatial_scale); // Skip if ROI doesn't include (h, w) const bool in_roi = (w >= roi_start_w && w <= roi_end_w && h >= roi_start_h && h <= roi_end_h); if (!in_roi) { continue; } int offset = roi_n * pooled_height * pooled_width * channels; const float* offset_top_diff = top_diff + offset; const int* offset_argmax_data = argmax_data + offset; // Compute feasible set of pooled units that could have pooled // this bottom unit // Force malformed ROIs to be 1x1 int roi_width = fmaxf(roi_end_w - roi_start_w + 1, 1); int roi_height = fmaxf(roi_end_h - roi_start_h + 1, 1); float bin_size_h = (float)(roi_height) / (float)(pooled_height); float bin_size_w = (float)(roi_width) / (float)(pooled_width); int phstart = floor((float)(h - roi_start_h) / bin_size_h); int phend = ceil((float)(h - roi_start_h + 1) / bin_size_h); int pwstart = floor((float)(w - roi_start_w) / bin_size_w); int pwend = ceil((float)(w - roi_start_w + 1) / bin_size_w); phstart = fminf(fmaxf(phstart, 0), pooled_height); phend = fminf(fmaxf(phend, 0), pooled_height); pwstart = fminf(fmaxf(pwstart, 0), pooled_width); pwend = fminf(fmaxf(pwend, 0), pooled_width); for (int ph = phstart; ph < phend; ++ph) { for (int pw = pwstart; pw < pwend; ++pw) { if (offset_argmax_data[(c * pooled_height + ph) * pooled_width + pw] == index) { gradient += offset_top_diff[(c * pooled_height + ph) * pooled_width + pw]; } } } } bottom_diff[index] = gradient; } } int ROIPoolBackwardLaucher(const float* top_diff, const float spatial_scale, const int batch_size, const int num_rois, const int height, const int width, const int channels, const int pooled_height, const int pooled_width, const float* bottom_rois, float* bottom_diff, const int* argmax_data, cudaStream_t stream) { const int kThreadsPerBlock = 1024; int output_size = batch_size * height * width * channels; cudaError_t err; ROIPoolBackward<<<(output_size + kThreadsPerBlock - 1) / kThreadsPerBlock, kThreadsPerBlock, 0, stream>>>( output_size, top_diff, argmax_data, num_rois, spatial_scale, height, width, channels, pooled_height, pooled_width, bottom_diff, bottom_rois); // dim3 blocks(DIVUP(output_size, kThreadsPerBlock), // DIVUP(output_size, kThreadsPerBlock)); // dim3 threads(kThreadsPerBlock); // // ROIPoolBackward<<>>( // output_size, top_diff, argmax_data, num_rois, spatial_scale, height, width, channels, pooled_height, // pooled_width, bottom_diff, bottom_rois); err = cudaGetLastError(); if(cudaSuccess != err) { fprintf( stderr, "cudaCheckError() failed : %s\n", cudaGetErrorString( err ) ); exit( -1 ); } return 1; } // #ifdef __cplusplus // } // #endif ================================================ FILE: lib/model/roi_pooling/src/roi_pooling_kernel.h ================================================ #ifndef _ROI_POOLING_KERNEL #define _ROI_POOLING_KERNEL #ifdef __cplusplus extern "C" { #endif int ROIPoolForwardLaucher( const float* bottom_data, const float spatial_scale, const int num_rois, const int height, const int width, const int channels, const int pooled_height, const int pooled_width, const float* bottom_rois, float* top_data, int* argmax_data, cudaStream_t stream); int ROIPoolBackwardLaucher(const float* top_diff, const float spatial_scale, const int batch_size, const int num_rois, const int height, const int width, const int channels, const int pooled_height, const int pooled_width, const float* bottom_rois, float* bottom_diff, const int* argmax_data, cudaStream_t stream); #ifdef __cplusplus } #endif #endif ================================================ FILE: lib/model/rpn/__init__.py ================================================ ================================================ FILE: lib/model/rpn/anchor_target_layer.py ================================================ from __future__ import absolute_import # -------------------------------------------------------- # Faster R-CNN # Copyright (c) 2015 Microsoft # Licensed under The MIT License [see LICENSE for details] # Written by Ross Girshick and Sean Bell # -------------------------------------------------------- # -------------------------------------------------------- # Reorganized and modified by Jianwei Yang and Jiasen Lu # -------------------------------------------------------- import torch import torch.nn as nn import numpy as np import numpy.random as npr from model.utils.config import cfg from .generate_anchors import generate_anchors from .bbox_transform import clip_boxes, bbox_overlaps_batch, bbox_transform_batch import pdb DEBUG = False try: long # Python 2 except NameError: long = int # Python 3 class _AnchorTargetLayer(nn.Module): """ Assign anchors to ground-truth targets. Produces anchor classification labels and bounding-box regression targets. """ def __init__(self, feat_stride, scales, ratios): super(_AnchorTargetLayer, self).__init__() self._feat_stride = feat_stride self._scales = scales anchor_scales = scales self._anchors = torch.from_numpy(generate_anchors(scales=np.array(anchor_scales), ratios=np.array(ratios))).float() self._num_anchors = self._anchors.size(0) # allow boxes to sit over the edge by a small amount self._allowed_border = 0 # default is 0 def forward(self, input): # Algorithm: # # for each (H, W) location i # generate 9 anchor boxes centered on cell i # apply predicted bbox deltas at cell i to each of the 9 anchors # filter out-of-image anchors rpn_cls_score = input[0] gt_boxes = input[1] im_info = input[2] num_boxes = input[3] # map of shape (..., H, W) height, width = rpn_cls_score.size(2), rpn_cls_score.size(3) batch_size = gt_boxes.size(0) feat_height, feat_width = rpn_cls_score.size(2), rpn_cls_score.size(3) shift_x = np.arange(0, feat_width) * self._feat_stride shift_y = np.arange(0, feat_height) * self._feat_stride shift_x, shift_y = np.meshgrid(shift_x, shift_y) shifts = torch.from_numpy(np.vstack((shift_x.ravel(), shift_y.ravel(), shift_x.ravel(), shift_y.ravel())).transpose()) shifts = shifts.contiguous().type_as(rpn_cls_score).float() A = self._num_anchors K = shifts.size(0) self._anchors = self._anchors.type_as(gt_boxes) # move to specific gpu. all_anchors = self._anchors.view(1, A, 4) + shifts.view(K, 1, 4) all_anchors = all_anchors.view(K * A, 4) total_anchors = int(K * A) keep = ((all_anchors[:, 0] >= -self._allowed_border) & (all_anchors[:, 1] >= -self._allowed_border) & (all_anchors[:, 2] < long(im_info[0][1]) + self._allowed_border) & (all_anchors[:, 3] < long(im_info[0][0]) + self._allowed_border)) inds_inside = torch.nonzero(keep).view(-1) # keep only inside anchors anchors = all_anchors[inds_inside, :] # label: 1 is positive, 0 is negative, -1 is dont care labels = gt_boxes.new(batch_size, inds_inside.size(0)).fill_(-1) bbox_inside_weights = gt_boxes.new(batch_size, inds_inside.size(0)).zero_() bbox_outside_weights = gt_boxes.new(batch_size, inds_inside.size(0)).zero_() overlaps = bbox_overlaps_batch(anchors, gt_boxes) max_overlaps, argmax_overlaps = torch.max(overlaps, 2) gt_max_overlaps, _ = torch.max(overlaps, 1) if not cfg.TRAIN.RPN_CLOBBER_POSITIVES: labels[max_overlaps < cfg.TRAIN.RPN_NEGATIVE_OVERLAP] = 0 gt_max_overlaps[gt_max_overlaps==0] = 1e-5 keep = torch.sum(overlaps.eq(gt_max_overlaps.view(batch_size,1,-1).expand_as(overlaps)), 2) if torch.sum(keep) > 0: labels[keep>0] = 1 # fg label: above threshold IOU labels[max_overlaps >= cfg.TRAIN.RPN_POSITIVE_OVERLAP] = 1 if cfg.TRAIN.RPN_CLOBBER_POSITIVES: labels[max_overlaps < cfg.TRAIN.RPN_NEGATIVE_OVERLAP] = 0 num_fg = int(cfg.TRAIN.RPN_FG_FRACTION * cfg.TRAIN.RPN_BATCHSIZE) sum_fg = torch.sum((labels == 1).int(), 1) sum_bg = torch.sum((labels == 0).int(), 1) for i in range(batch_size): # subsample positive labels if we have too many if sum_fg[i] > num_fg: fg_inds = torch.nonzero(labels[i] == 1).view(-1) # torch.randperm seems has a bug on multi-gpu setting that cause the segfault. # See https://github.com/pytorch/pytorch/issues/1868 for more details. # use numpy instead. #rand_num = torch.randperm(fg_inds.size(0)).type_as(gt_boxes).long() rand_num = torch.from_numpy(np.random.permutation(fg_inds.size(0))).type_as(gt_boxes).long() disable_inds = fg_inds[rand_num[:fg_inds.size(0)-num_fg]] labels[i][disable_inds] = -1 # num_bg = cfg.TRAIN.RPN_BATCHSIZE - sum_fg[i] num_bg = cfg.TRAIN.RPN_BATCHSIZE - torch.sum((labels == 1).int(), 1)[i] # subsample negative labels if we have too many if sum_bg[i] > num_bg: bg_inds = torch.nonzero(labels[i] == 0).view(-1) #rand_num = torch.randperm(bg_inds.size(0)).type_as(gt_boxes).long() rand_num = torch.from_numpy(np.random.permutation(bg_inds.size(0))).type_as(gt_boxes).long() disable_inds = bg_inds[rand_num[:bg_inds.size(0)-num_bg]] labels[i][disable_inds] = -1 offset = torch.arange(0, batch_size)*gt_boxes.size(1) argmax_overlaps = argmax_overlaps + offset.view(batch_size, 1).type_as(argmax_overlaps) bbox_targets = _compute_targets_batch(anchors, gt_boxes.view(-1,5)[argmax_overlaps.view(-1), :].view(batch_size, -1, 5)) # use a single value instead of 4 values for easy index. bbox_inside_weights[labels==1] = cfg.TRAIN.RPN_BBOX_INSIDE_WEIGHTS[0] if cfg.TRAIN.RPN_POSITIVE_WEIGHT < 0: num_examples = torch.sum(labels[i] >= 0) positive_weights = 1.0 / num_examples.item() negative_weights = 1.0 / num_examples.item() else: assert ((cfg.TRAIN.RPN_POSITIVE_WEIGHT > 0) & (cfg.TRAIN.RPN_POSITIVE_WEIGHT < 1)) bbox_outside_weights[labels == 1] = positive_weights bbox_outside_weights[labels == 0] = negative_weights labels = _unmap(labels, total_anchors, inds_inside, batch_size, fill=-1) bbox_targets = _unmap(bbox_targets, total_anchors, inds_inside, batch_size, fill=0) bbox_inside_weights = _unmap(bbox_inside_weights, total_anchors, inds_inside, batch_size, fill=0) bbox_outside_weights = _unmap(bbox_outside_weights, total_anchors, inds_inside, batch_size, fill=0) outputs = [] labels = labels.view(batch_size, height, width, A).permute(0,3,1,2).contiguous() labels = labels.view(batch_size, 1, A * height, width) outputs.append(labels) bbox_targets = bbox_targets.view(batch_size, height, width, A*4).permute(0,3,1,2).contiguous() outputs.append(bbox_targets) anchors_count = bbox_inside_weights.size(1) bbox_inside_weights = bbox_inside_weights.view(batch_size,anchors_count,1).expand(batch_size, anchors_count, 4) bbox_inside_weights = bbox_inside_weights.contiguous().view(batch_size, height, width, 4*A)\ .permute(0,3,1,2).contiguous() outputs.append(bbox_inside_weights) bbox_outside_weights = bbox_outside_weights.view(batch_size,anchors_count,1).expand(batch_size, anchors_count, 4) bbox_outside_weights = bbox_outside_weights.contiguous().view(batch_size, height, width, 4*A)\ .permute(0,3,1,2).contiguous() outputs.append(bbox_outside_weights) return outputs def backward(self, top, propagate_down, bottom): """This layer does not propagate gradients.""" pass def reshape(self, bottom, top): """Reshaping happens during the call to forward.""" pass def _unmap(data, count, inds, batch_size, fill=0): """ Unmap a subset of item (data) back to the original set of items (of size count) """ if data.dim() == 2: ret = torch.Tensor(batch_size, count).fill_(fill).type_as(data) ret[:, inds] = data else: ret = torch.Tensor(batch_size, count, data.size(2)).fill_(fill).type_as(data) ret[:, inds,:] = data return ret def _compute_targets_batch(ex_rois, gt_rois): """Compute bounding-box regression targets for an image.""" return bbox_transform_batch(ex_rois, gt_rois[:, :, :4]) ================================================ FILE: lib/model/rpn/bbox_transform.py ================================================ # -------------------------------------------------------- # Fast R-CNN # Copyright (c) 2015 Microsoft # Licensed under The MIT License [see LICENSE for details] # Written by Ross Girshick # -------------------------------------------------------- # -------------------------------------------------------- # Reorganized and modified by Jianwei Yang and Jiasen Lu # -------------------------------------------------------- import torch import numpy as np import pdb def bbox_transform(ex_rois, gt_rois): ex_widths = ex_rois[:, 2] - ex_rois[:, 0] + 1.0 ex_heights = ex_rois[:, 3] - ex_rois[:, 1] + 1.0 ex_ctr_x = ex_rois[:, 0] + 0.5 * ex_widths ex_ctr_y = ex_rois[:, 1] + 0.5 * ex_heights gt_widths = gt_rois[:, 2] - gt_rois[:, 0] + 1.0 gt_heights = gt_rois[:, 3] - gt_rois[:, 1] + 1.0 gt_ctr_x = gt_rois[:, 0] + 0.5 * gt_widths gt_ctr_y = gt_rois[:, 1] + 0.5 * gt_heights targets_dx = (gt_ctr_x - ex_ctr_x) / ex_widths targets_dy = (gt_ctr_y - ex_ctr_y) / ex_heights targets_dw = torch.log(gt_widths / ex_widths) targets_dh = torch.log(gt_heights / ex_heights) targets = torch.stack( (targets_dx, targets_dy, targets_dw, targets_dh),1) return targets def bbox_transform_batch(ex_rois, gt_rois): if ex_rois.dim() == 2: ex_widths = ex_rois[:, 2] - ex_rois[:, 0] + 1.0 ex_heights = ex_rois[:, 3] - ex_rois[:, 1] + 1.0 ex_ctr_x = ex_rois[:, 0] + 0.5 * ex_widths ex_ctr_y = ex_rois[:, 1] + 0.5 * ex_heights gt_widths = gt_rois[:, :, 2] - gt_rois[:, :, 0] + 1.0 gt_heights = gt_rois[:, :, 3] - gt_rois[:, :, 1] + 1.0 gt_ctr_x = gt_rois[:, :, 0] + 0.5 * gt_widths gt_ctr_y = gt_rois[:, :, 1] + 0.5 * gt_heights targets_dx = (gt_ctr_x - ex_ctr_x.view(1,-1).expand_as(gt_ctr_x)) / ex_widths targets_dy = (gt_ctr_y - ex_ctr_y.view(1,-1).expand_as(gt_ctr_y)) / ex_heights targets_dw = torch.log(gt_widths / ex_widths.view(1,-1).expand_as(gt_widths)) targets_dh = torch.log(gt_heights / ex_heights.view(1,-1).expand_as(gt_heights)) elif ex_rois.dim() == 3: ex_widths = ex_rois[:, :, 2] - ex_rois[:, :, 0] + 1.0 ex_heights = ex_rois[:,:, 3] - ex_rois[:,:, 1] + 1.0 ex_ctr_x = ex_rois[:, :, 0] + 0.5 * ex_widths ex_ctr_y = ex_rois[:, :, 1] + 0.5 * ex_heights gt_widths = gt_rois[:, :, 2] - gt_rois[:, :, 0] + 1.0 gt_heights = gt_rois[:, :, 3] - gt_rois[:, :, 1] + 1.0 gt_ctr_x = gt_rois[:, :, 0] + 0.5 * gt_widths gt_ctr_y = gt_rois[:, :, 1] + 0.5 * gt_heights targets_dx = (gt_ctr_x - ex_ctr_x) / ex_widths targets_dy = (gt_ctr_y - ex_ctr_y) / ex_heights targets_dw = torch.log(gt_widths / ex_widths) targets_dh = torch.log(gt_heights / ex_heights) else: raise ValueError('ex_roi input dimension is not correct.') targets = torch.stack( (targets_dx, targets_dy, targets_dw, targets_dh),2) return targets def bbox_transform_inv(boxes, deltas, batch_size): widths = boxes[:, :, 2] - boxes[:, :, 0] + 1.0 heights = boxes[:, :, 3] - boxes[:, :, 1] + 1.0 ctr_x = boxes[:, :, 0] + 0.5 * widths ctr_y = boxes[:, :, 1] + 0.5 * heights dx = deltas[:, :, 0::4] dy = deltas[:, :, 1::4] dw = deltas[:, :, 2::4] dh = deltas[:, :, 3::4] pred_ctr_x = dx * widths.unsqueeze(2) + ctr_x.unsqueeze(2) pred_ctr_y = dy * heights.unsqueeze(2) + ctr_y.unsqueeze(2) pred_w = torch.exp(dw) * widths.unsqueeze(2) pred_h = torch.exp(dh) * heights.unsqueeze(2) pred_boxes = deltas.clone() # x1 pred_boxes[:, :, 0::4] = pred_ctr_x - 0.5 * pred_w # y1 pred_boxes[:, :, 1::4] = pred_ctr_y - 0.5 * pred_h # x2 pred_boxes[:, :, 2::4] = pred_ctr_x + 0.5 * pred_w # y2 pred_boxes[:, :, 3::4] = pred_ctr_y + 0.5 * pred_h return pred_boxes def clip_boxes_batch(boxes, im_shape, batch_size): """ Clip boxes to image boundaries. """ num_rois = boxes.size(1) boxes[boxes < 0] = 0 # batch_x = (im_shape[:,0]-1).view(batch_size, 1).expand(batch_size, num_rois) # batch_y = (im_shape[:,1]-1).view(batch_size, 1).expand(batch_size, num_rois) batch_x = im_shape[:, 1] - 1 batch_y = im_shape[:, 0] - 1 boxes[:,:,0][boxes[:,:,0] > batch_x] = batch_x boxes[:,:,1][boxes[:,:,1] > batch_y] = batch_y boxes[:,:,2][boxes[:,:,2] > batch_x] = batch_x boxes[:,:,3][boxes[:,:,3] > batch_y] = batch_y return boxes def clip_boxes(boxes, im_shape, batch_size): for i in range(batch_size): boxes[i,:,0::4].clamp_(0, im_shape[i, 1]-1) boxes[i,:,1::4].clamp_(0, im_shape[i, 0]-1) boxes[i,:,2::4].clamp_(0, im_shape[i, 1]-1) boxes[i,:,3::4].clamp_(0, im_shape[i, 0]-1) return boxes def bbox_overlaps(anchors, gt_boxes): """ anchors: (N, 4) ndarray of float gt_boxes: (K, 4) ndarray of float overlaps: (N, K) ndarray of overlap between boxes and query_boxes """ N = anchors.size(0) K = gt_boxes.size(0) gt_boxes_area = ((gt_boxes[:,2] - gt_boxes[:,0] + 1) * (gt_boxes[:,3] - gt_boxes[:,1] + 1)).view(1, K) anchors_area = ((anchors[:,2] - anchors[:,0] + 1) * (anchors[:,3] - anchors[:,1] + 1)).view(N, 1) boxes = anchors.view(N, 1, 4).expand(N, K, 4) query_boxes = gt_boxes.view(1, K, 4).expand(N, K, 4) iw = (torch.min(boxes[:,:,2], query_boxes[:,:,2]) - torch.max(boxes[:,:,0], query_boxes[:,:,0]) + 1) iw[iw < 0] = 0 ih = (torch.min(boxes[:,:,3], query_boxes[:,:,3]) - torch.max(boxes[:,:,1], query_boxes[:,:,1]) + 1) ih[ih < 0] = 0 ua = anchors_area + gt_boxes_area - (iw * ih) overlaps = iw * ih / ua return overlaps def bbox_overlaps_batch(anchors, gt_boxes): """ anchors: (N, 4) ndarray of float gt_boxes: (b, K, 5) ndarray of float overlaps: (N, K) ndarray of overlap between boxes and query_boxes """ batch_size = gt_boxes.size(0) if anchors.dim() == 2: N = anchors.size(0) K = gt_boxes.size(1) anchors = anchors.view(1, N, 4).expand(batch_size, N, 4).contiguous() gt_boxes = gt_boxes[:,:,:4].contiguous() gt_boxes_x = (gt_boxes[:,:,2] - gt_boxes[:,:,0] + 1) gt_boxes_y = (gt_boxes[:,:,3] - gt_boxes[:,:,1] + 1) gt_boxes_area = (gt_boxes_x * gt_boxes_y).view(batch_size, 1, K) anchors_boxes_x = (anchors[:,:,2] - anchors[:,:,0] + 1) anchors_boxes_y = (anchors[:,:,3] - anchors[:,:,1] + 1) anchors_area = (anchors_boxes_x * anchors_boxes_y).view(batch_size, N, 1) gt_area_zero = (gt_boxes_x == 1) & (gt_boxes_y == 1) anchors_area_zero = (anchors_boxes_x == 1) & (anchors_boxes_y == 1) boxes = anchors.view(batch_size, N, 1, 4).expand(batch_size, N, K, 4) query_boxes = gt_boxes.view(batch_size, 1, K, 4).expand(batch_size, N, K, 4) iw = (torch.min(boxes[:,:,:,2], query_boxes[:,:,:,2]) - torch.max(boxes[:,:,:,0], query_boxes[:,:,:,0]) + 1) iw[iw < 0] = 0 ih = (torch.min(boxes[:,:,:,3], query_boxes[:,:,:,3]) - torch.max(boxes[:,:,:,1], query_boxes[:,:,:,1]) + 1) ih[ih < 0] = 0 ua = anchors_area + gt_boxes_area - (iw * ih) overlaps = iw * ih / ua # mask the overlap here. overlaps.masked_fill_(gt_area_zero.view(batch_size, 1, K).expand(batch_size, N, K), 0) overlaps.masked_fill_(anchors_area_zero.view(batch_size, N, 1).expand(batch_size, N, K), -1) elif anchors.dim() == 3: N = anchors.size(1) K = gt_boxes.size(1) if anchors.size(2) == 4: anchors = anchors[:,:,:4].contiguous() else: anchors = anchors[:,:,1:5].contiguous() gt_boxes = gt_boxes[:,:,:4].contiguous() gt_boxes_x = (gt_boxes[:,:,2] - gt_boxes[:,:,0] + 1) gt_boxes_y = (gt_boxes[:,:,3] - gt_boxes[:,:,1] + 1) gt_boxes_area = (gt_boxes_x * gt_boxes_y).view(batch_size, 1, K) anchors_boxes_x = (anchors[:,:,2] - anchors[:,:,0] + 1) anchors_boxes_y = (anchors[:,:,3] - anchors[:,:,1] + 1) anchors_area = (anchors_boxes_x * anchors_boxes_y).view(batch_size, N, 1) gt_area_zero = (gt_boxes_x == 1) & (gt_boxes_y == 1) anchors_area_zero = (anchors_boxes_x == 1) & (anchors_boxes_y == 1) boxes = anchors.view(batch_size, N, 1, 4).expand(batch_size, N, K, 4) query_boxes = gt_boxes.view(batch_size, 1, K, 4).expand(batch_size, N, K, 4) iw = (torch.min(boxes[:,:,:,2], query_boxes[:,:,:,2]) - torch.max(boxes[:,:,:,0], query_boxes[:,:,:,0]) + 1) iw[iw < 0] = 0 ih = (torch.min(boxes[:,:,:,3], query_boxes[:,:,:,3]) - torch.max(boxes[:,:,:,1], query_boxes[:,:,:,1]) + 1) ih[ih < 0] = 0 ua = anchors_area + gt_boxes_area - (iw * ih) # Intersection (iw * ih) divided by Union (ua) overlaps = iw * ih / ua # mask the overlap here. overlaps.masked_fill_(gt_area_zero.view(batch_size, 1, K).expand(batch_size, N, K), 0) overlaps.masked_fill_(anchors_area_zero.view(batch_size, N, 1).expand(batch_size, N, K), -1) else: raise ValueError('anchors input dimension is not correct.') return overlaps ================================================ FILE: lib/model/rpn/generate_anchors.py ================================================ from __future__ import print_function # -------------------------------------------------------- # Faster R-CNN # Copyright (c) 2015 Microsoft # Licensed under The MIT License [see LICENSE for details] # Written by Ross Girshick and Sean Bell # -------------------------------------------------------- import numpy as np import pdb # Verify that we compute the same anchors as Shaoqing's matlab implementation: # # >> load output/rpn_cachedir/faster_rcnn_VOC2007_ZF_stage1_rpn/anchors.mat # >> anchors # # anchors = # # -83 -39 100 56 # -175 -87 192 104 # -359 -183 376 200 # -55 -55 72 72 # -119 -119 136 136 # -247 -247 264 264 # -35 -79 52 96 # -79 -167 96 184 # -167 -343 184 360 #array([[ -83., -39., 100., 56.], # [-175., -87., 192., 104.], # [-359., -183., 376., 200.], # [ -55., -55., 72., 72.], # [-119., -119., 136., 136.], # [-247., -247., 264., 264.], # [ -35., -79., 52., 96.], # [ -79., -167., 96., 184.], # [-167., -343., 184., 360.]]) try: xrange # Python 2 except NameError: xrange = range # Python 3 def generate_anchors(base_size=16, ratios=[0.5, 1, 2], scales=2**np.arange(3, 6)): """ Generate anchor (reference) windows by enumerating aspect ratios X scales wrt a reference (0, 0, 15, 15) window. """ base_anchor = np.array([1, 1, base_size, base_size]) - 1 ratio_anchors = _ratio_enum(base_anchor, ratios) anchors = np.vstack([_scale_enum(ratio_anchors[i, :], scales) for i in xrange(ratio_anchors.shape[0])]) return anchors def _whctrs(anchor): """ Return width, height, x center, and y center for an anchor (window). """ w = anchor[2] - anchor[0] + 1 h = anchor[3] - anchor[1] + 1 x_ctr = anchor[0] + 0.5 * (w - 1) y_ctr = anchor[1] + 0.5 * (h - 1) return w, h, x_ctr, y_ctr def _mkanchors(ws, hs, x_ctr, y_ctr): """ Given a vector of widths (ws) and heights (hs) around a center (x_ctr, y_ctr), output a set of anchors (windows). """ ws = ws[:, np.newaxis] hs = hs[:, np.newaxis] anchors = np.hstack((x_ctr - 0.5 * (ws - 1), y_ctr - 0.5 * (hs - 1), x_ctr + 0.5 * (ws - 1), y_ctr + 0.5 * (hs - 1))) return anchors def _ratio_enum(anchor, ratios): """ Enumerate a set of anchors for each aspect ratio wrt an anchor. """ w, h, x_ctr, y_ctr = _whctrs(anchor) size = w * h size_ratios = size / ratios ws = np.round(np.sqrt(size_ratios)) hs = np.round(ws * ratios) anchors = _mkanchors(ws, hs, x_ctr, y_ctr) return anchors def _scale_enum(anchor, scales): """ Enumerate a set of anchors for each scale wrt an anchor. """ w, h, x_ctr, y_ctr = _whctrs(anchor) ws = w * scales hs = h * scales anchors = _mkanchors(ws, hs, x_ctr, y_ctr) return anchors if __name__ == '__main__': import time t = time.time() a = generate_anchors() print(time.time() - t) print(a) from IPython import embed; embed() ================================================ FILE: lib/model/rpn/proposal_layer.py ================================================ from __future__ import absolute_import # -------------------------------------------------------- # Faster R-CNN # Copyright (c) 2015 Microsoft # Licensed under The MIT License [see LICENSE for details] # Written by Ross Girshick and Sean Bell # -------------------------------------------------------- # -------------------------------------------------------- # Reorganized and modified by Jianwei Yang and Jiasen Lu # -------------------------------------------------------- import torch import torch.nn as nn import numpy as np import math import yaml from model.utils.config import cfg from .generate_anchors import generate_anchors from .bbox_transform import bbox_transform_inv, clip_boxes, clip_boxes_batch from model.nms.nms_wrapper import nms import pdb DEBUG = False class _ProposalLayer(nn.Module): """ Outputs object detection proposals by applying estimated bounding-box transformations to a set of regular boxes (called "anchors"). """ def __init__(self, feat_stride, scales, ratios): super(_ProposalLayer, self).__init__() self._feat_stride = feat_stride self._anchors = torch.from_numpy(generate_anchors(scales=np.array(scales), ratios=np.array(ratios))).float() self._num_anchors = self._anchors.size(0) # rois blob: holds R regions of interest, each is a 5-tuple # (n, x1, y1, x2, y2) specifying an image batch index n and a # rectangle (x1, y1, x2, y2) # top[0].reshape(1, 5) # # # scores blob: holds scores for R regions of interest # if len(top) > 1: # top[1].reshape(1, 1, 1, 1) def forward(self, input): # Algorithm: # # for each (H, W) location i # generate A anchor boxes centered on cell i # apply predicted bbox deltas at cell i to each of the A anchors # clip predicted boxes to image # remove predicted boxes with either height or width < threshold # sort all (proposal, score) pairs by score from highest to lowest # take top pre_nms_topN proposals before NMS # apply NMS with threshold 0.7 to remaining proposals # take after_nms_topN proposals after NMS # return the top proposals (-> RoIs top, scores top) # the first set of _num_anchors channels are bg probs # the second set are the fg probs scores = input[0][:, self._num_anchors:, :, :] bbox_deltas = input[1] im_info = input[2] cfg_key = input[3] pre_nms_topN = cfg[cfg_key].RPN_PRE_NMS_TOP_N post_nms_topN = cfg[cfg_key].RPN_POST_NMS_TOP_N nms_thresh = cfg[cfg_key].RPN_NMS_THRESH min_size = cfg[cfg_key].RPN_MIN_SIZE batch_size = bbox_deltas.size(0) feat_height, feat_width = scores.size(2), scores.size(3) shift_x = np.arange(0, feat_width) * self._feat_stride shift_y = np.arange(0, feat_height) * self._feat_stride shift_x, shift_y = np.meshgrid(shift_x, shift_y) shifts = torch.from_numpy(np.vstack((shift_x.ravel(), shift_y.ravel(), shift_x.ravel(), shift_y.ravel())).transpose()) shifts = shifts.contiguous().type_as(scores).float() A = self._num_anchors K = shifts.size(0) self._anchors = self._anchors.type_as(scores) # anchors = self._anchors.view(1, A, 4) + shifts.view(1, K, 4).permute(1, 0, 2).contiguous() anchors = self._anchors.view(1, A, 4) + shifts.view(K, 1, 4) anchors = anchors.view(1, K * A, 4).expand(batch_size, K * A, 4) # Transpose and reshape predicted bbox transformations to get them # into the same order as the anchors: bbox_deltas = bbox_deltas.permute(0, 2, 3, 1).contiguous() bbox_deltas = bbox_deltas.view(batch_size, -1, 4) # Same story for the scores: scores = scores.permute(0, 2, 3, 1).contiguous() scores = scores.view(batch_size, -1) # Convert anchors into proposals via bbox transformations proposals = bbox_transform_inv(anchors, bbox_deltas, batch_size) # 2. clip predicted boxes to image proposals = clip_boxes(proposals, im_info, batch_size) # proposals = clip_boxes_batch(proposals, im_info, batch_size) # assign the score to 0 if it's non keep. # keep = self._filter_boxes(proposals, min_size * im_info[:, 2]) # trim keep index to make it euqal over batch # keep_idx = torch.cat(tuple(keep_idx), 0) # scores_keep = scores.view(-1)[keep_idx].view(batch_size, trim_size) # proposals_keep = proposals.view(-1, 4)[keep_idx, :].contiguous().view(batch_size, trim_size, 4) # _, order = torch.sort(scores_keep, 1, True) scores_keep = scores proposals_keep = proposals _, order = torch.sort(scores_keep, 1, True) output = scores.new(batch_size, post_nms_topN, 5).zero_() for i in range(batch_size): # # 3. remove predicted boxes with either height or width < threshold # # (NOTE: convert min_size to input image scale stored in im_info[2]) proposals_single = proposals_keep[i] scores_single = scores_keep[i] # # 4. sort all (proposal, score) pairs by score from highest to lowest # # 5. take top pre_nms_topN (e.g. 6000) order_single = order[i] if pre_nms_topN > 0 and pre_nms_topN < scores_keep.numel(): order_single = order_single[:pre_nms_topN] proposals_single = proposals_single[order_single, :] scores_single = scores_single[order_single].view(-1,1) # 6. apply nms (e.g. threshold = 0.7) # 7. take after_nms_topN (e.g. 300) # 8. return the top proposals (-> RoIs top) keep_idx_i = nms(torch.cat((proposals_single, scores_single), 1), nms_thresh, force_cpu=not cfg.USE_GPU_NMS) keep_idx_i = keep_idx_i.long().view(-1) if post_nms_topN > 0: keep_idx_i = keep_idx_i[:post_nms_topN] proposals_single = proposals_single[keep_idx_i, :] scores_single = scores_single[keep_idx_i, :] # padding 0 at the end. num_proposal = proposals_single.size(0) output[i,:,0] = i output[i,:num_proposal,1:] = proposals_single return output def backward(self, top, propagate_down, bottom): """This layer does not propagate gradients.""" pass def reshape(self, bottom, top): """Reshaping happens during the call to forward.""" pass def _filter_boxes(self, boxes, min_size): """Remove all boxes with any side smaller than min_size.""" ws = boxes[:, :, 2] - boxes[:, :, 0] + 1 hs = boxes[:, :, 3] - boxes[:, :, 1] + 1 keep = ((ws >= min_size.view(-1,1).expand_as(ws)) & (hs >= min_size.view(-1,1).expand_as(hs))) return keep ================================================ FILE: lib/model/rpn/proposal_target_layer_cascade.py ================================================ from __future__ import absolute_import # -------------------------------------------------------- # Faster R-CNN # Copyright (c) 2015 Microsoft # Licensed under The MIT License [see LICENSE for details] # Written by Ross Girshick and Sean Bell # -------------------------------------------------------- # -------------------------------------------------------- # Reorganized and modified by Jianwei Yang and Jiasen Lu # -------------------------------------------------------- import torch import torch.nn as nn import numpy as np import numpy.random as npr from ..utils.config import cfg from .bbox_transform import bbox_overlaps_batch, bbox_transform_batch import pdb class _ProposalTargetLayer(nn.Module): """ Assign object detection proposals to ground-truth targets. Produces proposal classification labels and bounding-box regression targets. """ def __init__(self, nclasses): super(_ProposalTargetLayer, self).__init__() self._num_classes = nclasses self.BBOX_NORMALIZE_MEANS = torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_MEANS) self.BBOX_NORMALIZE_STDS = torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_STDS) self.BBOX_INSIDE_WEIGHTS = torch.FloatTensor(cfg.TRAIN.BBOX_INSIDE_WEIGHTS) def forward(self, all_rois, gt_boxes, num_boxes): self.BBOX_NORMALIZE_MEANS = self.BBOX_NORMALIZE_MEANS.type_as(gt_boxes) self.BBOX_NORMALIZE_STDS = self.BBOX_NORMALIZE_STDS.type_as(gt_boxes) self.BBOX_INSIDE_WEIGHTS = self.BBOX_INSIDE_WEIGHTS.type_as(gt_boxes) gt_boxes_append = gt_boxes.new(gt_boxes.size()).zero_() gt_boxes_append[:,:,1:5] = gt_boxes[:,:,:4] # Include ground-truth boxes in the set of candidate rois all_rois = torch.cat([all_rois, gt_boxes_append], 1) num_images = 1 rois_per_image = int(cfg.TRAIN.BATCH_SIZE / num_images) fg_rois_per_image = int(np.round(cfg.TRAIN.FG_FRACTION * rois_per_image)) fg_rois_per_image = 1 if fg_rois_per_image == 0 else fg_rois_per_image labels, rois, bbox_targets, bbox_inside_weights = self._sample_rois_pytorch( all_rois, gt_boxes, fg_rois_per_image, rois_per_image, self._num_classes) bbox_outside_weights = (bbox_inside_weights > 0).float() return rois, labels, bbox_targets, bbox_inside_weights, bbox_outside_weights def backward(self, top, propagate_down, bottom): """This layer does not propagate gradients.""" pass def reshape(self, bottom, top): """Reshaping happens during the call to forward.""" pass def _get_bbox_regression_labels_pytorch(self, bbox_target_data, labels_batch, num_classes): """Bounding-box regression targets (bbox_target_data) are stored in a compact form b x N x (class, tx, ty, tw, th) This function expands those targets into the 4-of-4*K representation used by the network (i.e. only one class has non-zero targets). Returns: bbox_target (ndarray): b x N x 4K blob of regression targets bbox_inside_weights (ndarray): b x N x 4K blob of loss weights """ batch_size = labels_batch.size(0) rois_per_image = labels_batch.size(1) clss = labels_batch bbox_targets = bbox_target_data.new(batch_size, rois_per_image, 4).zero_() bbox_inside_weights = bbox_target_data.new(bbox_targets.size()).zero_() for b in range(batch_size): # assert clss[b].sum() > 0 if clss[b].sum() == 0: continue inds = torch.nonzero(clss[b] > 0).view(-1) for i in range(inds.numel()): ind = inds[i] bbox_targets[b, ind, :] = bbox_target_data[b, ind, :] bbox_inside_weights[b, ind, :] = self.BBOX_INSIDE_WEIGHTS return bbox_targets, bbox_inside_weights def _compute_targets_pytorch(self, ex_rois, gt_rois): """Compute bounding-box regression targets for an image.""" assert ex_rois.size(1) == gt_rois.size(1) assert ex_rois.size(2) == 4 assert gt_rois.size(2) == 4 batch_size = ex_rois.size(0) rois_per_image = ex_rois.size(1) targets = bbox_transform_batch(ex_rois, gt_rois) if cfg.TRAIN.BBOX_NORMALIZE_TARGETS_PRECOMPUTED: # Optionally normalize targets by a precomputed mean and stdev targets = ((targets - self.BBOX_NORMALIZE_MEANS.expand_as(targets)) / self.BBOX_NORMALIZE_STDS.expand_as(targets)) return targets def _sample_rois_pytorch(self, all_rois, gt_boxes, fg_rois_per_image, rois_per_image, num_classes): """Generate a random sample of RoIs comprising foreground and background examples. """ # overlaps: (rois x gt_boxes) overlaps = bbox_overlaps_batch(all_rois, gt_boxes) max_overlaps, gt_assignment = torch.max(overlaps, 2) batch_size = overlaps.size(0) num_proposal = overlaps.size(1) num_boxes_per_img = overlaps.size(2) offset = torch.arange(0, batch_size)*gt_boxes.size(1) offset = offset.view(-1, 1).type_as(gt_assignment) + gt_assignment labels = gt_boxes[:,:,4].contiguous().view(-1).index((offset.view(-1),)).view(batch_size, -1) labels_batch = labels.new(batch_size, rois_per_image).zero_() rois_batch = all_rois.new(batch_size, rois_per_image, 5).zero_() gt_rois_batch = all_rois.new(batch_size, rois_per_image, 5).zero_() # Guard against the case when an image has fewer than max_fg_rois_per_image # foreground RoIs for i in range(batch_size): fg_inds = torch.nonzero(max_overlaps[i] >= cfg.TRAIN.FG_THRESH).view(-1) fg_num_rois = fg_inds.numel() # Select background RoIs as those within [BG_THRESH_LO, BG_THRESH_HI) bg_inds = torch.nonzero((max_overlaps[i] < cfg.TRAIN.BG_THRESH_HI) & (max_overlaps[i] >= cfg.TRAIN.BG_THRESH_LO)).view(-1) bg_num_rois = bg_inds.numel() if fg_num_rois > 0 and bg_num_rois > 0: # sampling fg fg_rois_per_this_image = min(fg_rois_per_image, fg_num_rois) # torch.randperm seems has a bug on multi-gpu setting that cause the segfault. # See https://github.com/pytorch/pytorch/issues/1868 for more details. # use numpy instead. #rand_num = torch.randperm(fg_num_rois).long().cuda() rand_num = torch.from_numpy(np.random.permutation(fg_num_rois)).type_as(gt_boxes).long() fg_inds = fg_inds[rand_num[:fg_rois_per_this_image]] # sampling bg bg_rois_per_this_image = rois_per_image - fg_rois_per_this_image # Seems torch.rand has a bug, it will generate very large number and make an error. # We use numpy rand instead. #rand_num = (torch.rand(bg_rois_per_this_image) * bg_num_rois).long().cuda() rand_num = np.floor(np.random.rand(bg_rois_per_this_image) * bg_num_rois) rand_num = torch.from_numpy(rand_num).type_as(gt_boxes).long() bg_inds = bg_inds[rand_num] elif fg_num_rois > 0 and bg_num_rois == 0: # sampling fg #rand_num = torch.floor(torch.rand(rois_per_image) * fg_num_rois).long().cuda() rand_num = np.floor(np.random.rand(rois_per_image) * fg_num_rois) rand_num = torch.from_numpy(rand_num).type_as(gt_boxes).long() fg_inds = fg_inds[rand_num] fg_rois_per_this_image = rois_per_image bg_rois_per_this_image = 0 elif bg_num_rois > 0 and fg_num_rois == 0: # sampling bg #rand_num = torch.floor(torch.rand(rois_per_image) * bg_num_rois).long().cuda() rand_num = np.floor(np.random.rand(rois_per_image) * bg_num_rois) rand_num = torch.from_numpy(rand_num).type_as(gt_boxes).long() bg_inds = bg_inds[rand_num] bg_rois_per_this_image = rois_per_image fg_rois_per_this_image = 0 else: raise ValueError("bg_num_rois = 0 and fg_num_rois = 0, this should not happen!") # The indices that we're selecting (both fg and bg) keep_inds = torch.cat([fg_inds, bg_inds], 0) # Select sampled values from various arrays: labels_batch[i].copy_(labels[i][keep_inds]) # Clamp labels for the background RoIs to 0 if fg_rois_per_this_image < rois_per_image: labels_batch[i][fg_rois_per_this_image:] = 0 rois_batch[i] = all_rois[i][keep_inds] rois_batch[i,:,0] = i gt_rois_batch[i] = gt_boxes[i][gt_assignment[i][keep_inds]] bbox_target_data = self._compute_targets_pytorch( rois_batch[:,:,1:5], gt_rois_batch[:,:,:4]) bbox_targets, bbox_inside_weights = \ self._get_bbox_regression_labels_pytorch(bbox_target_data, labels_batch, num_classes) return labels_batch, rois_batch, bbox_targets, bbox_inside_weights ================================================ FILE: lib/model/rpn/rpn.py ================================================ from __future__ import absolute_import import torch import torch.nn as nn import torch.nn.functional as F from torch.autograd import Variable from model.utils.config import cfg from .proposal_layer import _ProposalLayer from .anchor_target_layer import _AnchorTargetLayer from model.utils.net_utils import _smooth_l1_loss import numpy as np import math import pdb import time class _RPN(nn.Module): """ region proposal network """ def __init__(self, din): super(_RPN, self).__init__() self.din = din # get depth of input feature map, e.g., 512 self.anchor_scales = cfg.ANCHOR_SCALES self.anchor_ratios = cfg.ANCHOR_RATIOS self.feat_stride = cfg.FEAT_STRIDE[0] # define the convrelu layers processing input feature map self.RPN_Conv = nn.Conv2d(self.din, 512, 3, 1, 1, bias=True) # define bg/fg classifcation score layer self.nc_score_out = len(self.anchor_scales) * len(self.anchor_ratios) * 2 # 2(bg/fg) * 9 (anchors) self.RPN_cls_score = nn.Conv2d(512, self.nc_score_out, 1, 1, 0) # define anchor box offset prediction layer self.nc_bbox_out = len(self.anchor_scales) * len(self.anchor_ratios) * 4 # 4(coords) * 9 (anchors) self.RPN_bbox_pred = nn.Conv2d(512, self.nc_bbox_out, 1, 1, 0) # define proposal layer self.RPN_proposal = _ProposalLayer(self.feat_stride, self.anchor_scales, self.anchor_ratios) # define anchor target layer self.RPN_anchor_target = _AnchorTargetLayer(self.feat_stride, self.anchor_scales, self.anchor_ratios) self.rpn_loss_cls = 0 self.rpn_loss_box = 0 @staticmethod def reshape(x, d): input_shape = x.size() x = x.view( input_shape[0], int(d), int(float(input_shape[1] * input_shape[2]) / float(d)), input_shape[3] ) return x def forward(self, base_feat, im_info, gt_boxes, num_boxes): batch_size = base_feat.size(0) # return feature map after convrelu layer rpn_conv1 = F.relu(self.RPN_Conv(base_feat), inplace=True) # get rpn classification score rpn_cls_score = self.RPN_cls_score(rpn_conv1) rpn_cls_score_reshape = self.reshape(rpn_cls_score, 2) rpn_cls_prob_reshape = F.softmax(rpn_cls_score_reshape, 1) rpn_cls_prob = self.reshape(rpn_cls_prob_reshape, self.nc_score_out) # get rpn offsets to the anchor boxes rpn_bbox_pred = self.RPN_bbox_pred(rpn_conv1) # proposal layer cfg_key = 'TRAIN' if self.training else 'TEST' rois = self.RPN_proposal((rpn_cls_prob.data, rpn_bbox_pred.data, im_info, cfg_key)) self.rpn_loss_cls = 0 self.rpn_loss_box = 0 # generating training labels and build the rpn loss if self.training: assert gt_boxes is not None rpn_data = self.RPN_anchor_target((rpn_cls_score.data, gt_boxes, im_info, num_boxes)) # compute classification loss rpn_cls_score = rpn_cls_score_reshape.permute(0, 2, 3, 1).contiguous().view(batch_size, -1, 2) rpn_label = rpn_data[0].view(batch_size, -1) rpn_keep = Variable(rpn_label.view(-1).ne(-1).nonzero().view(-1)) rpn_cls_score = torch.index_select(rpn_cls_score.view(-1,2), 0, rpn_keep) rpn_label = torch.index_select(rpn_label.view(-1), 0, rpn_keep.data) rpn_label = Variable(rpn_label.long()) self.rpn_loss_cls = F.cross_entropy(rpn_cls_score, rpn_label) fg_cnt = torch.sum(rpn_label.data.ne(0)) rpn_bbox_targets, rpn_bbox_inside_weights, rpn_bbox_outside_weights = rpn_data[1:] # compute bbox regression loss rpn_bbox_inside_weights = Variable(rpn_bbox_inside_weights) rpn_bbox_outside_weights = Variable(rpn_bbox_outside_weights) rpn_bbox_targets = Variable(rpn_bbox_targets) self.rpn_loss_box = _smooth_l1_loss(rpn_bbox_pred, rpn_bbox_targets, rpn_bbox_inside_weights, rpn_bbox_outside_weights, sigma=3, dim=[1,2,3]) return rois, self.rpn_loss_cls, self.rpn_loss_box ================================================ FILE: lib/model/utils/.gitignore ================================================ *.c *.cpp *.so ================================================ FILE: lib/model/utils/__init__.py ================================================ ================================================ FILE: lib/model/utils/bbox.pyx ================================================ # -------------------------------------------------------- # Fast R-CNN # Copyright (c) 2015 Microsoft # Licensed under The MIT License [see LICENSE for details] # Written by Sergey Karayev # -------------------------------------------------------- cimport cython import numpy as np cimport numpy as np DTYPE = np.float ctypedef np.float_t DTYPE_t def bbox_overlaps(np.ndarray[DTYPE_t, ndim=2] boxes, np.ndarray[DTYPE_t, ndim=2] query_boxes): return bbox_overlaps_c(boxes, query_boxes) cdef np.ndarray[DTYPE_t, ndim=2] bbox_overlaps_c( np.ndarray[DTYPE_t, ndim=2] boxes, np.ndarray[DTYPE_t, ndim=2] query_boxes): """ Parameters ---------- boxes: (N, 4) ndarray of float query_boxes: (K, 4) ndarray of float Returns ------- overlaps: (N, K) ndarray of overlap between boxes and query_boxes """ cdef unsigned int N = boxes.shape[0] cdef unsigned int K = query_boxes.shape[0] cdef np.ndarray[DTYPE_t, ndim=2] overlaps = np.zeros((N, K), dtype=DTYPE) cdef DTYPE_t iw, ih, box_area cdef DTYPE_t ua cdef unsigned int k, n for k in range(K): box_area = ( (query_boxes[k, 2] - query_boxes[k, 0] + 1) * (query_boxes[k, 3] - query_boxes[k, 1] + 1) ) for n in range(N): iw = ( min(boxes[n, 2], query_boxes[k, 2]) - max(boxes[n, 0], query_boxes[k, 0]) + 1 ) if iw > 0: ih = ( min(boxes[n, 3], query_boxes[k, 3]) - max(boxes[n, 1], query_boxes[k, 1]) + 1 ) if ih > 0: ua = float( (boxes[n, 2] - boxes[n, 0] + 1) * (boxes[n, 3] - boxes[n, 1] + 1) + box_area - iw * ih ) overlaps[n, k] = iw * ih / ua return overlaps def bbox_intersections( np.ndarray[DTYPE_t, ndim=2] boxes, np.ndarray[DTYPE_t, ndim=2] query_boxes): return bbox_intersections_c(boxes, query_boxes) cdef np.ndarray[DTYPE_t, ndim=2] bbox_intersections_c( np.ndarray[DTYPE_t, ndim=2] boxes, np.ndarray[DTYPE_t, ndim=2] query_boxes): """ For each query box compute the intersection ratio covered by boxes ---------- Parameters ---------- boxes: (N, 4) ndarray of float query_boxes: (K, 4) ndarray of float Returns ------- overlaps: (N, K) ndarray of intersec between boxes and query_boxes """ cdef unsigned int N = boxes.shape[0] cdef unsigned int K = query_boxes.shape[0] cdef np.ndarray[DTYPE_t, ndim=2] intersec = np.zeros((N, K), dtype=DTYPE) cdef DTYPE_t iw, ih, box_area cdef DTYPE_t ua cdef unsigned int k, n for k in range(K): box_area = ( (query_boxes[k, 2] - query_boxes[k, 0] + 1) * (query_boxes[k, 3] - query_boxes[k, 1] + 1) ) for n in range(N): iw = ( min(boxes[n, 2], query_boxes[k, 2]) - max(boxes[n, 0], query_boxes[k, 0]) + 1 ) if iw > 0: ih = ( min(boxes[n, 3], query_boxes[k, 3]) - max(boxes[n, 1], query_boxes[k, 1]) + 1 ) if ih > 0: intersec[n, k] = iw * ih / box_area return intersec ================================================ FILE: lib/model/utils/blob.py ================================================ # -------------------------------------------------------- # Fast R-CNN # Copyright (c) 2015 Microsoft # Licensed under The MIT License [see LICENSE for details] # Written by Ross Girshick # -------------------------------------------------------- """Blob helper functions.""" import numpy as np # from scipy.misc import imread, imresize import cv2 try: xrange # Python 2 except NameError: xrange = range # Python 3 def im_list_to_blob(ims): """Convert a list of images into a network input. Assumes images are already prepared (means subtracted, BGR order, ...). """ max_shape = np.array([im.shape for im in ims]).max(axis=0) num_images = len(ims) blob = np.zeros((num_images, max_shape[0], max_shape[1], 3), dtype=np.float32) for i in xrange(num_images): im = ims[i] blob[i, 0:im.shape[0], 0:im.shape[1], :] = im return blob def prep_im_for_blob(im, pixel_means, target_size, max_size): """Mean subtract and scale an image for use in a blob.""" im = im.astype(np.float32, copy=False) im -= pixel_means # im = im[:, :, ::-1] im_shape = im.shape im_size_min = np.min(im_shape[0:2]) im_size_max = np.max(im_shape[0:2]) im_scale = float(target_size) / float(im_size_min) # Prevent the biggest axis from being more than MAX_SIZE # if np.round(im_scale * im_size_max) > max_size: # im_scale = float(max_size) / float(im_size_max) # im = imresize(im, im_scale) im = cv2.resize(im, None, None, fx=im_scale, fy=im_scale, interpolation=cv2.INTER_LINEAR) return im, im_scale ================================================ FILE: lib/model/utils/config.py ================================================ from __future__ import absolute_import from __future__ import division from __future__ import print_function import os import os.path as osp import numpy as np # `pip install easydict` if you don't have it from easydict import EasyDict as edict __C = edict() # Consumers can get config by: # from fast_rcnn_config import cfg cfg = __C # # Training options # __C.TRAIN = edict() # Initial learning rate __C.TRAIN.LEARNING_RATE = 0.001 # Momentum __C.TRAIN.MOMENTUM = 0.9 # Weight decay, for regularization __C.TRAIN.WEIGHT_DECAY = 0.0005 # Factor for reducing the learning rate __C.TRAIN.GAMMA = 0.1 # Step size for reducing the learning rate, currently only support one step __C.TRAIN.STEPSIZE = [30000] # Iteration intervals for showing the loss during training, on command line interface __C.TRAIN.DISPLAY = 10 # Whether to double the learning rate for bias __C.TRAIN.DOUBLE_BIAS = True # Whether to initialize the weights with truncated normal distribution __C.TRAIN.TRUNCATED = False # Whether to have weight decay on bias as well __C.TRAIN.BIAS_DECAY = False # Whether to add ground truth boxes to the pool when sampling regions __C.TRAIN.USE_GT = False # Whether to use aspect-ratio grouping of training images, introduced merely for saving # GPU memory __C.TRAIN.ASPECT_GROUPING = False # The number of snapshots kept, older ones are deleted to save space __C.TRAIN.SNAPSHOT_KEPT = 3 # The time interval for saving tensorflow summaries __C.TRAIN.SUMMARY_INTERVAL = 180 # Scale to use during training (can list multiple scales) # The scale is the pixel size of an image's shortest side __C.TRAIN.SCALES = (600,) # Max pixel size of the longest side of a scaled input image __C.TRAIN.MAX_SIZE = 1000 # Trim size for input images to create minibatch __C.TRAIN.TRIM_HEIGHT = 600 __C.TRAIN.TRIM_WIDTH = 600 # Images to use per minibatch __C.TRAIN.IMS_PER_BATCH = 1 # Minibatch size (number of regions of interest [ROIs]) __C.TRAIN.BATCH_SIZE = 128 # Fraction of minibatch that is labeled foreground (i.e. class > 0) __C.TRAIN.FG_FRACTION = 0.25 # Overlap threshold for a ROI to be considered foreground (if >= FG_THRESH) __C.TRAIN.FG_THRESH = 0.5 # Overlap threshold for a ROI to be considered background (class = 0 if # overlap in [LO, HI)) __C.TRAIN.BG_THRESH_HI = 0.5 __C.TRAIN.BG_THRESH_LO = 0.1 # Use horizontally-flipped images during training? __C.TRAIN.USE_FLIPPED = True # Train bounding-box regressors __C.TRAIN.BBOX_REG = True # Overlap required between a ROI and ground-truth box in order for that ROI to # be used as a bounding-box regression training example __C.TRAIN.BBOX_THRESH = 0.5 # Iterations between snapshots __C.TRAIN.SNAPSHOT_ITERS = 5000 # solver.prototxt specifies the snapshot path prefix, this adds an optional # infix to yield the path: [_]_iters_XYZ.caffemodel __C.TRAIN.SNAPSHOT_PREFIX = 'res101_faster_rcnn' # __C.TRAIN.SNAPSHOT_INFIX = '' # Use a prefetch thread in roi_data_layer.layer # So far I haven't found this useful; likely more engineering work is required # __C.TRAIN.USE_PREFETCH = False # Normalize the targets (subtract empirical mean, divide by empirical stddev) __C.TRAIN.BBOX_NORMALIZE_TARGETS = True # Deprecated (inside weights) __C.TRAIN.BBOX_INSIDE_WEIGHTS = (1.0, 1.0, 1.0, 1.0) # Normalize the targets using "precomputed" (or made up) means and stdevs # (BBOX_NORMALIZE_TARGETS must also be True) __C.TRAIN.BBOX_NORMALIZE_TARGETS_PRECOMPUTED = True __C.TRAIN.BBOX_NORMALIZE_MEANS = (0.0, 0.0, 0.0, 0.0) __C.TRAIN.BBOX_NORMALIZE_STDS = (0.1, 0.1, 0.2, 0.2) # Train using these proposals __C.TRAIN.PROPOSAL_METHOD = 'gt' # Make minibatches from images that have similar aspect ratios (i.e. both # tall and thin or both short and wide) in order to avoid wasting computation # on zero-padding. # Use RPN to detect objects __C.TRAIN.HAS_RPN = True # IOU >= thresh: positive example __C.TRAIN.RPN_POSITIVE_OVERLAP = 0.7 # IOU < thresh: negative example __C.TRAIN.RPN_NEGATIVE_OVERLAP = 0.3 # If an anchor statisfied by positive and negative conditions set to negative __C.TRAIN.RPN_CLOBBER_POSITIVES = False # Max number of foreground examples __C.TRAIN.RPN_FG_FRACTION = 0.5 # Total number of examples __C.TRAIN.RPN_BATCHSIZE = 256 # NMS threshold used on RPN proposals __C.TRAIN.RPN_NMS_THRESH = 0.7 # Number of top scoring boxes to keep before apply NMS to RPN proposals __C.TRAIN.RPN_PRE_NMS_TOP_N = 12000 # Number of top scoring boxes to keep after applying NMS to RPN proposals __C.TRAIN.RPN_POST_NMS_TOP_N = 2000 # Proposal height and width both need to be greater than RPN_MIN_SIZE (at orig image scale) __C.TRAIN.RPN_MIN_SIZE = 8 # Deprecated (outside weights) __C.TRAIN.RPN_BBOX_INSIDE_WEIGHTS = (1.0, 1.0, 1.0, 1.0) # Give the positive RPN examples weight of p * 1 / {num positives} # and give negatives a weight of (1 - p) # Set to -1.0 to use uniform example weighting __C.TRAIN.RPN_POSITIVE_WEIGHT = -1.0 # Whether to use all ground truth bounding boxes for training, # For COCO, setting USE_ALL_GT to False will exclude boxes that are flagged as ''iscrowd'' __C.TRAIN.USE_ALL_GT = True # Whether to tune the batch normalization parameters during training __C.TRAIN.BN_TRAIN = False # # Testing options # __C.TEST = edict() # Scale to use during testing (can NOT list multiple scales) # The scale is the pixel size of an image's shortest side __C.TEST.SCALES = (600,) # Max pixel size of the longest side of a scaled input image __C.TEST.MAX_SIZE = 1000 # Overlap threshold used for non-maximum suppression (suppress boxes with # IoU >= this threshold) __C.TEST.NMS = 0.3 # Experimental: treat the (K+1) units in the cls_score layer as linear # predictors (trained, eg, with one-vs-rest SVMs). __C.TEST.SVM = False # Test using bounding-box regressors __C.TEST.BBOX_REG = True # Propose boxes __C.TEST.HAS_RPN = False # Test using these proposals __C.TEST.PROPOSAL_METHOD = 'gt' ## NMS threshold used on RPN proposals __C.TEST.RPN_NMS_THRESH = 0.7 ## Number of top scoring boxes to keep before apply NMS to RPN proposals __C.TEST.RPN_PRE_NMS_TOP_N = 6000 ## Number of top scoring boxes to keep after applying NMS to RPN proposals __C.TEST.RPN_POST_NMS_TOP_N = 300 # Proposal height and width both need to be greater than RPN_MIN_SIZE (at orig image scale) __C.TEST.RPN_MIN_SIZE = 16 # Testing mode, default to be 'nms', 'top' is slower but better # See report for details __C.TEST.MODE = 'nms' # Only useful when TEST.MODE is 'top', specifies the number of top proposals to select __C.TEST.RPN_TOP_N = 5000 # # ResNet options # __C.RESNET = edict() # Option to set if max-pooling is appended after crop_and_resize. # if true, the region will be resized to a square of 2xPOOLING_SIZE, # then 2x2 max-pooling is applied; otherwise the region will be directly # resized to a square of POOLING_SIZE __C.RESNET.MAX_POOL = False # Number of fixed blocks during training, by default the first of all 4 blocks is fixed # Range: 0 (none) to 3 (all) __C.RESNET.FIXED_BLOCKS = 1 # # MobileNet options # __C.MOBILENET = edict() # Whether to regularize the depth-wise filters during training __C.MOBILENET.REGU_DEPTH = False # Number of fixed layers during training, by default the first of all 14 layers is fixed # Range: 0 (none) to 12 (all) __C.MOBILENET.FIXED_LAYERS = 5 # Weight decay for the mobilenet weights __C.MOBILENET.WEIGHT_DECAY = 0.00004 # Depth multiplier __C.MOBILENET.DEPTH_MULTIPLIER = 1. # # MISC # # The mapping from image coordinates to feature map coordinates might cause # some boxes that are distinct in image space to become identical in feature # coordinates. If DEDUP_BOXES > 0, then DEDUP_BOXES is used as the scale factor # for identifying duplicate boxes. # 1/16 is correct for {Alex,Caffe}Net, VGG_CNN_M_1024, and VGG16 __C.DEDUP_BOXES = 1. / 16. # Pixel mean values (BGR order) as a (1, 1, 3) array # We use the same pixel mean for all networks even though it's not exactly what # they were trained with __C.PIXEL_MEANS = np.array([[[102.9801, 115.9465, 122.7717]]]) # For reproducibility __C.RNG_SEED = 3 # A small number that's used many times __C.EPS = 1e-14 # Root directory of project __C.ROOT_DIR = osp.abspath(osp.join(osp.dirname(__file__), '..', '..', '..')) # Data directory __C.DATA_DIR = osp.abspath(osp.join(__C.ROOT_DIR, 'data')) # Name (or path to) the matlab executable __C.MATLAB = 'matlab' # Place outputs under an experiments directory __C.EXP_DIR = 'default' # Use GPU implementation of non-maximum suppression __C.USE_GPU_NMS = True # Default GPU device id __C.GPU_ID = 0 __C.POOLING_MODE = 'crop' # Size of the pooled region after RoI pooling __C.POOLING_SIZE = 7 # Maximal number of gt rois in an image during Training __C.MAX_NUM_GT_BOXES = 20 # Anchor scales for RPN __C.ANCHOR_SCALES = [8,16,32] # Anchor ratios for RPN __C.ANCHOR_RATIOS = [0.5,1,2] # Feature stride for RPN __C.FEAT_STRIDE = [16, ] __C.CUDA = False __C.CROP_RESIZE_WITH_MAX_POOL = True import pdb def get_output_dir(imdb, weights_filename): """Return the directory where experimental artifacts are placed. If the directory does not exist, it is created. A canonical path is built using the name from an imdb and a network (if not None). """ outdir = osp.abspath(osp.join(__C.ROOT_DIR, 'output', __C.EXP_DIR, imdb.name)) if weights_filename is None: weights_filename = 'default' outdir = osp.join(outdir, weights_filename) if not os.path.exists(outdir): os.makedirs(outdir) return outdir def get_output_tb_dir(imdb, weights_filename): """Return the directory where tensorflow summaries are placed. If the directory does not exist, it is created. A canonical path is built using the name from an imdb and a network (if not None). """ outdir = osp.abspath(osp.join(__C.ROOT_DIR, 'tensorboard', __C.EXP_DIR, imdb.name)) if weights_filename is None: weights_filename = 'default' outdir = osp.join(outdir, weights_filename) if not os.path.exists(outdir): os.makedirs(outdir) return outdir def _merge_a_into_b(a, b): """Merge config dictionary a into config dictionary b, clobbering the options in b whenever they are also specified in a. """ if type(a) is not edict: return for k, v in a.items(): # a must specify keys that are in b if k not in b: raise KeyError('{} is not a valid config key'.format(k)) # the types must match, too old_type = type(b[k]) if old_type is not type(v): if isinstance(b[k], np.ndarray): v = np.array(v, dtype=b[k].dtype) else: raise ValueError(('Type mismatch ({} vs. {}) ' 'for config key: {}').format(type(b[k]), type(v), k)) # recursively merge dicts if type(v) is edict: try: _merge_a_into_b(a[k], b[k]) except: print(('Error under config key: {}'.format(k))) raise else: b[k] = v def cfg_from_file(filename): """Load a config file and merge it into the default options.""" import yaml with open(filename, 'r') as f: yaml_cfg = edict(yaml.load(f)) _merge_a_into_b(yaml_cfg, __C) def cfg_from_list(cfg_list): """Set config keys via list (e.g., from command line).""" from ast import literal_eval assert len(cfg_list) % 2 == 0 for k, v in zip(cfg_list[0::2], cfg_list[1::2]): key_list = k.split('.') d = __C for subkey in key_list[:-1]: assert subkey in d d = d[subkey] subkey = key_list[-1] assert subkey in d try: value = literal_eval(v) except: # handle the case when v is a string literal value = v assert type(value) == type(d[subkey]), \ 'type {} does not match original type {}'.format( type(value), type(d[subkey])) d[subkey] = value ================================================ FILE: lib/model/utils/logger.py ================================================ # Code referenced from https://gist.github.com/gyglim/1f8dfb1b5c82627ae3efcfbbadb9f514 import tensorflow as tf import numpy as np import scipy.misc try: from StringIO import StringIO # Python 2.7 except ImportError: from io import BytesIO # Python 3.x class Logger(object): def __init__(self, log_dir): """Create a summary writer logging to log_dir.""" self.writer = tf.summary.FileWriter(log_dir) def scalar_summary(self, tag, value, step): """Log a scalar variable.""" summary = tf.Summary(value=[tf.Summary.Value(tag=tag, simple_value=value)]) self.writer.add_summary(summary, step) def image_summary(self, tag, images, step): """Log a list of images.""" img_summaries = [] for i, img in enumerate(images): # Write the image to a string try: s = StringIO() except: s = BytesIO() scipy.misc.toimage(img).save(s, format="png") # Create an Image object img_sum = tf.Summary.Image(encoded_image_string=s.getvalue(), height=img.shape[0], width=img.shape[1]) # Create a Summary value img_summaries.append(tf.Summary.Value(tag='%s/%d' % (tag, i), image=img_sum)) # Create and write Summary summary = tf.Summary(value=img_summaries) self.writer.add_summary(summary, step) def histo_summary(self, tag, values, step, bins=1000): """Log a histogram of the tensor of values.""" # Create a histogram using numpy counts, bin_edges = np.histogram(values, bins=bins) # Fill the fields of the histogram proto hist = tf.HistogramProto() hist.min = float(np.min(values)) hist.max = float(np.max(values)) hist.num = int(np.prod(values.shape)) hist.sum = float(np.sum(values)) hist.sum_squares = float(np.sum(values**2)) # Drop the start of the first bin bin_edges = bin_edges[1:] # Add bin edges and counts for edge in bin_edges: hist.bucket_limit.append(edge) for c in counts: hist.bucket.append(c) # Create and write Summary summary = tf.Summary(value=[tf.Summary.Value(tag=tag, histo=hist)]) self.writer.add_summary(summary, step) self.writer.flush() ================================================ FILE: lib/model/utils/net_utils.py ================================================ import torch import torch.nn as nn import torch.nn.functional as F from torch.autograd import Variable import numpy as np import torchvision.models as models from model.utils.config import cfg from model.roi_crop.functions.roi_crop import RoICropFunction import cv2 import pdb import random def save_net(fname, net): import h5py h5f = h5py.File(fname, mode='w') for k, v in net.state_dict().items(): h5f.create_dataset(k, data=v.cpu().numpy()) def load_net(fname, net): import h5py h5f = h5py.File(fname, mode='r') for k, v in net.state_dict().items(): param = torch.from_numpy(np.asarray(h5f[k])) v.copy_(param) def weights_normal_init(model, dev=0.01): if isinstance(model, list): for m in model: weights_normal_init(m, dev) else: for m in model.modules(): if isinstance(m, nn.Conv2d): m.weight.data.normal_(0.0, dev) elif isinstance(m, nn.Linear): m.weight.data.normal_(0.0, dev) def clip_gradient(model, clip_norm): """Computes a gradient clipping coefficient based on gradient norm.""" totalnorm = 0 for p in model.parameters(): if p.requires_grad: modulenorm = p.grad.data.norm() totalnorm += modulenorm ** 2 totalnorm = torch.sqrt(totalnorm).item() norm = (clip_norm / max(totalnorm, clip_norm)) for p in model.parameters(): if p.requires_grad: p.grad.mul_(norm) def vis_detections(im, class_name, dets, thresh=0.8): """Visual debugging of detections.""" for i in range(np.minimum(10, dets.shape[0])): bbox = tuple(int(np.round(x)) for x in dets[i, :4]) score = dets[i, -1] if score > thresh: cv2.rectangle(im, bbox[0:2], bbox[2:4], (0, 204, 0), 2) cv2.putText(im, '%s: %.3f' % (class_name, score), (bbox[0], bbox[1] + 15), cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 0, 255), thickness=1) return im def adjust_learning_rate(optimizer, decay=0.1): """Sets the learning rate to the initial LR decayed by 0.5 every 20 epochs""" for param_group in optimizer.param_groups: param_group['lr'] = decay * param_group['lr'] def save_checkpoint(state, filename): torch.save(state, filename) def _smooth_l1_loss(bbox_pred, bbox_targets, bbox_inside_weights, bbox_outside_weights, sigma=1.0, dim=[1]): sigma_2 = sigma ** 2 box_diff = bbox_pred - bbox_targets in_box_diff = bbox_inside_weights * box_diff abs_in_box_diff = torch.abs(in_box_diff) smoothL1_sign = (abs_in_box_diff < 1. / sigma_2).detach().float() in_loss_box = torch.pow(in_box_diff, 2) * (sigma_2 / 2.) * smoothL1_sign \ + (abs_in_box_diff - (0.5 / sigma_2)) * (1. - smoothL1_sign) out_loss_box = bbox_outside_weights * in_loss_box loss_box = out_loss_box for i in sorted(dim, reverse=True): loss_box = loss_box.sum(i) loss_box = loss_box.mean() return loss_box def _crop_pool_layer(bottom, rois, max_pool=True): # code modified from # https://github.com/ruotianluo/pytorch-faster-rcnn # implement it using stn # box to affine # input (x1,y1,x2,y2) """ [ x2-x1 x1 + x2 - W + 1 ] [ ----- 0 --------------- ] [ W - 1 W - 1 ] [ ] [ y2-y1 y1 + y2 - H + 1 ] [ 0 ----- --------------- ] [ H - 1 H - 1 ] """ rois = rois.detach() batch_size = bottom.size(0) D = bottom.size(1) H = bottom.size(2) W = bottom.size(3) roi_per_batch = rois.size(0) / batch_size x1 = rois[:, 1::4] / 16.0 y1 = rois[:, 2::4] / 16.0 x2 = rois[:, 3::4] / 16.0 y2 = rois[:, 4::4] / 16.0 height = bottom.size(2) width = bottom.size(3) # affine theta zero = Variable(rois.data.new(rois.size(0), 1).zero_()) theta = torch.cat([\ (x2 - x1) / (width - 1), zero, (x1 + x2 - width + 1) / (width - 1), zero, (y2 - y1) / (height - 1), (y1 + y2 - height + 1) / (height - 1)], 1).view(-1, 2, 3) if max_pool: pre_pool_size = cfg.POOLING_SIZE * 2 grid = F.affine_grid(theta, torch.Size((rois.size(0), 1, pre_pool_size, pre_pool_size))) bottom = bottom.view(1, batch_size, D, H, W).contiguous().expand(roi_per_batch, batch_size, D, H, W)\ .contiguous().view(-1, D, H, W) crops = F.grid_sample(bottom, grid) crops = F.max_pool2d(crops, 2, 2) else: grid = F.affine_grid(theta, torch.Size((rois.size(0), 1, cfg.POOLING_SIZE, cfg.POOLING_SIZE))) bottom = bottom.view(1, batch_size, D, H, W).contiguous().expand(roi_per_batch, batch_size, D, H, W)\ .contiguous().view(-1, D, H, W) crops = F.grid_sample(bottom, grid) return crops, grid def _affine_grid_gen(rois, input_size, grid_size): rois = rois.detach() x1 = rois[:, 1::4] / 16.0 y1 = rois[:, 2::4] / 16.0 x2 = rois[:, 3::4] / 16.0 y2 = rois[:, 4::4] / 16.0 height = input_size[0] width = input_size[1] zero = Variable(rois.data.new(rois.size(0), 1).zero_()) theta = torch.cat([\ (x2 - x1) / (width - 1), zero, (x1 + x2 - width + 1) / (width - 1), zero, (y2 - y1) / (height - 1), (y1 + y2 - height + 1) / (height - 1)], 1).view(-1, 2, 3) grid = F.affine_grid(theta, torch.Size((rois.size(0), 1, grid_size, grid_size))) return grid def _affine_theta(rois, input_size): rois = rois.detach() x1 = rois[:, 1::4] / 16.0 y1 = rois[:, 2::4] / 16.0 x2 = rois[:, 3::4] / 16.0 y2 = rois[:, 4::4] / 16.0 height = input_size[0] width = input_size[1] zero = Variable(rois.data.new(rois.size(0), 1).zero_()) # theta = torch.cat([\ # (x2 - x1) / (width - 1), # zero, # (x1 + x2 - width + 1) / (width - 1), # zero, # (y2 - y1) / (height - 1), # (y1 + y2 - height + 1) / (height - 1)], 1).view(-1, 2, 3) theta = torch.cat([\ (y2 - y1) / (height - 1), zero, (y1 + y2 - height + 1) / (height - 1), zero, (x2 - x1) / (width - 1), (x1 + x2 - width + 1) / (width - 1)], 1).view(-1, 2, 3) return theta def compare_grid_sample(): # do gradcheck N = random.randint(1, 8) C = 2 # random.randint(1, 8) H = 5 # random.randint(1, 8) W = 4 # random.randint(1, 8) input = Variable(torch.randn(N, C, H, W).cuda(), requires_grad=True) input_p = input.clone().data.contiguous() grid = Variable(torch.randn(N, H, W, 2).cuda(), requires_grad=True) grid_clone = grid.clone().contiguous() out_offcial = F.grid_sample(input, grid) grad_outputs = Variable(torch.rand(out_offcial.size()).cuda()) grad_outputs_clone = grad_outputs.clone().contiguous() grad_inputs = torch.autograd.grad(out_offcial, (input, grid), grad_outputs.contiguous()) grad_input_off = grad_inputs[0] crf = RoICropFunction() grid_yx = torch.stack([grid_clone.data[:,:,:,1], grid_clone.data[:,:,:,0]], 3).contiguous().cuda() out_stn = crf.forward(input_p, grid_yx) grad_inputs = crf.backward(grad_outputs_clone.data) grad_input_stn = grad_inputs[0] pdb.set_trace() delta = (grad_input_off.data - grad_input_stn).sum() ================================================ FILE: lib/pycocotools/UPSTREAM_REV ================================================ https://github.com/pdollar/coco/commit/3ac47c77ebd5a1ed4254a98b7fbf2ef4765a3574 ================================================ FILE: lib/pycocotools/__init__.py ================================================ __author__ = 'tylin' ================================================ FILE: lib/pycocotools/_mask.pyx ================================================ # distutils: language = c # distutils: sources = ../MatlabAPI/private/maskApi.c #************************************************************************** # Microsoft COCO Toolbox. version 2.0 # Data, paper, and tutorials available at: http://mscoco.org/ # Code written by Piotr Dollar and Tsung-Yi Lin, 2015. # Licensed under the Simplified BSD License [see coco/license.txt] #************************************************************************** __author__ = 'tsungyi' # import both Python-level and C-level symbols of Numpy # the API uses Numpy to interface C and Python import numpy as np cimport numpy as np from libc.stdlib cimport malloc, free # intialized Numpy. must do. np.import_array() # import numpy C function # we use PyArray_ENABLEFLAGS to make Numpy ndarray responsible to memoery management cdef extern from "numpy/arrayobject.h": void PyArray_ENABLEFLAGS(np.ndarray arr, int flags) # Declare the prototype of the C functions in MaskApi.h cdef extern from "maskApi.h": ctypedef unsigned int uint ctypedef unsigned long siz ctypedef unsigned char byte ctypedef double* BB ctypedef struct RLE: siz h, siz w, siz m, uint* cnts, void rlesInit( RLE **R, siz n ) void rleEncode( RLE *R, const byte *M, siz h, siz w, siz n ) void rleDecode( const RLE *R, byte *mask, siz n ) void rleMerge( const RLE *R, RLE *M, siz n, bint intersect ) void rleArea( const RLE *R, siz n, uint *a ) void rleIou( RLE *dt, RLE *gt, siz m, siz n, byte *iscrowd, double *o ) void bbIou( BB dt, BB gt, siz m, siz n, byte *iscrowd, double *o ) void rleToBbox( const RLE *R, BB bb, siz n ) void rleFrBbox( RLE *R, const BB bb, siz h, siz w, siz n ) void rleFrPoly( RLE *R, const double *xy, siz k, siz h, siz w ) char* rleToString( const RLE *R ) void rleFrString( RLE *R, char *s, siz h, siz w ) # python class to wrap RLE array in C # the class handles the memory allocation and deallocation cdef class RLEs: cdef RLE *_R cdef siz _n def __cinit__(self, siz n =0): rlesInit(&self._R, n) self._n = n # free the RLE array here def __dealloc__(self): if self._R is not NULL: for i in range(self._n): free(self._R[i].cnts) free(self._R) def __getattr__(self, key): if key == 'n': return self._n raise AttributeError(key) # python class to wrap Mask array in C # the class handles the memory allocation and deallocation cdef class Masks: cdef byte *_mask cdef siz _h cdef siz _w cdef siz _n def __cinit__(self, h, w, n): self._mask = malloc(h*w*n* sizeof(byte)) self._h = h self._w = w self._n = n # def __dealloc__(self): # the memory management of _mask has been passed to np.ndarray # it doesn't need to be freed here # called when passing into np.array() and return an np.ndarray in column-major order def __array__(self): cdef np.npy_intp shape[1] shape[0] = self._h*self._w*self._n # Create a 1D array, and reshape it to fortran/Matlab column-major array ndarray = np.PyArray_SimpleNewFromData(1, shape, np.NPY_UINT8, self._mask).reshape((self._h, self._w, self._n), order='F') # The _mask allocated by Masks is now handled by ndarray PyArray_ENABLEFLAGS(ndarray, np.NPY_OWNDATA) return ndarray # internal conversion from Python RLEs object to compressed RLE format def _toString(RLEs Rs): cdef siz n = Rs.n cdef bytes py_string cdef char* c_string objs = [] for i in range(n): c_string = rleToString( &Rs._R[i] ) py_string = c_string objs.append({ 'size': [Rs._R[i].h, Rs._R[i].w], 'counts': py_string }) free(c_string) return objs # internal conversion from compressed RLE format to Python RLEs object def _frString(rleObjs): cdef siz n = len(rleObjs) Rs = RLEs(n) cdef bytes py_string cdef char* c_string for i, obj in enumerate(rleObjs): py_string = str(obj['counts']) c_string = py_string rleFrString( &Rs._R[i], c_string, obj['size'][0], obj['size'][1] ) return Rs # encode mask to RLEs objects # list of RLE string can be generated by RLEs member function def encode(np.ndarray[np.uint8_t, ndim=3, mode='fortran'] mask): h, w, n = mask.shape[0], mask.shape[1], mask.shape[2] cdef RLEs Rs = RLEs(n) rleEncode(Rs._R,mask.data,h,w,n) objs = _toString(Rs) return objs # decode mask from compressed list of RLE string or RLEs object def decode(rleObjs): cdef RLEs Rs = _frString(rleObjs) h, w, n = Rs._R[0].h, Rs._R[0].w, Rs._n masks = Masks(h, w, n) rleDecode( Rs._R, masks._mask, n ); return np.array(masks) def merge(rleObjs, bint intersect=0): cdef RLEs Rs = _frString(rleObjs) cdef RLEs R = RLEs(1) rleMerge(Rs._R, R._R, Rs._n, intersect) obj = _toString(R)[0] return obj def area(rleObjs): cdef RLEs Rs = _frString(rleObjs) cdef uint* _a = malloc(Rs._n* sizeof(uint)) rleArea(Rs._R, Rs._n, _a) cdef np.npy_intp shape[1] shape[0] = Rs._n a = np.array((Rs._n, ), dtype=np.uint8) a = np.PyArray_SimpleNewFromData(1, shape, np.NPY_UINT32, _a) PyArray_ENABLEFLAGS(a, np.NPY_OWNDATA) return a # iou computation. support function overload (RLEs-RLEs and bbox-bbox). def iou( dt, gt, pyiscrowd ): def _preproc(objs): if len(objs) == 0: return objs if type(objs) == np.ndarray: if len(objs.shape) == 1: objs = objs.reshape((objs[0], 1)) # check if it's Nx4 bbox if not len(objs.shape) == 2 or not objs.shape[1] == 4: raise Exception('numpy ndarray input is only for *bounding boxes* and should have Nx4 dimension') objs = objs.astype(np.double) elif type(objs) == list: # check if list is in box format and convert it to np.ndarray isbox = np.all(np.array([(len(obj)==4) and ((type(obj)==list) or (type(obj)==np.ndarray)) for obj in objs])) isrle = np.all(np.array([type(obj) == dict for obj in objs])) if isbox: objs = np.array(objs, dtype=np.double) if len(objs.shape) == 1: objs = objs.reshape((1,objs.shape[0])) elif isrle: objs = _frString(objs) else: raise Exception('list input can be bounding box (Nx4) or RLEs ([RLE])') else: raise Exception('unrecognized type. The following type: RLEs (rle), np.ndarray (box), and list (box) are supported.') return objs def _rleIou(RLEs dt, RLEs gt, np.ndarray[np.uint8_t, ndim=1] iscrowd, siz m, siz n, np.ndarray[np.double_t, ndim=1] _iou): rleIou( dt._R, gt._R, m, n, iscrowd.data, _iou.data ) def _bbIou(np.ndarray[np.double_t, ndim=2] dt, np.ndarray[np.double_t, ndim=2] gt, np.ndarray[np.uint8_t, ndim=1] iscrowd, siz m, siz n, np.ndarray[np.double_t, ndim=1] _iou): bbIou( dt.data, gt.data, m, n, iscrowd.data, _iou.data ) def _len(obj): cdef siz N = 0 if type(obj) == RLEs: N = obj.n elif len(obj)==0: pass elif type(obj) == np.ndarray: N = obj.shape[0] return N # convert iscrowd to numpy array cdef np.ndarray[np.uint8_t, ndim=1] iscrowd = np.array(pyiscrowd, dtype=np.uint8) # simple type checking cdef siz m, n dt = _preproc(dt) gt = _preproc(gt) m = _len(dt) n = _len(gt) if m == 0 or n == 0: return [] if not type(dt) == type(gt): raise Exception('The dt and gt should have the same data type, either RLEs, list or np.ndarray') # define local variables cdef double* _iou = 0 cdef np.npy_intp shape[1] # check type and assign iou function if type(dt) == RLEs: _iouFun = _rleIou elif type(dt) == np.ndarray: _iouFun = _bbIou else: raise Exception('input data type not allowed.') _iou = malloc(m*n* sizeof(double)) iou = np.zeros((m*n, ), dtype=np.double) shape[0] = m*n iou = np.PyArray_SimpleNewFromData(1, shape, np.NPY_DOUBLE, _iou) PyArray_ENABLEFLAGS(iou, np.NPY_OWNDATA) _iouFun(dt, gt, iscrowd, m, n, iou) return iou.reshape((m,n), order='F') def toBbox( rleObjs ): cdef RLEs Rs = _frString(rleObjs) cdef siz n = Rs.n cdef BB _bb = malloc(4*n* sizeof(double)) rleToBbox( Rs._R, _bb, n ) cdef np.npy_intp shape[1] shape[0] = 4*n bb = np.array((1,4*n), dtype=np.double) bb = np.PyArray_SimpleNewFromData(1, shape, np.NPY_DOUBLE, _bb).reshape((n, 4)) PyArray_ENABLEFLAGS(bb, np.NPY_OWNDATA) return bb def frBbox(np.ndarray[np.double_t, ndim=2] bb, siz h, siz w ): cdef siz n = bb.shape[0] Rs = RLEs(n) rleFrBbox( Rs._R, bb.data, h, w, n ) objs = _toString(Rs) return objs def frPoly( poly, siz h, siz w ): cdef np.ndarray[np.double_t, ndim=1] np_poly n = len(poly) Rs = RLEs(n) for i, p in enumerate(poly): np_poly = np.array(p, dtype=np.double, order='F') rleFrPoly( &Rs._R[i], np_poly.data, len(np_poly)/2, h, w ) objs = _toString(Rs) return objs def frUncompressedRLE(ucRles, siz h, siz w): cdef np.ndarray[np.uint32_t, ndim=1] cnts cdef RLE R cdef uint *data n = len(ucRles) objs = [] for i in range(n): Rs = RLEs(1) cnts = np.array(ucRles[i]['counts'], dtype=np.uint32) # time for malloc can be saved here but it's fine data = malloc(len(cnts)* sizeof(uint)) for j in range(len(cnts)): data[j] = cnts[j] R = RLE(ucRles[i]['size'][0], ucRles[i]['size'][1], len(cnts), data) Rs._R[0] = R objs.append(_toString(Rs)[0]) return objs def frPyObjects(pyobj, siz h, w): if type(pyobj) == np.ndarray: objs = frBbox(pyobj, h, w ) elif type(pyobj) == list and len(pyobj[0]) == 4: objs = frBbox(pyobj, h, w ) elif type(pyobj) == list and len(pyobj[0]) > 4: objs = frPoly(pyobj, h, w ) elif type(pyobj) == list and type(pyobj[0]) == dict: objs = frUncompressedRLE(pyobj, h, w) else: raise Exception('input type is not supported.') return objs ================================================ FILE: lib/pycocotools/coco.py ================================================ from __future__ import print_function from __future__ import absolute_import __author__ = 'tylin' __version__ = '1.0.1' # Interface for accessing the Microsoft COCO dataset. # Microsoft COCO is a large image dataset designed for object detection, # segmentation, and caption generation. pycocotools is a Python API that # assists in loading, parsing and visualizing the annotations in COCO. # Please visit http://mscoco.org/ for more information on COCO, including # for the data, paper, and tutorials. The exact format of the annotations # is also described on the COCO website. For example usage of the pycocotools # please see pycocotools_demo.ipynb. In addition to this API, please download both # the COCO images and annotations in order to run the demo. # An alternative to using the API is to load the annotations directly # into Python dictionary # Using the API provides additional utility functions. Note that this API # supports both *instance* and *caption* annotations. In the case of # captions not all functions are defined (e.g. categories are undefined). # The following API functions are defined: # COCO - COCO api class that loads COCO annotation file and prepare data structures. # decodeMask - Decode binary mask M encoded via run-length encoding. # encodeMask - Encode binary mask M using run-length encoding. # getAnnIds - Get ann ids that satisfy given filter conditions. # getCatIds - Get cat ids that satisfy given filter conditions. # getImgIds - Get img ids that satisfy given filter conditions. # loadAnns - Load anns with the specified ids. # loadCats - Load cats with the specified ids. # loadImgs - Load imgs with the specified ids. # segToMask - Convert polygon segmentation to binary mask. # showAnns - Display the specified annotations. # loadRes - Load algorithm results and create API for accessing them. # download - Download COCO images from mscoco.org server. # Throughout the API "ann"=annotation, "cat"=category, and "img"=image. # Help on each functions can be accessed by: "help COCO>function". # See also COCO>decodeMask, # COCO>encodeMask, COCO>getAnnIds, COCO>getCatIds, # COCO>getImgIds, COCO>loadAnns, COCO>loadCats, # COCO>loadImgs, COCO>segToMask, COCO>showAnns # Microsoft COCO Toolbox. version 2.0 # Data, paper, and tutorials available at: http://mscoco.org/ # Code written by Piotr Dollar and Tsung-Yi Lin, 2014. # Licensed under the Simplified BSD License [see bsd.txt] import json import datetime import time import matplotlib.pyplot as plt from matplotlib.collections import PatchCollection from matplotlib.patches import Polygon import numpy as np # from skimage.draw import polygon import urllib import copy import itertools from . import mask import os try: unicode # Python 2 except NameError: unicode = str # Python 3 class COCO: def __init__(self, annotation_file=None): """ Constructor of Microsoft COCO helper class for reading and visualizing annotations. :param annotation_file (str): location of annotation file :param image_folder (str): location to the folder that hosts images. :return: """ # load dataset self.dataset = {} self.anns = [] self.imgToAnns = {} self.catToImgs = {} self.imgs = {} self.cats = {} if not annotation_file == None: print('loading annotations into memory...') tic = time.time() dataset = json.load(open(annotation_file, 'r')) print('Done (t=%0.2fs)'%(time.time()- tic)) self.dataset = dataset self.createIndex() def createIndex(self): # create index print('creating index...') anns = {} imgToAnns = {} catToImgs = {} cats = {} imgs = {} if 'annotations' in self.dataset: imgToAnns = {ann['image_id']: [] for ann in self.dataset['annotations']} anns = {ann['id']: [] for ann in self.dataset['annotations']} for ann in self.dataset['annotations']: imgToAnns[ann['image_id']] += [ann] anns[ann['id']] = ann if 'images' in self.dataset: imgs = {im['id']: {} for im in self.dataset['images']} for img in self.dataset['images']: imgs[img['id']] = img if 'categories' in self.dataset: cats = {cat['id']: [] for cat in self.dataset['categories']} for cat in self.dataset['categories']: cats[cat['id']] = cat catToImgs = {cat['id']: [] for cat in self.dataset['categories']} if 'annotations' in self.dataset: for ann in self.dataset['annotations']: catToImgs[ann['category_id']] += [ann['image_id']] print('index created!') # create class members self.anns = anns self.imgToAnns = imgToAnns self.catToImgs = catToImgs self.imgs = imgs self.cats = cats def info(self): """ Print information about the annotation file. :return: """ for key, value in self.dataset['info'].items(): print('%s: %s'%(key, value)) def getAnnIds(self, imgIds=[], catIds=[], areaRng=[], iscrowd=None): """ Get ann ids that satisfy given filter conditions. default skips that filter :param imgIds (int array) : get anns for given imgs catIds (int array) : get anns for given cats areaRng (float array) : get anns for given area range (e.g. [0 inf]) iscrowd (boolean) : get anns for given crowd label (False or True) :return: ids (int array) : integer array of ann ids """ imgIds = imgIds if type(imgIds) == list else [imgIds] catIds = catIds if type(catIds) == list else [catIds] if len(imgIds) == len(catIds) == len(areaRng) == 0: anns = self.dataset['annotations'] else: if not len(imgIds) == 0: # this can be changed by defaultdict lists = [self.imgToAnns[imgId] for imgId in imgIds if imgId in self.imgToAnns] anns = list(itertools.chain.from_iterable(lists)) else: anns = self.dataset['annotations'] anns = anns if len(catIds) == 0 else [ann for ann in anns if ann['category_id'] in catIds] anns = anns if len(areaRng) == 0 else [ann for ann in anns if ann['area'] > areaRng[0] and ann['area'] < areaRng[1]] if not iscrowd == None: ids = [ann['id'] for ann in anns if ann['iscrowd'] == iscrowd] else: ids = [ann['id'] for ann in anns] return ids def getCatIds(self, catNms=[], supNms=[], catIds=[]): """ filtering parameters. default skips that filter. :param catNms (str array) : get cats for given cat names :param supNms (str array) : get cats for given supercategory names :param catIds (int array) : get cats for given cat ids :return: ids (int array) : integer array of cat ids """ catNms = catNms if type(catNms) == list else [catNms] supNms = supNms if type(supNms) == list else [supNms] catIds = catIds if type(catIds) == list else [catIds] if len(catNms) == len(supNms) == len(catIds) == 0: cats = self.dataset['categories'] else: cats = self.dataset['categories'] cats = cats if len(catNms) == 0 else [cat for cat in cats if cat['name'] in catNms] cats = cats if len(supNms) == 0 else [cat for cat in cats if cat['supercategory'] in supNms] cats = cats if len(catIds) == 0 else [cat for cat in cats if cat['id'] in catIds] ids = [cat['id'] for cat in cats] return ids def getImgIds(self, imgIds=[], catIds=[]): ''' Get img ids that satisfy given filter conditions. :param imgIds (int array) : get imgs for given ids :param catIds (int array) : get imgs with all given cats :return: ids (int array) : integer array of img ids ''' imgIds = imgIds if type(imgIds) == list else [imgIds] catIds = catIds if type(catIds) == list else [catIds] if len(imgIds) == len(catIds) == 0: ids = self.imgs.keys() else: ids = set(imgIds) for i, catId in enumerate(catIds): if i == 0 and len(ids) == 0: ids = set(self.catToImgs[catId]) else: ids &= set(self.catToImgs[catId]) return list(ids) def loadAnns(self, ids=[]): """ Load anns with the specified ids. :param ids (int array) : integer ids specifying anns :return: anns (object array) : loaded ann objects """ if type(ids) == list: return [self.anns[id] for id in ids] elif type(ids) == int: return [self.anns[ids]] def loadCats(self, ids=[]): """ Load cats with the specified ids. :param ids (int array) : integer ids specifying cats :return: cats (object array) : loaded cat objects """ if type(ids) == list: return [self.cats[id] for id in ids] elif type(ids) == int: return [self.cats[ids]] def loadImgs(self, ids=[]): """ Load anns with the specified ids. :param ids (int array) : integer ids specifying img :return: imgs (object array) : loaded img objects """ if type(ids) == list: return [self.imgs[id] for id in ids] elif type(ids) == int: return [self.imgs[ids]] def showAnns(self, anns): """ Display the specified annotations. :param anns (array of object): annotations to display :return: None """ if len(anns) == 0: return 0 if 'segmentation' in anns[0]: datasetType = 'instances' elif 'caption' in anns[0]: datasetType = 'captions' if datasetType == 'instances': ax = plt.gca() polygons = [] color = [] for ann in anns: c = np.random.random((1, 3)).tolist()[0] if type(ann['segmentation']) == list: # polygon for seg in ann['segmentation']: poly = np.array(seg).reshape((len(seg)/2, 2)) polygons.append(Polygon(poly, True,alpha=0.4)) color.append(c) else: # mask t = self.imgs[ann['image_id']] if type(ann['segmentation']['counts']) == list: rle = mask.frPyObjects([ann['segmentation']], t['height'], t['width']) else: rle = [ann['segmentation']] m = mask.decode(rle) img = np.ones( (m.shape[0], m.shape[1], 3) ) if ann['iscrowd'] == 1: color_mask = np.array([2.0,166.0,101.0])/255 if ann['iscrowd'] == 0: color_mask = np.random.random((1, 3)).tolist()[0] for i in range(3): img[:,:,i] = color_mask[i] ax.imshow(np.dstack( (img, m*0.5) )) p = PatchCollection(polygons, facecolors=color, edgecolors=(0,0,0,1), linewidths=3, alpha=0.4) ax.add_collection(p) elif datasetType == 'captions': for ann in anns: print(ann['caption']) def loadRes(self, resFile): """ Load result file and return a result api object. :param resFile (str) : file name of result file :return: res (obj) : result api object """ res = COCO() res.dataset['images'] = [img for img in self.dataset['images']] # res.dataset['info'] = copy.deepcopy(self.dataset['info']) # res.dataset['licenses'] = copy.deepcopy(self.dataset['licenses']) print('Loading and preparing results... ') tic = time.time() anns = json.load(open(resFile)) assert type(anns) == list, 'results in not an array of objects' annsImgIds = [ann['image_id'] for ann in anns] assert set(annsImgIds) == (set(annsImgIds) & set(self.getImgIds())), \ 'Results do not correspond to current coco set' if 'caption' in anns[0]: imgIds = set([img['id'] for img in res.dataset['images']]) & set([ann['image_id'] for ann in anns]) res.dataset['images'] = [img for img in res.dataset['images'] if img['id'] in imgIds] for id, ann in enumerate(anns): ann['id'] = id+1 elif 'bbox' in anns[0] and not anns[0]['bbox'] == []: res.dataset['categories'] = copy.deepcopy(self.dataset['categories']) for id, ann in enumerate(anns): bb = ann['bbox'] x1, x2, y1, y2 = [bb[0], bb[0]+bb[2], bb[1], bb[1]+bb[3]] if not 'segmentation' in ann: ann['segmentation'] = [[x1, y1, x1, y2, x2, y2, x2, y1]] ann['area'] = bb[2]*bb[3] ann['id'] = id+1 ann['iscrowd'] = 0 elif 'segmentation' in anns[0]: res.dataset['categories'] = copy.deepcopy(self.dataset['categories']) for id, ann in enumerate(anns): # now only support compressed RLE format as segmentation results ann['area'] = mask.area([ann['segmentation']])[0] if not 'bbox' in ann: ann['bbox'] = mask.toBbox([ann['segmentation']])[0] ann['id'] = id+1 ann['iscrowd'] = 0 print('DONE (t=%0.2fs)'%(time.time()- tic)) res.dataset['annotations'] = anns res.createIndex() return res def download( self, tarDir = None, imgIds = [] ): ''' Download COCO images from mscoco.org server. :param tarDir (str): COCO results directory name imgIds (list): images to be downloaded :return: ''' if tarDir is None: print('Please specify target directory') return -1 if len(imgIds) == 0: imgs = self.imgs.values() else: imgs = self.loadImgs(imgIds) N = len(imgs) if not os.path.exists(tarDir): os.makedirs(tarDir) for i, img in enumerate(imgs): tic = time.time() fname = os.path.join(tarDir, img['file_name']) if not os.path.exists(fname): urllib.urlretrieve(img['coco_url'], fname) print('downloaded %d/%d images (t=%.1fs)'%(i, N, time.time()- tic)) ================================================ FILE: lib/pycocotools/cocoeval.py ================================================ from __future__ import print_function from __future__ import absolute_import __author__ = 'tsungyi' import numpy as np import datetime import time from collections import defaultdict from . import mask import copy try: string_types = (str, unicode) # Python 2 except NameError: string_types = (str, ) # Python 3 class COCOeval: # Interface for evaluating detection on the Microsoft COCO dataset. # # The usage for CocoEval is as follows: # cocoGt=..., cocoDt=... # load dataset and results # E = CocoEval(cocoGt,cocoDt); # initialize CocoEval object # E.params.recThrs = ...; # set parameters as desired # E.evaluate(); # run per image evaluation # E.accumulate(); # accumulate per image results # E.summarize(); # display summary metrics of results # For example usage see evalDemo.m and http://mscoco.org/. # # The evaluation parameters are as follows (defaults in brackets): # imgIds - [all] N img ids to use for evaluation # catIds - [all] K cat ids to use for evaluation # iouThrs - [.5:.05:.95] T=10 IoU thresholds for evaluation # recThrs - [0:.01:1] R=101 recall thresholds for evaluation # areaRng - [...] A=4 object area ranges for evaluation # maxDets - [1 10 100] M=3 thresholds on max detections per image # useSegm - [1] if true evaluate against ground-truth segments # useCats - [1] if true use category labels for evaluation # Note: if useSegm=0 the evaluation is run on bounding boxes. # Note: if useCats=0 category labels are ignored as in proposal scoring. # Note: multiple areaRngs [Ax2] and maxDets [Mx1] can be specified. # # evaluate(): evaluates detections on every image and every category and # concats the results into the "evalImgs" with fields: # dtIds - [1xD] id for each of the D detections (dt) # gtIds - [1xG] id for each of the G ground truths (gt) # dtMatches - [TxD] matching gt id at each IoU or 0 # gtMatches - [TxG] matching dt id at each IoU or 0 # dtScores - [1xD] confidence of each dt # gtIgnore - [1xG] ignore flag for each gt # dtIgnore - [TxD] ignore flag for each dt at each IoU # # accumulate(): accumulates the per-image, per-category evaluation # results in "evalImgs" into the dictionary "eval" with fields: # params - parameters used for evaluation # date - date evaluation was performed # counts - [T,R,K,A,M] parameter dimensions (see above) # precision - [TxRxKxAxM] precision for every evaluation setting # recall - [TxKxAxM] max recall for every evaluation setting # Note: precision and recall==-1 for settings with no gt objects. # # See also coco, mask, pycocoDemo, pycocoEvalDemo # # Microsoft COCO Toolbox. version 2.0 # Data, paper, and tutorials available at: http://mscoco.org/ # Code written by Piotr Dollar and Tsung-Yi Lin, 2015. # Licensed under the Simplified BSD License [see coco/license.txt] def __init__(self, cocoGt=None, cocoDt=None): ''' Initialize CocoEval using coco APIs for gt and dt :param cocoGt: coco object with ground truth annotations :param cocoDt: coco object with detection results :return: None ''' self.cocoGt = cocoGt # ground truth COCO API self.cocoDt = cocoDt # detections COCO API self.params = {} # evaluation parameters self.evalImgs = defaultdict(list) # per-image per-category evaluation results [KxAxI] elements self.eval = {} # accumulated evaluation results self._gts = defaultdict(list) # gt for evaluation self._dts = defaultdict(list) # dt for evaluation self.params = Params() # parameters self._paramsEval = {} # parameters for evaluation self.stats = [] # result summarization self.ious = {} # ious between all gts and dts if not cocoGt is None: self.params.imgIds = sorted(cocoGt.getImgIds()) self.params.catIds = sorted(cocoGt.getCatIds()) def _prepare(self): ''' Prepare ._gts and ._dts for evaluation based on params :return: None ''' # def _toMask(objs, coco): # modify segmentation by reference for obj in objs: t = coco.imgs[obj['image_id']] if type(obj['segmentation']) == list: if type(obj['segmentation'][0]) == dict: print('debug') obj['segmentation'] = mask.frPyObjects(obj['segmentation'],t['height'],t['width']) if len(obj['segmentation']) == 1: obj['segmentation'] = obj['segmentation'][0] else: # an object can have multiple polygon regions # merge them into one RLE mask obj['segmentation'] = mask.merge(obj['segmentation']) elif type(obj['segmentation']) == dict and type(obj['segmentation']['counts']) == list: obj['segmentation'] = mask.frPyObjects([obj['segmentation']],t['height'],t['width'])[0] elif type(obj['segmentation']) == dict and \ type(obj['segmentation']['counts']) in string_types: pass else: raise Exception('segmentation format not supported.') p = self.params if p.useCats: gts=self.cocoGt.loadAnns(self.cocoGt.getAnnIds(imgIds=p.imgIds, catIds=p.catIds)) dts=self.cocoDt.loadAnns(self.cocoDt.getAnnIds(imgIds=p.imgIds, catIds=p.catIds)) else: gts=self.cocoGt.loadAnns(self.cocoGt.getAnnIds(imgIds=p.imgIds)) dts=self.cocoDt.loadAnns(self.cocoDt.getAnnIds(imgIds=p.imgIds)) if p.useSegm: _toMask(gts, self.cocoGt) _toMask(dts, self.cocoDt) self._gts = defaultdict(list) # gt for evaluation self._dts = defaultdict(list) # dt for evaluation for gt in gts: self._gts[gt['image_id'], gt['category_id']].append(gt) for dt in dts: self._dts[dt['image_id'], dt['category_id']].append(dt) self.evalImgs = defaultdict(list) # per-image per-category evaluation results self.eval = {} # accumulated evaluation results def evaluate(self): ''' Run per image evaluation on given images and store results (a list of dict) in self.evalImgs :return: None ''' tic = time.time() print('Running per image evaluation... ') p = self.params p.imgIds = list(np.unique(p.imgIds)) if p.useCats: p.catIds = list(np.unique(p.catIds)) p.maxDets = sorted(p.maxDets) self.params=p self._prepare() # loop through images, area range, max detection number catIds = p.catIds if p.useCats else [-1] computeIoU = self.computeIoU self.ious = {(imgId, catId): computeIoU(imgId, catId) \ for imgId in p.imgIds for catId in catIds} evaluateImg = self.evaluateImg maxDet = p.maxDets[-1] self.evalImgs = [evaluateImg(imgId, catId, areaRng, maxDet) for catId in catIds for areaRng in p.areaRng for imgId in p.imgIds ] self._paramsEval = copy.deepcopy(self.params) toc = time.time() print('DONE (t=%0.2fs).'%(toc-tic)) def computeIoU(self, imgId, catId): p = self.params if p.useCats: gt = self._gts[imgId,catId] dt = self._dts[imgId,catId] else: gt = [_ for cId in p.catIds for _ in self._gts[imgId,cId]] dt = [_ for cId in p.catIds for _ in self._dts[imgId,cId]] if len(gt) == 0 and len(dt) ==0: return [] dt = sorted(dt, key=lambda x: -x['score']) if len(dt) > p.maxDets[-1]: dt=dt[0:p.maxDets[-1]] if p.useSegm: g = [g['segmentation'] for g in gt] d = [d['segmentation'] for d in dt] else: g = [g['bbox'] for g in gt] d = [d['bbox'] for d in dt] # compute iou between each dt and gt region iscrowd = [int(o['iscrowd']) for o in gt] ious = mask.iou(d,g,iscrowd) return ious def evaluateImg(self, imgId, catId, aRng, maxDet): ''' perform evaluation for single category and image :return: dict (single image results) ''' # p = self.params if p.useCats: gt = self._gts[imgId,catId] dt = self._dts[imgId,catId] else: gt = [_ for cId in p.catIds for _ in self._gts[imgId,cId]] dt = [_ for cId in p.catIds for _ in self._dts[imgId,cId]] if len(gt) == 0 and len(dt) ==0: return None for g in gt: if 'ignore' not in g: g['ignore'] = 0 if g['iscrowd'] == 1 or g['ignore'] or (g['area']aRng[1]): g['_ignore'] = 1 else: g['_ignore'] = 0 # sort dt highest score first, sort gt ignore last # gt = sorted(gt, key=lambda x: x['_ignore']) gtind = [ind for (ind, g) in sorted(enumerate(gt), key=lambda ind_g: ind_g[1]['_ignore']) ] gt = [gt[ind] for ind in gtind] dt = sorted(dt, key=lambda x: -x['score'])[0:maxDet] iscrowd = [int(o['iscrowd']) for o in gt] # load computed ious N_iou = len(self.ious[imgId, catId]) ious = self.ious[imgId, catId][0:maxDet, np.array(gtind)] if N_iou >0 else self.ious[imgId, catId] T = len(p.iouThrs) G = len(gt) D = len(dt) gtm = np.zeros((T,G)) dtm = np.zeros((T,D)) gtIg = np.array([g['_ignore'] for g in gt]) dtIg = np.zeros((T,D)) if not len(ious)==0: for tind, t in enumerate(p.iouThrs): for dind, d in enumerate(dt): # information about best match so far (m=-1 -> unmatched) iou = min([t,1-1e-10]) m = -1 for gind, g in enumerate(gt): # if this gt already matched, and not a crowd, continue if gtm[tind,gind]>0 and not iscrowd[gind]: continue # if dt matched to reg gt, and on ignore gt, stop if m>-1 and gtIg[m]==0 and gtIg[gind]==1: break # continue to next gt unless better match made if ious[dind,gind] < iou: continue # match successful and best so far, store appropriately iou=ious[dind,gind] m=gind # if match made store id of match for both dt and gt if m ==-1: continue dtIg[tind,dind] = gtIg[m] dtm[tind,dind] = gt[m]['id'] gtm[tind,m] = d['id'] # set unmatched detections outside of area range to ignore a = np.array([d['area']aRng[1] for d in dt]).reshape((1, len(dt))) dtIg = np.logical_or(dtIg, np.logical_and(dtm==0, np.repeat(a,T,0))) # store results for given image and category return { 'image_id': imgId, 'category_id': catId, 'aRng': aRng, 'maxDet': maxDet, 'dtIds': [d['id'] for d in dt], 'gtIds': [g['id'] for g in gt], 'dtMatches': dtm, 'gtMatches': gtm, 'dtScores': [d['score'] for d in dt], 'gtIgnore': gtIg, 'dtIgnore': dtIg, } def accumulate(self, p = None): ''' Accumulate per image evaluation results and store the result in self.eval :param p: input params for evaluation :return: None ''' print('Accumulating evaluation results... ') tic = time.time() if not self.evalImgs: print('Please run evaluate() first') # allows input customized parameters if p is None: p = self.params p.catIds = p.catIds if p.useCats == 1 else [-1] T = len(p.iouThrs) R = len(p.recThrs) K = len(p.catIds) if p.useCats else 1 A = len(p.areaRng) M = len(p.maxDets) precision = -np.ones((T,R,K,A,M)) # -1 for the precision of absent categories recall = -np.ones((T,K,A,M)) # create dictionary for future indexing _pe = self._paramsEval catIds = _pe.catIds if _pe.useCats else [-1] setK = set(catIds) setA = set(map(tuple, _pe.areaRng)) setM = set(_pe.maxDets) setI = set(_pe.imgIds) # get inds to evaluate k_list = [n for n, k in enumerate(p.catIds) if k in setK] m_list = [m for n, m in enumerate(p.maxDets) if m in setM] a_list = [n for n, a in enumerate(map(lambda x: tuple(x), p.areaRng)) if a in setA] i_list = [n for n, i in enumerate(p.imgIds) if i in setI] # K0 = len(_pe.catIds) I0 = len(_pe.imgIds) A0 = len(_pe.areaRng) # retrieve E at each category, area range, and max number of detections for k, k0 in enumerate(k_list): Nk = k0*A0*I0 for a, a0 in enumerate(a_list): Na = a0*I0 for m, maxDet in enumerate(m_list): E = [self.evalImgs[Nk+Na+i] for i in i_list] E = filter(None, E) if len(E) == 0: continue dtScores = np.concatenate([e['dtScores'][0:maxDet] for e in E]) # different sorting method generates slightly different results. # mergesort is used to be consistent as Matlab implementation. inds = np.argsort(-dtScores, kind='mergesort') dtm = np.concatenate([e['dtMatches'][:,0:maxDet] for e in E], axis=1)[:,inds] dtIg = np.concatenate([e['dtIgnore'][:,0:maxDet] for e in E], axis=1)[:,inds] gtIg = np.concatenate([e['gtIgnore'] for e in E]) npig = len([ig for ig in gtIg if ig == 0]) if npig == 0: continue tps = np.logical_and( dtm, np.logical_not(dtIg) ) fps = np.logical_and(np.logical_not(dtm), np.logical_not(dtIg) ) tp_sum = np.cumsum(tps, axis=1).astype(dtype=np.float) fp_sum = np.cumsum(fps, axis=1).astype(dtype=np.float) for t, (tp, fp) in enumerate(zip(tp_sum, fp_sum)): tp = np.array(tp) fp = np.array(fp) nd = len(tp) rc = tp / npig pr = tp / (fp+tp+np.spacing(1)) q = np.zeros((R,)) if nd: recall[t,k,a,m] = rc[-1] else: recall[t,k,a,m] = 0 # numpy is slow without cython optimization for accessing elements # use python array gets significant speed improvement pr = pr.tolist(); q = q.tolist() for i in range(nd-1, 0, -1): if pr[i] > pr[i-1]: pr[i-1] = pr[i] inds = np.searchsorted(rc, p.recThrs) try: for ri, pi in enumerate(inds): q[ri] = pr[pi] except: pass precision[t,:,k,a,m] = np.array(q) self.eval = { 'params': p, 'counts': [T, R, K, A, M], 'date': datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"), 'precision': precision, 'recall': recall, } toc = time.time() print('DONE (t=%0.2fs).'%( toc-tic )) def summarize(self): ''' Compute and display summary metrics for evaluation results. Note this functin can *only* be applied on the default parameter setting ''' def _summarize( ap=1, iouThr=None, areaRng='all', maxDets=100 ): p = self.params iStr = ' {:<18} {} @[ IoU={:<9} | area={:>6} | maxDets={:>3} ] = {}' titleStr = 'Average Precision' if ap == 1 else 'Average Recall' typeStr = '(AP)' if ap==1 else '(AR)' iouStr = '%0.2f:%0.2f'%(p.iouThrs[0], p.iouThrs[-1]) if iouThr is None else '%0.2f'%(iouThr) areaStr = areaRng maxDetsStr = '%d'%(maxDets) aind = [i for i, aRng in enumerate(['all', 'small', 'medium', 'large']) if aRng == areaRng] mind = [i for i, mDet in enumerate([1, 10, 100]) if mDet == maxDets] if ap == 1: # dimension of precision: [TxRxKxAxM] s = self.eval['precision'] # IoU if iouThr is not None: t = np.where(iouThr == p.iouThrs)[0] s = s[t] # areaRng s = s[:,:,:,aind,mind] else: # dimension of recall: [TxKxAxM] s = self.eval['recall'] s = s[:,:,aind,mind] if len(s[s>-1])==0: mean_s = -1 else: mean_s = np.mean(s[s>-1]) print(iStr.format(titleStr, typeStr, iouStr, areaStr, maxDetsStr, '%.3f'%(float(mean_s)))) return mean_s if not self.eval: raise Exception('Please run accumulate() first') self.stats = np.zeros((12,)) self.stats[0] = _summarize(1) self.stats[1] = _summarize(1,iouThr=.5) self.stats[2] = _summarize(1,iouThr=.75) self.stats[3] = _summarize(1,areaRng='small') self.stats[4] = _summarize(1,areaRng='medium') self.stats[5] = _summarize(1,areaRng='large') self.stats[6] = _summarize(0,maxDets=1) self.stats[7] = _summarize(0,maxDets=10) self.stats[8] = _summarize(0,maxDets=100) self.stats[9] = _summarize(0,areaRng='small') self.stats[10] = _summarize(0,areaRng='medium') self.stats[11] = _summarize(0,areaRng='large') def __str__(self): self.summarize() class Params: ''' Params for coco evaluation api ''' def __init__(self): self.imgIds = [] self.catIds = [] # np.arange causes trouble. the data point on arange is slightly larger than the true value self.iouThrs = np.linspace(.5, 0.95, np.round((0.95-.5)/.05)+1, endpoint=True) self.recThrs = np.linspace(.0, 1.00, np.round((1.00-.0)/.01)+1, endpoint=True) self.maxDets = [1,10,100] self.areaRng = [ [0**2,1e5**2], [0**2, 32**2], [32**2, 96**2], [96**2, 1e5**2] ] self.useSegm = 0 self.useCats = 1 ================================================ FILE: lib/pycocotools/license.txt ================================================ Copyright (c) 2014, Piotr Dollar and Tsung-Yi Lin All rights reserved. Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. The views and conclusions contained in the software and documentation are those of the authors and should not be interpreted as representing official policies, either expressed or implied, of the FreeBSD Project. ================================================ FILE: lib/pycocotools/mask.py ================================================ __author__ = 'tsungyi' from . import _mask # Interface for manipulating masks stored in RLE format. # # RLE is a simple yet efficient format for storing binary masks. RLE # first divides a vector (or vectorized image) into a series of piecewise # constant regions and then for each piece simply stores the length of # that piece. For example, given M=[0 0 1 1 1 0 1] the RLE counts would # be [2 3 1 1], or for M=[1 1 1 1 1 1 0] the counts would be [0 6 1] # (note that the odd counts are always the numbers of zeros). Instead of # storing the counts directly, additional compression is achieved with a # variable bitrate representation based on a common scheme called LEB128. # # Compression is greatest given large piecewise constant regions. # Specifically, the size of the RLE is proportional to the number of # *boundaries* in M (or for an image the number of boundaries in the y # direction). Assuming fairly simple shapes, the RLE representation is # O(sqrt(n)) where n is number of pixels in the object. Hence space usage # is substantially lower, especially for large simple objects (large n). # # Many common operations on masks can be computed directly using the RLE # (without need for decoding). This includes computations such as area, # union, intersection, etc. All of these operations are linear in the # size of the RLE, in other words they are O(sqrt(n)) where n is the area # of the object. Computing these operations on the original mask is O(n). # Thus, using the RLE can result in substantial computational savings. # # The following API functions are defined: # encode - Encode binary masks using RLE. # decode - Decode binary masks encoded via RLE. # merge - Compute union or intersection of encoded masks. # iou - Compute intersection over union between masks. # area - Compute area of encoded masks. # toBbox - Get bounding boxes surrounding encoded masks. # frPyObjects - Convert polygon, bbox, and uncompressed RLE to encoded RLE mask. # # Usage: # Rs = encode( masks ) # masks = decode( Rs ) # R = merge( Rs, intersect=false ) # o = iou( dt, gt, iscrowd ) # a = area( Rs ) # bbs = toBbox( Rs ) # Rs = frPyObjects( [pyObjects], h, w ) # # In the API the following formats are used: # Rs - [dict] Run-length encoding of binary masks # R - dict Run-length encoding of binary mask # masks - [hxwxn] Binary mask(s) (must have type np.ndarray(dtype=uint8) in column-major order) # iscrowd - [nx1] list of np.ndarray. 1 indicates corresponding gt image has crowd region to ignore # bbs - [nx4] Bounding box(es) stored as [x y w h] # poly - Polygon stored as [[x1 y1 x2 y2...],[x1 y1 ...],...] (2D list) # dt,gt - May be either bounding boxes or encoded masks # Both poly and bbs are 0-indexed (bbox=[0 0 1 1] encloses first pixel). # # Finally, a note about the intersection over union (iou) computation. # The standard iou of a ground truth (gt) and detected (dt) object is # iou(gt,dt) = area(intersect(gt,dt)) / area(union(gt,dt)) # For "crowd" regions, we use a modified criteria. If a gt object is # marked as "iscrowd", we allow a dt to match any subregion of the gt. # Choosing gt' in the crowd gt that best matches the dt can be done using # gt'=intersect(dt,gt). Since by definition union(gt',dt)=dt, computing # iou(gt,dt,iscrowd) = iou(gt',dt) = area(intersect(gt,dt)) / area(dt) # For crowd gt regions we use this modified criteria above for the iou. # # To compile run "python setup.py build_ext --inplace" # Please do not contact us for help with compiling. # # Microsoft COCO Toolbox. version 2.0 # Data, paper, and tutorials available at: http://mscoco.org/ # Code written by Piotr Dollar and Tsung-Yi Lin, 2015. # Licensed under the Simplified BSD License [see coco/license.txt] encode = _mask.encode decode = _mask.decode iou = _mask.iou merge = _mask.merge area = _mask.area toBbox = _mask.toBbox frPyObjects = _mask.frPyObjects ================================================ FILE: lib/pycocotools/maskApi.c ================================================ /************************************************************************** * Microsoft COCO Toolbox. version 2.0 * Data, paper, and tutorials available at: http://mscoco.org/ * Code written by Piotr Dollar and Tsung-Yi Lin, 2015. * Licensed under the Simplified BSD License [see coco/license.txt] **************************************************************************/ #include "maskApi.h" #include #include uint umin( uint a, uint b ) { return (ab) ? a : b; } void rleInit( RLE *R, siz h, siz w, siz m, uint *cnts ) { R->h=h; R->w=w; R->m=m; R->cnts=(m==0)?0:malloc(sizeof(uint)*m); if(cnts) for(siz j=0; jcnts[j]=cnts[j]; } void rleFree( RLE *R ) { free(R->cnts); R->cnts=0; } void rlesInit( RLE **R, siz n ) { *R = (RLE*) malloc(sizeof(RLE)*n); for(siz i=0; i0 ) { c=umin(ca,cb); cc+=c; ct=0; ca-=c; if(!ca && a0) { crowd=iscrowd!=NULL && iscrowd[g]; if(dt[d].h!=gt[g].h || dt[d].w!=gt[g].w) { o[g*m+d]=-1; continue; } siz ka, kb, a, b; uint c, ca, cb, ct, i, u; bool va, vb; ca=dt[d].cnts[0]; ka=dt[d].m; va=vb=0; cb=gt[g].cnts[0]; kb=gt[g].m; a=b=1; i=u=0; ct=1; while( ct>0 ) { c=umin(ca,cb); if(va||vb) { u+=c; if(va&&vb) i+=c; } ct=0; ca-=c; if(!ca && ad?1:c=dy && xs>xe) || (dxye); if(flip) { t=xs; xs=xe; xe=t; t=ys; ys=ye; ye=t; } s = dx>=dy ? (double)(ye-ys)/dx : (double)(xe-xs)/dy; if(dx>=dy) for( int d=0; d<=dx; d++ ) { t=flip?dx-d:d; u[m]=t+xs; v[m]=(int)(ys+s*t+.5); m++; } else for( int d=0; d<=dy; d++ ) { t=flip?dy-d:d; v[m]=t+ys; u[m]=(int)(xs+s*t+.5); m++; } } // get points along y-boundary and downsample free(x); free(y); k=m; m=0; double xd, yd; x=malloc(sizeof(int)*k); y=malloc(sizeof(int)*k); for( j=1; jw-1 ) continue; yd=(double)(v[j]h) yd=h; yd=ceil(yd); x[m]=(int) xd; y[m]=(int) yd; m++; } // compute rle encoding given y-boundary points k=m; a=malloc(sizeof(uint)*(k+1)); for( j=0; j0) b[m++]=a[j++]; else { j++; if(jm, p=0; long x; bool more; char *s=malloc(sizeof(char)*m*6); for( i=0; icnts[i]; if(i>2) x-=(long) R->cnts[i-2]; more=1; while( more ) { char c=x & 0x1f; x >>= 5; more=(c & 0x10) ? x!=-1 : x!=0; if(more) c |= 0x20; c+=48; s[p++]=c; } } s[p]=0; return s; } void rleFrString( RLE *R, char *s, siz h, siz w ) { siz m=0, p=0, k; long x; bool more; uint *cnts; while( s[m] ) m++; cnts=malloc(sizeof(uint)*m); m=0; while( s[p] ) { x=0; k=0; more=1; while( more ) { char c=s[p]-48; x |= (c & 0x1f) << 5*k; more = c & 0x20; p++; k++; if(!more && (c & 0x10)) x |= -1 << 5*k; } if(m>2) x+=(long) cnts[m-2]; cnts[m++]=(uint) x; } rleInit(R,h,w,m,cnts); free(cnts); } ================================================ FILE: lib/pycocotools/maskApi.h ================================================ /************************************************************************** * Microsoft COCO Toolbox. version 2.0 * Data, paper, and tutorials available at: http://mscoco.org/ * Code written by Piotr Dollar and Tsung-Yi Lin, 2015. * Licensed under the Simplified BSD License [see coco/license.txt] **************************************************************************/ #pragma once #include typedef unsigned int uint; typedef unsigned long siz; typedef unsigned char byte; typedef double* BB; typedef struct { siz h, w, m; uint *cnts; } RLE; // Initialize/destroy RLE. void rleInit( RLE *R, siz h, siz w, siz m, uint *cnts ); void rleFree( RLE *R ); // Initialize/destroy RLE array. void rlesInit( RLE **R, siz n ); void rlesFree( RLE **R, siz n ); // Encode binary masks using RLE. void rleEncode( RLE *R, const byte *mask, siz h, siz w, siz n ); // Decode binary masks encoded via RLE. void rleDecode( const RLE *R, byte *mask, siz n ); // Compute union or intersection of encoded masks. void rleMerge( const RLE *R, RLE *M, siz n, bool intersect ); // Compute area of encoded masks. void rleArea( const RLE *R, siz n, uint *a ); // Compute intersection over union between masks. void rleIou( RLE *dt, RLE *gt, siz m, siz n, byte *iscrowd, double *o ); // Compute intersection over union between bounding boxes. void bbIou( BB dt, BB gt, siz m, siz n, byte *iscrowd, double *o ); // Get bounding boxes surrounding encoded masks. void rleToBbox( const RLE *R, BB bb, siz n ); // Convert bounding boxes to encoded masks. void rleFrBbox( RLE *R, const BB bb, siz h, siz w, siz n ); // Convert polygon to encoded mask. void rleFrPoly( RLE *R, const double *xy, siz k, siz h, siz w ); // Get compressed string representation of encoded mask. char* rleToString( const RLE *R ); // Convert from compressed string representation of encoded mask. void rleFrString( RLE *R, char *s, siz h, siz w ); ================================================ FILE: lib/roi_data_layer/__init__.py ================================================ # -------------------------------------------------------- # Fast R-CNN # Copyright (c) 2015 Microsoft # Licensed under The MIT License [see LICENSE for details] # Written by Ross Girshick # -------------------------------------------------------- ================================================ FILE: lib/roi_data_layer/minibatch.py ================================================ # -------------------------------------------------------- # Fast R-CNN # Copyright (c) 2015 Microsoft # Licensed under The MIT License [see LICENSE for details] # Written by Ross Girshick and Xinlei Chen # -------------------------------------------------------- """Compute minibatch blobs for training a Fast R-CNN network.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function import numpy as np import numpy.random as npr from scipy.misc import imread from model.utils.config import cfg from model.utils.blob import prep_im_for_blob, im_list_to_blob import pdb def get_minibatch(roidb, num_classes): """Given a roidb, construct a minibatch sampled from it.""" num_images = len(roidb) # Sample random scales to use for each image in this batch random_scale_inds = npr.randint(0, high=len(cfg.TRAIN.SCALES), size=num_images) assert(cfg.TRAIN.BATCH_SIZE % num_images == 0), \ 'num_images ({}) must divide BATCH_SIZE ({})'. \ format(num_images, cfg.TRAIN.BATCH_SIZE) # Get the input image blob, formatted for caffe im_blob, im_scales = _get_image_blob(roidb, random_scale_inds) blobs = {'data': im_blob} assert len(im_scales) == 1, "Single batch only" assert len(roidb) == 1, "Single batch only" # gt boxes: (x1, y1, x2, y2, cls) if cfg.TRAIN.USE_ALL_GT: # Include all ground truth boxes gt_inds = np.where(roidb[0]['gt_classes'] != 0)[0] else: # For the COCO ground truth boxes, exclude the ones that are ''iscrowd'' gt_inds = np.where((roidb[0]['gt_classes'] != 0) & np.all(roidb[0]['gt_overlaps'].toarray() > -1.0, axis=1))[0] gt_boxes = np.empty((len(gt_inds), 5), dtype=np.float32) gt_boxes[:, 0:4] = roidb[0]['boxes'][gt_inds, :] * im_scales[0] gt_boxes[:, 4] = roidb[0]['gt_classes'][gt_inds] blobs['gt_boxes'] = gt_boxes blobs['im_info'] = np.array( [[im_blob.shape[1], im_blob.shape[2], im_scales[0]]], dtype=np.float32) blobs['img_id'] = roidb[0]['img_id'] return blobs def _get_image_blob(roidb, scale_inds): """Builds an input blob from the images in the roidb at the specified scales. """ num_images = len(roidb) processed_ims = [] im_scales = [] for i in range(num_images): #im = cv2.imread(roidb[i]['image']) im = imread(roidb[i]['image']) if len(im.shape) == 2: im = im[:,:,np.newaxis] im = np.concatenate((im,im,im), axis=2) # flip the channel, since the original one using cv2 # rgb -> bgr im = im[:,:,::-1] if roidb[i]['flipped']: im = im[:, ::-1, :] target_size = cfg.TRAIN.SCALES[scale_inds[i]] im, im_scale = prep_im_for_blob(im, cfg.PIXEL_MEANS, target_size, cfg.TRAIN.MAX_SIZE) im_scales.append(im_scale) processed_ims.append(im) # Create a blob to hold the input images blob = im_list_to_blob(processed_ims) return blob, im_scales ================================================ FILE: lib/roi_data_layer/roibatchLoader.py ================================================ """The data layer used during training to train a Fast R-CNN network. """ from __future__ import absolute_import from __future__ import division from __future__ import print_function import torch.utils.data as data from PIL import Image import torch from model.utils.config import cfg from roi_data_layer.minibatch import get_minibatch, get_minibatch from model.rpn.bbox_transform import bbox_transform_inv, clip_boxes import numpy as np import random import time import pdb class roibatchLoader(data.Dataset): def __init__(self, roidb, ratio_list, ratio_index, batch_size, num_classes, training=True, normalize=None): self._roidb = roidb self._num_classes = num_classes # we make the height of image consistent to trim_height, trim_width self.trim_height = cfg.TRAIN.TRIM_HEIGHT self.trim_width = cfg.TRAIN.TRIM_WIDTH self.max_num_box = cfg.MAX_NUM_GT_BOXES self.training = training self.normalize = normalize self.ratio_list = ratio_list self.ratio_index = ratio_index self.batch_size = batch_size self.data_size = len(self.ratio_list) # given the ratio_list, we want to make the ratio same for each batch. self.ratio_list_batch = torch.Tensor(self.data_size).zero_() num_batch = int(np.ceil(len(ratio_index) / batch_size)) for i in range(num_batch): left_idx = i*batch_size right_idx = min((i+1)*batch_size-1, self.data_size-1) if ratio_list[right_idx] < 1: # for ratio < 1, we preserve the leftmost in each batch. target_ratio = ratio_list[left_idx] elif ratio_list[left_idx] > 1: # for ratio > 1, we preserve the rightmost in each batch. target_ratio = ratio_list[right_idx] else: # for ratio cross 1, we make it to be 1. target_ratio = 1 self.ratio_list_batch[left_idx:(right_idx+1)] = torch.tensor(target_ratio.astype(np.float64)) # trainset ratio list ,each batch is same number def __getitem__(self, index): if self.training: index_ratio = int(self.ratio_index[index]) else: index_ratio = index # get the anchor index for current sample index # here we set the anchor index to the last one # sample in this group minibatch_db = [self._roidb[index_ratio]] blobs = get_minibatch(minibatch_db, self._num_classes) data = torch.from_numpy(blobs['data']) im_info = torch.from_numpy(blobs['im_info']) # we need to random shuffle the bounding box. data_height, data_width = data.size(1), data.size(2) if self.training: np.random.shuffle(blobs['gt_boxes']) gt_boxes = torch.from_numpy(blobs['gt_boxes']) ######################################################## # padding the input image to fixed size for each group # ######################################################## # NOTE1: need to cope with the case where a group cover both conditions. (done) # NOTE2: need to consider the situation for the tail samples. (no worry) # NOTE3: need to implement a parallel data loader. (no worry) # get the index range # if the image need to crop, crop to the target size. ratio = self.ratio_list_batch[index] if self._roidb[index_ratio]['need_crop']: if ratio < 1: # this means that data_width << data_height, we need to crop the # data_height min_y = int(torch.min(gt_boxes[:,1])) max_y = int(torch.max(gt_boxes[:,3])) trim_size = int(np.floor(data_width / ratio)) if trim_size > data_height: trim_size = data_height box_region = max_y - min_y + 1 if min_y == 0: y_s = 0 else: if (box_region-trim_size) < 0: y_s_min = max(max_y-trim_size, 0) y_s_max = min(min_y, data_height-trim_size) if y_s_min == y_s_max: y_s = y_s_min else: y_s = np.random.choice(range(y_s_min, y_s_max)) else: y_s_add = int((box_region-trim_size)/2) if y_s_add == 0: y_s = min_y else: y_s = np.random.choice(range(min_y, min_y+y_s_add)) # crop the image data = data[:, y_s:(y_s + trim_size), :, :] # shift y coordiante of gt_boxes gt_boxes[:, 1] = gt_boxes[:, 1] - float(y_s) gt_boxes[:, 3] = gt_boxes[:, 3] - float(y_s) # update gt bounding box according the trip gt_boxes[:, 1].clamp_(0, trim_size - 1) gt_boxes[:, 3].clamp_(0, trim_size - 1) else: # this means that data_width >> data_height, we need to crop the # data_width min_x = int(torch.min(gt_boxes[:,0])) max_x = int(torch.max(gt_boxes[:,2])) trim_size = int(np.ceil(data_height * ratio)) if trim_size > data_width: trim_size = data_width box_region = max_x - min_x + 1 if min_x == 0: x_s = 0 else: if (box_region-trim_size) < 0: x_s_min = max(max_x-trim_size, 0) x_s_max = min(min_x, data_width-trim_size) if x_s_min == x_s_max: x_s = x_s_min else: x_s = np.random.choice(range(x_s_min, x_s_max)) else: x_s_add = int((box_region-trim_size)/2) if x_s_add == 0: x_s = min_x else: x_s = np.random.choice(range(min_x, min_x+x_s_add)) # crop the image data = data[:, :, x_s:(x_s + trim_size), :] # shift x coordiante of gt_boxes gt_boxes[:, 0] = gt_boxes[:, 0] - float(x_s) gt_boxes[:, 2] = gt_boxes[:, 2] - float(x_s) # update gt bounding box according the trip gt_boxes[:, 0].clamp_(0, trim_size - 1) gt_boxes[:, 2].clamp_(0, trim_size - 1) # based on the ratio, padding the image. if ratio < 1: # this means that data_width < data_height trim_size = int(np.floor(data_width / ratio)) padding_data = torch.FloatTensor(int(np.ceil(data_width / ratio)), \ data_width, 3).zero_() padding_data[:data_height, :, :] = data[0] # update im_info im_info[0, 0] = padding_data.size(0) # print("height %d %d \n" %(index, anchor_idx)) elif ratio > 1: # this means that data_width > data_height # if the image need to crop. padding_data = torch.FloatTensor(data_height, \ int(np.ceil(data_height * ratio)), 3).zero_() padding_data[:, :data_width, :] = data[0] im_info[0, 1] = padding_data.size(1) else: trim_size = min(data_height, data_width) padding_data = torch.FloatTensor(trim_size, trim_size, 3).zero_() padding_data = data[0][:trim_size, :trim_size, :] # gt_boxes.clamp_(0, trim_size) gt_boxes[:, :4].clamp_(0, trim_size) im_info[0, 0] = trim_size im_info[0, 1] = trim_size # check the bounding box: not_keep = (gt_boxes[:,0] == gt_boxes[:,2]) | (gt_boxes[:,1] == gt_boxes[:,3]) keep = torch.nonzero(not_keep == 0).view(-1) gt_boxes_padding = torch.FloatTensor(self.max_num_box, gt_boxes.size(1)).zero_() if keep.numel() != 0: gt_boxes = gt_boxes[keep] num_boxes = min(gt_boxes.size(0), self.max_num_box) gt_boxes_padding[:num_boxes,:] = gt_boxes[:num_boxes] else: num_boxes = 0 # permute trim_data to adapt to downstream processing padding_data = padding_data.permute(2, 0, 1).contiguous() im_info = im_info.view(3) return padding_data, im_info, gt_boxes_padding, num_boxes else: data = data.permute(0, 3, 1, 2).contiguous().view(3, data_height, data_width) im_info = im_info.view(3) gt_boxes = torch.FloatTensor([1,1,1,1,1]) num_boxes = 0 return data, im_info, gt_boxes, num_boxes def __len__(self): return len(self._roidb) ================================================ FILE: lib/roi_data_layer/roidb.py ================================================ """Transform a roidb into a trainable roidb by adding a bunch of metadata.""" from __future__ import absolute_import from __future__ import division from __future__ import print_function import os import pickle import datasets import numpy as np from model.utils.config import cfg from datasets.factory import get_imdb import PIL import pdb def prepare_roidb(imdb): """Enrich the imdb's roidb by adding some derived quantities that are useful for training. This function precomputes the maximum overlap, taken over ground-truth boxes, between each ROI and each ground-truth box. The class with maximum overlap is also recorded. """ roidb = imdb.roidb if not (imdb.name.startswith('coco')): cache_file = os.path.join(imdb.cache_path, imdb.name + '_sizes.pkl') if os.path.exists(cache_file): print('Image sizes loaded from %s' % cache_file) with open(cache_file, 'rb') as f: sizes = pickle.load(f) else: print('Extracting image sizes... (It may take long time)') sizes = [PIL.Image.open(imdb.image_path_at(i)).size for i in range(imdb.num_images)] with open(cache_file, 'wb') as f: pickle.dump(sizes, f) print('Done!!') for i in range(len(imdb.image_index)): roidb[i]['img_id'] = imdb.image_id_at(i) roidb[i]['image'] = imdb.image_path_at(i) if not (imdb.name.startswith('coco')): roidb[i]['width'] = sizes[i][0] roidb[i]['height'] = sizes[i][1] # need gt_overlaps as a dense array for argmax gt_overlaps = roidb[i]['gt_overlaps'].toarray() # max overlap with gt over classes (columns) max_overlaps = gt_overlaps.max(axis=1) # gt class that had the max overlap max_classes = gt_overlaps.argmax(axis=1) roidb[i]['max_classes'] = max_classes roidb[i]['max_overlaps'] = max_overlaps # sanity checks # max overlap of 0 => class should be zero (background) zero_inds = np.where(max_overlaps == 0)[0] assert all(max_classes[zero_inds] == 0) # max overlap > 0 => class should not be zero (must be a fg class) nonzero_inds = np.where(max_overlaps > 0)[0] assert all(max_classes[nonzero_inds] != 0) def rank_roidb_ratio(roidb): # rank roidb based on the ratio between width and height. ratio_large = 2 # largest ratio to preserve. ratio_small = 0.5 # smallest ratio to preserve. ratio_list = [] for i in range(len(roidb)): width = roidb[i]['width'] height = roidb[i]['height'] ratio = width / float(height) if ratio > ratio_large: roidb[i]['need_crop'] = 1 ratio = ratio_large elif ratio < ratio_small: roidb[i]['need_crop'] = 1 ratio = ratio_small else: roidb[i]['need_crop'] = 0 ratio_list.append(ratio) ratio_list = np.array(ratio_list) ratio_index = np.argsort(ratio_list) return ratio_list[ratio_index], ratio_index def filter_roidb(roidb): # filter the image without bounding box. print('before filtering, there are %d images...' % (len(roidb))) i = 0 while i < len(roidb): if len(roidb[i]['boxes']) == 0: del roidb[i] i -= 1 i += 1 print('after filtering, there are %d images...' % (len(roidb))) return roidb def combined_roidb(imdb_names, training=True): """ Combine multiple roidbs """ def get_training_roidb(imdb): """Returns a roidb (Region of Interest database) for use in training.""" if cfg.TRAIN.USE_FLIPPED: print('Appending horizontally-flipped training examples...') imdb.append_flipped_images() print('done') print('Preparing training data...') prepare_roidb(imdb) #ratio_index = rank_roidb_ratio(imdb) print('done') return imdb.roidb def get_roidb(imdb_name): imdb = get_imdb(imdb_name) print('Loaded dataset `{:s}`'.format(imdb.name)) imdb.set_proposal_method(cfg.TRAIN.PROPOSAL_METHOD) print('Set proposal method: {:s}'.format(cfg.TRAIN.PROPOSAL_METHOD)) roidb = get_training_roidb(imdb) return roidb roidbs = [get_roidb(s) for s in imdb_names.split('+')] roidb = roidbs[0] if len(roidbs) > 1: for r in roidbs[1:]: roidb.extend(r) tmp = get_imdb(imdb_names.split('+')[1]) imdb = datasets.imdb.imdb(imdb_names, tmp.classes) else: imdb = get_imdb(imdb_names) if training: roidb = filter_roidb(roidb) ratio_list, ratio_index = rank_roidb_ratio(roidb) return imdb, roidb, ratio_list, ratio_index ================================================ FILE: lib/setup.py ================================================ from __future__ import print_function # -------------------------------------------------------- # Fast R-CNN # Copyright (c) 2015 Microsoft # Licensed under The MIT License [see LICENSE for details] # Written by Ross Girshick # -------------------------------------------------------- import os from os.path import join as pjoin import numpy as np from distutils.core import setup from distutils.extension import Extension from Cython.Distutils import build_ext def find_in_path(name, path): "Find a file in a search path" # adapted fom http://code.activestate.com/recipes/52224-find-a-file-given-a-search-path/ for dir in path.split(os.pathsep): binpath = pjoin(dir, name) if os.path.exists(binpath): return os.path.abspath(binpath) return None # def locate_cuda(): # """Locate the CUDA environment on the system # # Returns a dict with keys 'home', 'nvcc', 'include', and 'lib64' # and values giving the absolute path to each directory. # # Starts by looking for the CUDAHOME env variable. If not found, everything # is based on finding 'nvcc' in the PATH. # """ # # # first check if the CUDAHOME env variable is in use # if 'CUDAHOME' in os.environ: # home = os.environ['CUDAHOME'] # nvcc = pjoin(home, 'bin', 'nvcc') # else: # # otherwise, search the PATH for NVCC # default_path = pjoin(os.sep, 'usr', 'local', 'cuda', 'bin') # nvcc = find_in_path('nvcc', os.environ['PATH'] + os.pathsep + default_path) # if nvcc is None: # raise EnvironmentError('The nvcc binary could not be ' # 'located in your $PATH. Either add it to your path, or set $CUDAHOME') # home = os.path.dirname(os.path.dirname(nvcc)) # # cudaconfig = {'home': home, 'nvcc': nvcc, # 'include': pjoin(home, 'include'), # 'lib64': pjoin(home, 'lib64')} # for k, v in cudaconfig.iteritems(): # if not os.path.exists(v): # raise EnvironmentError('The CUDA %s path could not be located in %s' % (k, v)) # # return cudaconfig # CUDA = locate_cuda() # Obtain the numpy include directory. This logic works across numpy versions. try: numpy_include = np.get_include() except AttributeError: numpy_include = np.get_numpy_include() def customize_compiler_for_nvcc(self): """inject deep into distutils to customize how the dispatch to gcc/nvcc works. If you subclass UnixCCompiler, it's not trivial to get your subclass injected in, and still have the right customizations (i.e. distutils.sysconfig.customize_compiler) run on it. So instead of going the OO route, I have this. Note, it's kindof like a wierd functional subclassing going on.""" # tell the compiler it can processes .cu self.src_extensions.append('.cu') # save references to the default compiler_so and _comple methods default_compiler_so = self.compiler_so super = self._compile # now redefine the _compile method. This gets executed for each # object but distutils doesn't have the ability to change compilers # based on source extension: we add it. def _compile(obj, src, ext, cc_args, extra_postargs, pp_opts): print(extra_postargs) if os.path.splitext(src)[1] == '.cu': # use the cuda for .cu files self.set_executable('compiler_so', CUDA['nvcc']) # use only a subset of the extra_postargs, which are 1-1 translated # from the extra_compile_args in the Extension class postargs = extra_postargs['nvcc'] else: postargs = extra_postargs['gcc'] super(obj, src, ext, cc_args, postargs, pp_opts) # reset the default compiler_so, which we might have changed for cuda self.compiler_so = default_compiler_so # inject our redefined _compile method into the class self._compile = _compile # run the customize_compiler class custom_build_ext(build_ext): def build_extensions(self): customize_compiler_for_nvcc(self.compiler) build_ext.build_extensions(self) ext_modules = [ Extension( "model.utils.cython_bbox", ["model/utils/bbox.pyx"], extra_compile_args={'gcc': ["-Wno-cpp", "-Wno-unused-function"]}, include_dirs=[numpy_include] ), Extension( 'pycocotools._mask', sources=['pycocotools/maskApi.c', 'pycocotools/_mask.pyx'], include_dirs=[numpy_include, 'pycocotools'], extra_compile_args={ 'gcc': ['-Wno-cpp', '-Wno-unused-function', '-std=c99']}, ), ] setup( name='faster_rcnn', ext_modules=ext_modules, # inject our custom trigger cmdclass={'build_ext': custom_build_ext}, ) ================================================ FILE: requirements.txt ================================================ cython cffi opencv-python scipy msgpack easydict matplotlib pyyaml tensorboardX ================================================ FILE: test_net.py ================================================ # -------------------------------------------------------- # Tensorflow Faster R-CNN # Licensed under The MIT License [see LICENSE for details] # Written by Jiasen Lu, Jianwei Yang, based on code from Ross Girshick # -------------------------------------------------------- from __future__ import absolute_import from __future__ import division from __future__ import print_function import _init_paths import os import sys import numpy as np import argparse import pprint import pdb import time import cv2 import torch from torch.autograd import Variable import torch.nn as nn import torch.optim as optim import pickle from roi_data_layer.roidb import combined_roidb from roi_data_layer.roibatchLoader import roibatchLoader from model.utils.config import cfg, cfg_from_file, cfg_from_list, get_output_dir from model.rpn.bbox_transform import clip_boxes from model.nms.nms_wrapper import nms from model.rpn.bbox_transform import bbox_transform_inv from model.utils.net_utils import save_net, load_net, vis_detections from model.faster_rcnn.vgg16 import vgg16 from model.faster_rcnn.resnet import resnet try: xrange # Python 2 except NameError: xrange = range # Python 3 def parse_args(): """ Parse input arguments """ parser = argparse.ArgumentParser(description='Train a Fast R-CNN network') parser.add_argument('--dataset', dest='dataset', help='training dataset', default='pascal_voc', type=str) parser.add_argument('--cfg', dest='cfg_file', help='optional config file', default='cfgs/vgg16.yml', type=str) parser.add_argument('--net', dest='net', help='vgg16, res50, res101, res152', default='res101', type=str) parser.add_argument('--set', dest='set_cfgs', help='set config keys', default=None, nargs=argparse.REMAINDER) parser.add_argument('--load_dir', dest='load_dir', help='directory to load models', default="models", type=str) parser.add_argument('--cuda', dest='cuda', help='whether use CUDA', action='store_true') parser.add_argument('--ls', dest='large_scale', help='whether use large imag scale', action='store_true') parser.add_argument('--mGPUs', dest='mGPUs', help='whether use multiple GPUs', action='store_true') parser.add_argument('--cag', dest='class_agnostic', help='whether perform class_agnostic bbox regression', action='store_true') parser.add_argument('--parallel_type', dest='parallel_type', help='which part of model to parallel, 0: all, 1: model before roi pooling', default=0, type=int) parser.add_argument('--checksession', dest='checksession', help='checksession to load model', default=1, type=int) parser.add_argument('--checkepoch', dest='checkepoch', help='checkepoch to load network', default=1, type=int) parser.add_argument('--checkpoint', dest='checkpoint', help='checkpoint to load network', default=10021, type=int) parser.add_argument('--vis', dest='vis', help='visualization mode', action='store_true') args = parser.parse_args() return args lr = cfg.TRAIN.LEARNING_RATE momentum = cfg.TRAIN.MOMENTUM weight_decay = cfg.TRAIN.WEIGHT_DECAY if __name__ == '__main__': args = parse_args() print('Called with args:') print(args) if torch.cuda.is_available() and not args.cuda: print("WARNING: You have a CUDA device, so you should probably run with --cuda") np.random.seed(cfg.RNG_SEED) if args.dataset == "pascal_voc": args.imdb_name = "voc_2007_trainval" args.imdbval_name = "voc_2007_test" args.set_cfgs = ['ANCHOR_SCALES', '[8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]'] elif args.dataset == "pascal_voc_0712": args.imdb_name = "voc_2007_trainval+voc_2012_trainval" args.imdbval_name = "voc_2007_test" args.set_cfgs = ['ANCHOR_SCALES', '[8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]'] elif args.dataset == "coco": args.imdb_name = "coco_2014_train+coco_2014_valminusminival" args.imdbval_name = "coco_2014_minival" args.set_cfgs = ['ANCHOR_SCALES', '[4, 8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]'] elif args.dataset == "imagenet": args.imdb_name = "imagenet_train" args.imdbval_name = "imagenet_val" args.set_cfgs = ['ANCHOR_SCALES', '[8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]'] elif args.dataset == "vg": args.imdb_name = "vg_150-50-50_minitrain" args.imdbval_name = "vg_150-50-50_minival" args.set_cfgs = ['ANCHOR_SCALES', '[4, 8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]'] args.cfg_file = "cfgs/{}_ls.yml".format(args.net) if args.large_scale else "cfgs/{}.yml".format(args.net) if args.cfg_file is not None: cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) print('Using config:') pprint.pprint(cfg) cfg.TRAIN.USE_FLIPPED = False imdb, roidb, ratio_list, ratio_index = combined_roidb(args.imdbval_name, False) imdb.competition_mode(on=True) print('{:d} roidb entries'.format(len(roidb))) input_dir = args.load_dir + "/" + args.net + "/" + args.dataset if not os.path.exists(input_dir): raise Exception('There is no input directory for loading network from ' + input_dir) load_name = os.path.join(input_dir, 'faster_rcnn_{}_{}_{}.pth'.format(args.checksession, args.checkepoch, args.checkpoint)) # initilize the network here. if args.net == 'vgg16': fasterRCNN = vgg16(imdb.classes, pretrained=False, class_agnostic=args.class_agnostic) elif args.net == 'res101': fasterRCNN = resnet(imdb.classes, 101, pretrained=False, class_agnostic=args.class_agnostic) elif args.net == 'res50': fasterRCNN = resnet(imdb.classes, 50, pretrained=False, class_agnostic=args.class_agnostic) elif args.net == 'res152': fasterRCNN = resnet(imdb.classes, 152, pretrained=False, class_agnostic=args.class_agnostic) else: print("network is not defined") pdb.set_trace() fasterRCNN.create_architecture() print("load checkpoint %s" % (load_name)) checkpoint = torch.load(load_name) fasterRCNN.load_state_dict(checkpoint['model']) if 'pooling_mode' in checkpoint.keys(): cfg.POOLING_MODE = checkpoint['pooling_mode'] print('load model successfully!') # initilize the tensor holder here. im_data = torch.FloatTensor(1) im_info = torch.FloatTensor(1) num_boxes = torch.LongTensor(1) gt_boxes = torch.FloatTensor(1) # ship to cuda if args.cuda: im_data = im_data.cuda() im_info = im_info.cuda() num_boxes = num_boxes.cuda() gt_boxes = gt_boxes.cuda() # make variable im_data = Variable(im_data) im_info = Variable(im_info) num_boxes = Variable(num_boxes) gt_boxes = Variable(gt_boxes) if args.cuda: cfg.CUDA = True if args.cuda: fasterRCNN.cuda() start = time.time() max_per_image = 100 vis = args.vis if vis: thresh = 0.05 else: thresh = 0.0 save_name = 'faster_rcnn_10' num_images = len(imdb.image_index) all_boxes = [[[] for _ in xrange(num_images)] for _ in xrange(imdb.num_classes)] output_dir = get_output_dir(imdb, save_name) dataset = roibatchLoader(roidb, ratio_list, ratio_index, 1, \ imdb.num_classes, training=False, normalize = False) dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=0, pin_memory=True) data_iter = iter(dataloader) _t = {'im_detect': time.time(), 'misc': time.time()} det_file = os.path.join(output_dir, 'detections.pkl') fasterRCNN.eval() empty_array = np.transpose(np.array([[],[],[],[],[]]), (1,0)) for i in range(num_images): data = next(data_iter) im_data.data.resize_(data[0].size()).copy_(data[0]) im_info.data.resize_(data[1].size()).copy_(data[1]) gt_boxes.data.resize_(data[2].size()).copy_(data[2]) num_boxes.data.resize_(data[3].size()).copy_(data[3]) det_tic = time.time() rois, cls_prob, bbox_pred, \ rpn_loss_cls, rpn_loss_box, \ RCNN_loss_cls, RCNN_loss_bbox, \ rois_label = fasterRCNN(im_data, im_info, gt_boxes, num_boxes) scores = cls_prob.data boxes = rois.data[:, :, 1:5] if cfg.TEST.BBOX_REG: # Apply bounding-box regression deltas box_deltas = bbox_pred.data if cfg.TRAIN.BBOX_NORMALIZE_TARGETS_PRECOMPUTED: # Optionally normalize targets by a precomputed mean and stdev if args.class_agnostic: box_deltas = box_deltas.view(-1, 4) * torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_STDS).cuda() \ + torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_MEANS).cuda() box_deltas = box_deltas.view(1, -1, 4) else: box_deltas = box_deltas.view(-1, 4) * torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_STDS).cuda() \ + torch.FloatTensor(cfg.TRAIN.BBOX_NORMALIZE_MEANS).cuda() box_deltas = box_deltas.view(1, -1, 4 * len(imdb.classes)) pred_boxes = bbox_transform_inv(boxes, box_deltas, 1) pred_boxes = clip_boxes(pred_boxes, im_info.data, 1) else: # Simply repeat the boxes, once for each class _ = torch.from_numpy(np.tile(boxes, (1, scores.shape[1]))) pred_boxes = _.cuda() if args.cuda > 0 else _ pred_boxes /= data[1][0][2].item() scores = scores.squeeze() pred_boxes = pred_boxes.squeeze() det_toc = time.time() detect_time = det_toc - det_tic misc_tic = time.time() if vis: im = cv2.imread(imdb.image_path_at(i)) im2show = np.copy(im) for j in xrange(1, imdb.num_classes): inds = torch.nonzero(scores[:,j]>thresh).view(-1) # if there is det if inds.numel() > 0: cls_scores = scores[:,j][inds] _, order = torch.sort(cls_scores, 0, True) if args.class_agnostic: cls_boxes = pred_boxes[inds, :] else: cls_boxes = pred_boxes[inds][:, j * 4:(j + 1) * 4] cls_dets = torch.cat((cls_boxes, cls_scores.unsqueeze(1)), 1) # cls_dets = torch.cat((cls_boxes, cls_scores), 1) cls_dets = cls_dets[order] keep = nms(cls_dets, cfg.TEST.NMS) cls_dets = cls_dets[keep.view(-1).long()] if vis: im2show = vis_detections(im2show, imdb.classes[j], cls_dets.cpu().numpy(), 0.3) all_boxes[j][i] = cls_dets.cpu().numpy() else: all_boxes[j][i] = empty_array # Limit to max_per_image detections *over all classes* if max_per_image > 0: image_scores = np.hstack([all_boxes[j][i][:, -1] for j in xrange(1, imdb.num_classes)]) if len(image_scores) > max_per_image: image_thresh = np.sort(image_scores)[-max_per_image] for j in xrange(1, imdb.num_classes): keep = np.where(all_boxes[j][i][:, -1] >= image_thresh)[0] all_boxes[j][i] = all_boxes[j][i][keep, :] misc_toc = time.time() nms_time = misc_toc - misc_tic sys.stdout.write('im_detect: {:d}/{:d} {:.3f}s {:.3f}s \r' \ .format(i + 1, num_images, detect_time, nms_time)) sys.stdout.flush() if vis: cv2.imwrite('result.png', im2show) pdb.set_trace() #cv2.imshow('test', im2show) #cv2.waitKey(0) with open(det_file, 'wb') as f: pickle.dump(all_boxes, f, pickle.HIGHEST_PROTOCOL) print('Evaluating detections') imdb.evaluate_detections(all_boxes, output_dir) end = time.time() print("test time: %0.4fs" % (end - start)) ================================================ FILE: trainval_net.py ================================================ # -------------------------------------------------------- # Pytorch multi-GPU Faster R-CNN # Licensed under The MIT License [see LICENSE for details] # Written by Jiasen Lu, Jianwei Yang, based on code from Ross Girshick # -------------------------------------------------------- from __future__ import absolute_import from __future__ import division from __future__ import print_function import _init_paths import os import sys import numpy as np import argparse import pprint import pdb import time import torch from torch.autograd import Variable import torch.nn as nn import torch.optim as optim import torchvision.transforms as transforms from torch.utils.data.sampler import Sampler from roi_data_layer.roidb import combined_roidb from roi_data_layer.roibatchLoader import roibatchLoader from model.utils.config import cfg, cfg_from_file, cfg_from_list, get_output_dir from model.utils.net_utils import weights_normal_init, save_net, load_net, \ adjust_learning_rate, save_checkpoint, clip_gradient from model.faster_rcnn.vgg16 import vgg16 from model.faster_rcnn.resnet import resnet def parse_args(): """ Parse input arguments """ parser = argparse.ArgumentParser(description='Train a Fast R-CNN network') parser.add_argument('--dataset', dest='dataset', help='training dataset', default='pascal_voc', type=str) parser.add_argument('--net', dest='net', help='vgg16, res101', default='vgg16', type=str) parser.add_argument('--start_epoch', dest='start_epoch', help='starting epoch', default=1, type=int) parser.add_argument('--epochs', dest='max_epochs', help='number of epochs to train', default=20, type=int) parser.add_argument('--disp_interval', dest='disp_interval', help='number of iterations to display', default=100, type=int) parser.add_argument('--checkpoint_interval', dest='checkpoint_interval', help='number of iterations to display', default=10000, type=int) parser.add_argument('--save_dir', dest='save_dir', help='directory to save models', default="models", type=str) parser.add_argument('--nw', dest='num_workers', help='number of workers to load data', default=0, type=int) parser.add_argument('--cuda', dest='cuda', help='whether use CUDA', action='store_true') parser.add_argument('--ls', dest='large_scale', help='whether use large imag scale', action='store_true') parser.add_argument('--mGPUs', dest='mGPUs', help='whether use multiple GPUs', action='store_true') parser.add_argument('--bs', dest='batch_size', help='batch_size', default=1, type=int) parser.add_argument('--cag', dest='class_agnostic', help='whether to perform class_agnostic bbox regression', action='store_true') # config optimization parser.add_argument('--o', dest='optimizer', help='training optimizer', default="sgd", type=str) parser.add_argument('--lr', dest='lr', help='starting learning rate', default=0.001, type=float) parser.add_argument('--lr_decay_step', dest='lr_decay_step', help='step to do learning rate decay, unit is epoch', default=5, type=int) parser.add_argument('--lr_decay_gamma', dest='lr_decay_gamma', help='learning rate decay ratio', default=0.1, type=float) # set training session parser.add_argument('--s', dest='session', help='training session', default=1, type=int) # resume trained model parser.add_argument('--r', dest='resume', help='resume checkpoint or not', default=False, type=bool) parser.add_argument('--checksession', dest='checksession', help='checksession to load model', default=1, type=int) parser.add_argument('--checkepoch', dest='checkepoch', help='checkepoch to load model', default=1, type=int) parser.add_argument('--checkpoint', dest='checkpoint', help='checkpoint to load model', default=0, type=int) # log and display parser.add_argument('--use_tfb', dest='use_tfboard', help='whether use tensorboard', action='store_true') args = parser.parse_args() return args class sampler(Sampler): def __init__(self, train_size, batch_size): self.num_data = train_size self.num_per_batch = int(train_size / batch_size) self.batch_size = batch_size self.range = torch.arange(0,batch_size).view(1, batch_size).long() self.leftover_flag = False if train_size % batch_size: self.leftover = torch.arange(self.num_per_batch*batch_size, train_size).long() self.leftover_flag = True def __iter__(self): rand_num = torch.randperm(self.num_per_batch).view(-1,1) * self.batch_size self.rand_num = rand_num.expand(self.num_per_batch, self.batch_size) + self.range self.rand_num_view = self.rand_num.view(-1) if self.leftover_flag: self.rand_num_view = torch.cat((self.rand_num_view, self.leftover),0) return iter(self.rand_num_view) def __len__(self): return self.num_data if __name__ == '__main__': args = parse_args() print('Called with args:') print(args) if args.dataset == "pascal_voc": args.imdb_name = "voc_2007_trainval" args.imdbval_name = "voc_2007_test" args.set_cfgs = ['ANCHOR_SCALES', '[8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]', 'MAX_NUM_GT_BOXES', '20'] elif args.dataset == "pascal_voc_0712": args.imdb_name = "voc_2007_trainval+voc_2012_trainval" args.imdbval_name = "voc_2007_test" args.set_cfgs = ['ANCHOR_SCALES', '[8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]', 'MAX_NUM_GT_BOXES', '20'] elif args.dataset == "coco": args.imdb_name = "coco_2014_train+coco_2014_valminusminival" args.imdbval_name = "coco_2014_minival" args.set_cfgs = ['ANCHOR_SCALES', '[4, 8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]', 'MAX_NUM_GT_BOXES', '50'] elif args.dataset == "imagenet": args.imdb_name = "imagenet_train" args.imdbval_name = "imagenet_val" args.set_cfgs = ['ANCHOR_SCALES', '[4, 8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]', 'MAX_NUM_GT_BOXES', '30'] elif args.dataset == "vg": # train sizes: train, smalltrain, minitrain # train scale: ['150-50-20', '150-50-50', '500-150-80', '750-250-150', '1750-700-450', '1600-400-20'] args.imdb_name = "vg_150-50-50_minitrain" args.imdbval_name = "vg_150-50-50_minival" args.set_cfgs = ['ANCHOR_SCALES', '[4, 8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]', 'MAX_NUM_GT_BOXES', '50'] args.cfg_file = "cfgs/{}_ls.yml".format(args.net) if args.large_scale else "cfgs/{}.yml".format(args.net) if args.cfg_file is not None: cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) print('Using config:') pprint.pprint(cfg) np.random.seed(cfg.RNG_SEED) #torch.backends.cudnn.benchmark = True if torch.cuda.is_available() and not args.cuda: print("WARNING: You have a CUDA device, so you should probably run with --cuda") # train set # -- Note: Use validation set and disable the flipped to enable faster loading. cfg.TRAIN.USE_FLIPPED = True cfg.USE_GPU_NMS = args.cuda imdb, roidb, ratio_list, ratio_index = combined_roidb(args.imdb_name) train_size = len(roidb) print('{:d} roidb entries'.format(len(roidb))) output_dir = args.save_dir + "/" + args.net + "/" + args.dataset if not os.path.exists(output_dir): os.makedirs(output_dir) sampler_batch = sampler(train_size, args.batch_size) dataset = roibatchLoader(roidb, ratio_list, ratio_index, args.batch_size, \ imdb.num_classes, training=True) dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, sampler=sampler_batch, num_workers=args.num_workers) # initilize the tensor holder here. im_data = torch.FloatTensor(1) im_info = torch.FloatTensor(1) num_boxes = torch.LongTensor(1) gt_boxes = torch.FloatTensor(1) # ship to cuda if args.cuda: im_data = im_data.cuda() im_info = im_info.cuda() num_boxes = num_boxes.cuda() gt_boxes = gt_boxes.cuda() # make variable im_data = Variable(im_data) im_info = Variable(im_info) num_boxes = Variable(num_boxes) gt_boxes = Variable(gt_boxes) if args.cuda: cfg.CUDA = True # initilize the network here. if args.net == 'vgg16': fasterRCNN = vgg16(imdb.classes, pretrained=True, class_agnostic=args.class_agnostic) elif args.net == 'res101': fasterRCNN = resnet(imdb.classes, 101, pretrained=True, class_agnostic=args.class_agnostic) elif args.net == 'res50': fasterRCNN = resnet(imdb.classes, 50, pretrained=True, class_agnostic=args.class_agnostic) elif args.net == 'res152': fasterRCNN = resnet(imdb.classes, 152, pretrained=True, class_agnostic=args.class_agnostic) else: print("network is not defined") pdb.set_trace() fasterRCNN.create_architecture() lr = cfg.TRAIN.LEARNING_RATE lr = args.lr #tr_momentum = cfg.TRAIN.MOMENTUM #tr_momentum = args.momentum params = [] for key, value in dict(fasterRCNN.named_parameters()).items(): if value.requires_grad: if 'bias' in key: params += [{'params':[value],'lr':lr*(cfg.TRAIN.DOUBLE_BIAS + 1), \ 'weight_decay': cfg.TRAIN.BIAS_DECAY and cfg.TRAIN.WEIGHT_DECAY or 0}] else: params += [{'params':[value],'lr':lr, 'weight_decay': cfg.TRAIN.WEIGHT_DECAY}] if args.optimizer == "adam": lr = lr * 0.1 optimizer = torch.optim.Adam(params) elif args.optimizer == "sgd": optimizer = torch.optim.SGD(params, momentum=cfg.TRAIN.MOMENTUM) if args.cuda: fasterRCNN.cuda() if args.resume: load_name = os.path.join(output_dir, 'faster_rcnn_{}_{}_{}.pth'.format(args.checksession, args.checkepoch, args.checkpoint)) print("loading checkpoint %s" % (load_name)) checkpoint = torch.load(load_name) args.session = checkpoint['session'] args.start_epoch = checkpoint['epoch'] fasterRCNN.load_state_dict(checkpoint['model']) optimizer.load_state_dict(checkpoint['optimizer']) lr = optimizer.param_groups[0]['lr'] if 'pooling_mode' in checkpoint.keys(): cfg.POOLING_MODE = checkpoint['pooling_mode'] print("loaded checkpoint %s" % (load_name)) if args.mGPUs: fasterRCNN = nn.DataParallel(fasterRCNN) iters_per_epoch = int(train_size / args.batch_size) if args.use_tfboard: from tensorboardX import SummaryWriter logger = SummaryWriter("logs") for epoch in range(args.start_epoch, args.max_epochs + 1): # setting to train mode fasterRCNN.train() loss_temp = 0 start = time.time() if epoch % (args.lr_decay_step + 1) == 0: adjust_learning_rate(optimizer, args.lr_decay_gamma) lr *= args.lr_decay_gamma data_iter = iter(dataloader) for step in range(iters_per_epoch): data = next(data_iter) im_data.data.resize_(data[0].size()).copy_(data[0]) im_info.data.resize_(data[1].size()).copy_(data[1]) gt_boxes.data.resize_(data[2].size()).copy_(data[2]) num_boxes.data.resize_(data[3].size()).copy_(data[3]) fasterRCNN.zero_grad() rois, cls_prob, bbox_pred, \ rpn_loss_cls, rpn_loss_box, \ RCNN_loss_cls, RCNN_loss_bbox, \ rois_label = fasterRCNN(im_data, im_info, gt_boxes, num_boxes) loss = rpn_loss_cls.mean() + rpn_loss_box.mean() \ + RCNN_loss_cls.mean() + RCNN_loss_bbox.mean() loss_temp += loss.item() # backward optimizer.zero_grad() loss.backward() if args.net == "vgg16": clip_gradient(fasterRCNN, 10.) optimizer.step() if step % args.disp_interval == 0: end = time.time() if step > 0: loss_temp /= (args.disp_interval + 1) if args.mGPUs: loss_rpn_cls = rpn_loss_cls.mean().item() loss_rpn_box = rpn_loss_box.mean().item() loss_rcnn_cls = RCNN_loss_cls.mean().item() loss_rcnn_box = RCNN_loss_bbox.mean().item() fg_cnt = torch.sum(rois_label.data.ne(0)) bg_cnt = rois_label.data.numel() - fg_cnt else: loss_rpn_cls = rpn_loss_cls.item() loss_rpn_box = rpn_loss_box.item() loss_rcnn_cls = RCNN_loss_cls.item() loss_rcnn_box = RCNN_loss_bbox.item() fg_cnt = torch.sum(rois_label.data.ne(0)) bg_cnt = rois_label.data.numel() - fg_cnt print("[session %d][epoch %2d][iter %4d/%4d] loss: %.4f, lr: %.2e" \ % (args.session, epoch, step, iters_per_epoch, loss_temp, lr)) print("\t\t\tfg/bg=(%d/%d), time cost: %f" % (fg_cnt, bg_cnt, end-start)) print("\t\t\trpn_cls: %.4f, rpn_box: %.4f, rcnn_cls: %.4f, rcnn_box %.4f" \ % (loss_rpn_cls, loss_rpn_box, loss_rcnn_cls, loss_rcnn_box)) if args.use_tfboard: info = { 'loss': loss_temp, 'loss_rpn_cls': loss_rpn_cls, 'loss_rpn_box': loss_rpn_box, 'loss_rcnn_cls': loss_rcnn_cls, 'loss_rcnn_box': loss_rcnn_box } logger.add_scalars("logs_s_{}/losses".format(args.session), info, (epoch - 1) * iters_per_epoch + step) loss_temp = 0 start = time.time() save_name = os.path.join(output_dir, 'faster_rcnn_{}_{}_{}.pth'.format(args.session, epoch, step)) save_checkpoint({ 'session': args.session, 'epoch': epoch + 1, 'model': fasterRCNN.module.state_dict() if args.mGPUs else fasterRCNN.state_dict(), 'optimizer': optimizer.state_dict(), 'pooling_mode': cfg.POOLING_MODE, 'class_agnostic': args.class_agnostic, }, save_name) print('save model: {}'.format(save_name)) if args.use_tfboard: logger.close()