diff --git a/.gitignore b/.gitignore new file mode 100755 index 0000000..fa7a7fe --- /dev/null +++ b/.gitignore @@ -0,0 +1,12 @@ +*DS_Store* +tools/vis* +cache +data +data/ +cache/ +results/ + +*.swp + +*.pyc +*.o* diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..247ff3c --- /dev/null +++ b/LICENSE @@ -0,0 +1,29 @@ +BSD 3-Clause License + +Copyright (c) 2018, University of Michigan +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +* Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +* 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. + +* Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +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 HOLDER 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. diff --git a/README.md b/README.md index 119427c..a04f7fe 100644 --- a/README.md +++ b/README.md @@ -1 +1,164 @@ -# ExtremeNet \ No newline at end of file +# ExtremeNet: Training and Evaluation Code +Code for **bottom-up** object detection by grouping extreme and center points: +![](readme/teaser.png) +> [**Bottom-up Object Detection by Grouping Extreme and Center Points**](https://arxiv.org/abs/xxxx.xxxxx), +> Xingyi Zhou, Jiacheng Zhuo, Philipp Krähenbühl, +> *arXiv technical report* + +This project is developed upon the [CornerNet code](https://github.com/princeton-vl/CornerNet) and contains the code from [Deep Extreme Cut](https://github.com/scaelles/DEXTR-PyTorch). Thanks to the original authors! + +Contact: [zhouxy2017@gmail.com](mailto:zhouxy2017@gmail.com). Any questions or discussions are welcomed! + +## Abstract + +With the advent of deep learning, object detection drifted from a bottom-up to a top-down recognition problem. State of the art algorithms enumerate a near-exhaustive list of object locations and classify each into: object or not. In this paper, we show that bottom-up approaches still perform competitively. We detect four extreme points (top-most, left-most, bottom-most, right-most) and one center point of objects using a standard keypoint estimation network. We group the five keypoints into a bounding box if they are geometrically aligned. Object detection is then a purely appearance-based keypoint estimation problem, without region classification or implicit feature learning. The proposed method performs on-par with the state-of-the-art region based detection methods, with a bounding box AP of 43.2% on COCO test-dev. In addition, our estimated extreme points directly span a coarse octagonal mask, with a COCO Mask AP of 18.9%, much better than the Mask AP of vanilla bounding boxes. Extreme point guided segmentation further improves this to 34.1% Mask AP. + +## Installation +The code was tested with [Anaconda](https://www.anaconda.com/download) Python 3.6 and [PyTorch]((http://pytorch.org/)) v0.4.1. After install Anaconda: + +1. Clone this repo: + + ~~~ + ExtremeNet_ROOT=/path/to/clone/ExtremeNet + git clone --recursive https://github.com/xingyizhou/ExtremeNet ExtremeNet_ROOT + ~~~ + + +2. Create an Anaconda environment using the provided package list from [Cornernet](https://github.com/princeton-vl/CornerNet). + + ~~~ + conda create --name CornerNet --file conda_packagelist.txt + source activate CornerNet + ~~~ + +3. Compiling NMS (originally from [Faster R-CNN](https://github.com/rbgirshick/py-faster-rcnn/blob/master/lib/nms/cpu_nms.pyx) and [Soft-NMS](https://github.com/bharatsingh430/soft-nms/blob/master/lib/nms/cpu_nms.pyx)). + + ~~~ + cd ExtremeNet_ROOT/external + make + ~~~ + +## Demo + +- Download our [pre-trained model](https://drive.google.com/file/d/1re-A74WRvuhE528X6sWsg1eEbMG8dmE4/view?usp=sharing) and put it in `cache/`. +- Optionally, if you want to test instance segmentation with [Deep Extreme Cut](https://github.com/scaelles/DEXTR-PyTorch), download their [PASCAL + SBD pertained model](https://data.vision.ee.ethz.ch/kmaninis/share/DEXTR/Downloads/models/dextr_pascal-sbd.pth) and put it in `cache/`. +- Run the demo + + ~~~ + Python demo.py ExtremeNet [--demo /path/to/image/or/folder] [--show_mask] + ~~~ + + Contents in `[]` are optional. By default, it runs the sample images provided in `ExtremeNet_ROOT/images/` (from [Detectron](https://github.com/facebookresearch/Detectron/tree/master/demo)). We show the predicted extreme point heatmaps (combined four heatmaps and overlaid on the input image), the predicted center point heatmap, and the detection and octagon mask results. If setup correctly, the output will look like: + + + +

+ +

+ + If `--show_mask` is turned on, it further pipelined with [Deep Extreme Cut](https://github.com/scaelles/DEXTR-PyTorch) for instance segmentation. The output will look like: +

+ +

+ + +## Data preparation + +If you want to reproduce the results in the paper for benchmark evaluation and training, you will need to setup dataset. + +### Installing MS COCO APIs +~~~ +cd ExtremeNet_ROOT/data +git clone https://github.com/cocodataset/cocoapi.git coco +cd ExtremeNet_ROOT/data/coco/PythonAPI +Make +python setup.up install --user +~~~ + +### Downloading MS COCO Data +- Download the images (2017 Train, 2017 Val, 2017 Test) from [coco website](http://cocodataset.org/#download). +- Download annotation files (2017 train/val and test image info) from [coco website](http://cocodataset.org/#download). +- Place the data (or create symlinks) to make the data folder like: + + ~~~ + ${ExtremeNet_ROOT} + |-- data + `-- |-- coco + `-- |-- annotations + | |-- instances_train2017.json + | |-- instances_val2017.json + | |-- image_info_test-dev2017.json + `-- images + |-- train2017 + |-- val2017 + |-- test2017 + ~~~ +- Generate extreme point annotation from segmentation: + + ~~~ + cd ExtremeNet_ROOT/tools/ + python gen_coco_extreme_points.py + ~~~ + It generates `instances_extreme_train2017.json` and `instances_extreme_val2017.json` in `data/coco/annotations/`. + +## Benchmark Evaluation + +After downloading our pre-trained model and the dataset, + +- Run the following command to evaluate object detection: + + ~~~ + python test.py ExtremeNet [--suffix multi_scale] + ~~~ + The results on COCO validation set should be `40.3` box AP without `--suffix multi_scale` and `43.3` box AP with `--suffix multi_scale`. + +- After obtaining the detection results, run the following commands for instance segmentation: + + ~~~ + python eval_dextr_mask.py results/ExtremeNet/250000/validation/multi_scale/results.json + ~~~ + The results on COCO validation set should be `34.6` mask AP (The evaluation will be slow). + +- You can test with other hyper-parameters by creating a new config file (`ExtremeNet-.json`) in `config/`. + + +## Training + +You will need 5 12GB GPUs to reproduce our training. Our model is fine-tuned on the 10-GPU pre-trained [CornerNet model](https://drive.google.com/file/d/1UHjVzSG27Ms0VfSFeGYJ2h2AYZ6d4Le_/view?usp=sharing) (training from scratch may result in 2 AP lower based on our preliminary experiments). After downloading the CornerNet model and put it in `cache/`, run + +~~~ +python train.py ExtremeNet +~~~ + +You can resume a half-trained model by + +~~~ +python train.py ExtremeNet --iter xxxx +~~~ + +## Citation +If you find this model useful for your resesarch, please use the following BibTeX entry. + + @inproceedings{zhou2019bottomup, + title={Bottom-up Object Detection by Grouping Extreme and Center Points}, + author={Zhou, Xingyi and Zhuo, Jiacheng and Kr{\"a}henb{\"u}hl, Philipp}, + booktitle={arXiv preprint arXiv:xxxx.xxxxx}, + year={2019} + } + +Please also considering citing the CornerNet paper (where this code is heavily borrowed from) and Deep Extreme Cut paper (if you use the instance segmentation part). + + @inproceedings{law2018cornernet, + title={CornerNet: Detecting Objects as Paired Keypoints}, + author={Law, Hei and Deng, Jia}, + booktitle={Proceedings of the European Conference on Computer Vision (ECCV)}, + pages={734--750}, + year={2018} + } + + @Inproceedings{Man+18, + Title = {Deep Extreme Cut: From Extreme Points to Object Segmentation}, + Author = {K.K. Maninis and S. Caelles and J. Pont-Tuset and L. {Van Gool}}, + Booktitle = {Computer Vision and Pattern Recognition (CVPR)}, + Year = {2018} + } \ No newline at end of file diff --git a/conda_packagelist.txt b/conda_packagelist.txt new file mode 100644 index 0000000..bd9fb12 --- /dev/null +++ b/conda_packagelist.txt @@ -0,0 +1,88 @@ +# This file may be used to create an environment using: +# $ conda create --name --file +# platform: linux-64 +@EXPLICIT +https://repo.continuum.io/pkgs/main/linux-64/blas-1.0-mkl.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/bzip2-1.0.6-h9a117a8_4.tar.bz2 +https://conda.anaconda.org/conda-forge/linux-64/ca-certificates-2018.4.16-0.tar.bz2 +https://conda.anaconda.org/caffe2/linux-64/caffe2-cuda8.0-cudnn7-0.8.dev-py36_2018.05.14.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/cairo-1.14.12-h7636065_2.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/certifi-2018.4.16-py36_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/cffi-1.11.5-py36h9745a5d_0.tar.bz2 +https://repo.continuum.io/pkgs/free/linux-64/cudatoolkit-8.0-3.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/cycler-0.10.0-py36h93f1223_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/dbus-1.13.2-h714fa37_1.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/expat-2.2.5-he0dffb1_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/ffmpeg-3.4-h7264315_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/fontconfig-2.12.6-h49f89f6_0.tar.bz2 +https://repo.continuum.io/pkgs/free/linux-64/freeglut-2.8.1-0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/freetype-2.8-hab7d2ae_1.tar.bz2 +https://repo.continuum.io/pkgs/free/linux-64/future-0.16.0-py36_1.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/gflags-2.2.1-hf484d3e_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/glib-2.56.1-h000015b_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/glog-0.3.5-hf484d3e_1.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/graphite2-1.3.11-hf63cedd_1.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/gst-plugins-base-1.14.0-hbbd80ab_1.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/gstreamer-1.14.0-hb453b48_1.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/h5py-2.8.0-py36hca9c191_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/harfbuzz-1.7.6-h5f0a787_1.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/hdf5-1.8.18-h6792536_1.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/icu-58.2-h9c2bf20_1.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/intel-openmp-2018.0.0-8.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/jasper-2.0.14-h07fcdf6_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/jpeg-9b-h024ee3a_2.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/kiwisolver-1.0.1-py36h764f252_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/libedit-3.1-heed3624_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/libffi-3.2.1-hd88cf55_4.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/libgcc-ng-7.2.0-hdf63c60_3.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/libgfortran-ng-7.2.0-hdf63c60_3.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/libglu-9.0.0-h0c0bdc1_1.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/libopus-1.2.1-hb9ed12e_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/libpng-1.6.34-hb9fc6fc_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/libprotobuf-3.5.2-h6f1eeef_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/libstdcxx-ng-7.2.0-hdf63c60_3.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/libtiff-4.0.9-h28f6b97_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/libvpx-1.6.1-h888fd40_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/libxcb-1.13-h1bed415_1.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/libxml2-2.9.8-hf84eae3_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/matplotlib-2.2.2-py36h0e671d2_1.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/mkl-2018.0.2-1.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/mkl_fft-1.0.1-py36h3010b51_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/mkl_random-1.0.1-py36h629b387_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/ncurses-6.0-h9df7e31_2.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/ninja-1.8.2-py36h6bb024c_1.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/numpy-1.14.3-py36hcd700cb_1.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/numpy-base-1.14.3-py36h9be14a7_1.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/olefile-0.45.1-py36_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/opencv-3.3.1-py36h9248ab4_2.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/openssl-1.0.2o-h20670df_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/pcre-8.42-h439df22_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/pillow-5.1.0-py36h3deb7b8_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/pip-10.0.1-py36_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/pixman-0.34.0-hceecf20_3.tar.bz2 +https://conda.anaconda.org/conda-forge/linux-64/protobuf-3.5.2-py36_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/pycparser-2.18-py36hf9f622e_1.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/pyparsing-2.2.0-py36hee85983_1.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/pyqt-5.9.2-py36h751905a_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/python-3.6.5-hc3d631a_2.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/python-dateutil-2.7.2-py36_0.tar.bz2 +https://conda.anaconda.org/pytorch/linux-64/pytorch-0.4.1-py36_cuda9.0.176_cudnn7.1.2_1.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/pytz-2018.4-py36_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/pyyaml-3.12-py36hafb9ca4_1.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/qt-5.9.5-h7e424d6_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/readline-7.0-ha6073c6_4.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/scikit-learn-0.19.1-py36h7aa7ec6_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/scipy-1.1.0-py36hfc37229_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/setuptools-39.1.0-py36_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/sip-4.19.8-py36hf484d3e_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/six-1.11.0-py36h372c433_1.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/sqlite-3.23.1-he433501_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/tk-8.6.7-hc745277_3.tar.bz2 +https://conda.anaconda.org/pytorch/linux-64/torchvision-0.2.1-py36_1.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/tornado-5.0.2-py36_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/tqdm-4.23.0-py36_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/wheel-0.31.0-py36_0.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/xz-5.2.3-h5e939de_4.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/yaml-0.1.7-had09818_2.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/zlib-1.2.11-ha838bed_2.tar.bz2 +https://repo.continuum.io/pkgs/main/linux-64/progress-1.4-py36_0.tar.bz2 diff --git a/config.py b/config.py new file mode 100755 index 0000000..e6567aa --- /dev/null +++ b/config.py @@ -0,0 +1,179 @@ +import os +import numpy as np + +class Config: + def __init__(self): + self._configs = {} + self._configs["dataset"] = None + self._configs["sampling_function"] = "kp_detection" + + # Training Config + self._configs["display"] = 5 + self._configs["snapshot"] = 5000 + self._configs["stepsize"] = 450000 + self._configs["learning_rate"] = 0.00025 + self._configs["decay_rate"] = 10 + self._configs["max_iter"] = 500000 + self._configs["val_iter"] = 100 + self._configs["batch_size"] = 1 + self._configs["snapshot_name"] = None + self._configs["prefetch_size"] = 100 + self._configs["weight_decay"] = False + self._configs["weight_decay_rate"] = 1e-5 + self._configs["weight_decay_type"] = "l2" + self._configs["pretrain"] = None + self._configs["opt_algo"] = "adam" + self._configs["chunk_sizes"] = None + + # Directories + self._configs["data_dir"] = "./data" + self._configs["cache_dir"] = "./cache" + self._configs["config_dir"] = "./config" + self._configs["result_dir"] = "./results" + + # Split + self._configs["train_split"] = "trainval" + self._configs["val_split"] = "minival" + self._configs["test_split"] = "testdev" + + # Rng + self._configs["data_rng"] = np.random.RandomState(123) + self._configs["nnet_rng"] = np.random.RandomState(317) + + @property + def chunk_sizes(self): + return self._configs["chunk_sizes"] + + @property + def train_split(self): + return self._configs["train_split"] + + @property + def val_split(self): + return self._configs["val_split"] + + @property + def test_split(self): + return self._configs["test_split"] + + @property + def full(self): + return self._configs + + @property + def sampling_function(self): + return self._configs["sampling_function"] + + @property + def data_rng(self): + return self._configs["data_rng"] + + @property + def nnet_rng(self): + return self._configs["nnet_rng"] + + @property + def opt_algo(self): + return self._configs["opt_algo"] + + @property + def weight_decay_type(self): + return self._configs["weight_decay_type"] + + @property + def prefetch_size(self): + return self._configs["prefetch_size"] + + @property + def pretrain(self): + return self._configs["pretrain"] + + @property + def weight_decay_rate(self): + return self._configs["weight_decay_rate"] + + @property + def weight_decay(self): + return self._configs["weight_decay"] + + @property + def result_dir(self): + result_dir = os.path.join(self._configs["result_dir"], self.snapshot_name) + if not os.path.exists(result_dir): + os.makedirs(result_dir) + return result_dir + + @property + def dataset(self): + return self._configs["dataset"] + + @property + def snapshot_name(self): + return self._configs["snapshot_name"] + + @property + def snapshot_dir(self): + snapshot_dir = os.path.join(self.cache_dir, "nnet", self.snapshot_name) + + if not os.path.exists(snapshot_dir): + os.makedirs(snapshot_dir) + + return snapshot_dir + + @property + def snapshot_file(self): + snapshot_file = os.path.join(self.snapshot_dir, self.snapshot_name + "_{}.pkl") + return snapshot_file + + @property + def config_dir(self): + return self._configs["config_dir"] + + @property + def batch_size(self): + return self._configs["batch_size"] + + @property + def max_iter(self): + return self._configs["max_iter"] + + @property + def learning_rate(self): + return self._configs["learning_rate"] + + @property + def decay_rate(self): + return self._configs["decay_rate"] + + @property + def stepsize(self): + return self._configs["stepsize"] + + @property + def snapshot(self): + return self._configs["snapshot"] + + @property + def display(self): + return self._configs["display"] + + @property + def val_iter(self): + return self._configs["val_iter"] + + @property + def data_dir(self): + return self._configs["data_dir"] + + @property + def cache_dir(self): + if not os.path.exists(self._configs["cache_dir"]): + os.makedirs(self._configs["cache_dir"]) + return self._configs["cache_dir"] + + def update_config(self, new): + for key in new: + if key in self._configs: + self._configs[key] = new[key] + +system_configs = Config() diff --git a/config/CornerNet-multi_scale.json b/config/CornerNet-multi_scale.json new file mode 100755 index 0000000..c23a13d --- /dev/null +++ b/config/CornerNet-multi_scale.json @@ -0,0 +1,54 @@ +{ + "system": { + "dataset": "MSCOCO", + "batch_size": 49, + "sampling_function": "kp_detection", + + "train_split": "trainval", + "val_split": "minival", + + "learning_rate": 0.00025, + "decay_rate": 10, + + "val_iter": 100, + + "opt_algo": "adam", + "prefetch_size": 5, + + "max_iter": 500000, + "stepsize": 450000, + "snapshot": 5000, + + "chunk_sizes": [4, 5, 5, 5, 5, 5, 5, 5, 5, 5], + + "data_dir": "./data" + }, + + "db": { + "rand_scale_min": 0.6, + "rand_scale_max": 1.4, + "rand_scale_step": 0.1, + "rand_scales": null, + + "rand_crop": true, + "rand_color": true, + + "border": 128, + "gaussian_bump": true, + + "input_size": [511, 511], + "output_sizes": [[128, 128]], + + "test_scales": [0.5, 0.75, 1, 1.25, 1.5], + + "top_k": 100, + "categories": 80, + "ae_threshold": 0.5, + "nms_threshold": 0.5, + + "merge_bbox": true, + "weight_exp": 10, + + "max_per_image": 100 + } +} diff --git a/config/CornerNet.json b/config/CornerNet.json new file mode 100755 index 0000000..bf24265 --- /dev/null +++ b/config/CornerNet.json @@ -0,0 +1,51 @@ +{ + "system": { + "dataset": "MSCOCO", + "batch_size": 49, + "sampling_function": "kp_detection", + + "train_split": "trainval", + "val_split": "minival", + + "learning_rate": 0.00025, + "decay_rate": 10, + + "val_iter": 100, + + "opt_algo": "adam", + "prefetch_size": 5, + + "max_iter": 500000, + "stepsize": 450000, + "snapshot": 5000, + + "chunk_sizes": [4, 5, 5, 5, 5, 5, 5, 5, 5, 5], + + "data_dir": "./data" + }, + + "db": { + "rand_scale_min": 0.6, + "rand_scale_max": 1.4, + "rand_scale_step": 0.1, + "rand_scales": null, + + "rand_crop": true, + "rand_color": true, + + "border": 128, + "gaussian_bump": true, + + "input_size": [511, 511], + "output_sizes": [[128, 128]], + + "test_scales": [1], + + "top_k": 100, + "categories": 80, + "ae_threshold": 0.5, + "nms_threshold": 0.5, + + "max_per_image": 100 + } +} diff --git a/config/ExtremeNet-multi_scale.json b/config/ExtremeNet-multi_scale.json new file mode 100755 index 0000000..f755468 --- /dev/null +++ b/config/ExtremeNet-multi_scale.json @@ -0,0 +1,55 @@ +{ + "system": { + "dataset": "MSCOCOExtreme", + "batch_size": 24, + "sampling_function": "kp_detection", + + "train_split": "train", + "val_split": "val", + + "learning_rate": 0.00025, + "decay_rate": 10, + + "val_iter": 100, + + "opt_algo": "adam", + "prefetch_size": 5, + + "max_iter": 250000, + "stepsize": 200000, + "snapshot": 50000, + + "chunk_sizes": [4, 5, 5, 5, 5], + + "data_dir": "./data", + + "pretrain": "./cache/CornerNet_500000.pkl" + }, + + "db": { + "rand_scale_min": 0.6, + "rand_scale_max": 1.4, + "rand_scale_step": 0.1, + "rand_scales": null, + + "rand_crop": true, + "rand_color": true, + + "border": 128, + "gaussian_bump": true, + + "input_size": [511, 511], + "output_sizes": [[128, 128]], + + "test_scales": [0.5, 0.75, 1, 1.25, 1.5], + + "top_k": 40, + "categories": 80, + "aggr_weight": 0.1, + "scores_thresh": 0.1, + "center_thresh": 0.1, + "nms_threshold": 0.5, + + "max_per_image": 100 + } +} diff --git a/config/ExtremeNet.json b/config/ExtremeNet.json new file mode 100755 index 0000000..75715f1 --- /dev/null +++ b/config/ExtremeNet.json @@ -0,0 +1,56 @@ +{ + "system": { + "dataset": "MSCOCOExtreme", + "batch_size": 38, + "sampling_function": "kp_detection", + + "train_split": "train", + "val_split": "val", + + "learning_rate": 0.00025, + "decay_rate": 10, + + "val_iter": 100, + + "opt_algo": "adam", + "prefetch_size": 10, + + "max_iter": 250000, + "stepsize": 200000, + "snapshot": 50000, + + "chunk_sizes": [3, 5, 5, 5, 5, 5, 5, 5], + + "data_dir": "./data", + + "pretrain": "./cache/CornerNet_500000.pkl" + }, + + "db": { + "rand_scale_min": 0.6, + "rand_scale_max": 1.4, + "rand_scale_step": 0.1, + "rand_scales": null, + + "rand_crop": true, + "rand_color": true, + + "border": 128, + "gaussian_bump": true, + + "input_size": [511, 511], + "output_sizes": [[128, 128]], + + "test_scales": [1], + + "top_k": 40, + "categories": 80, + "aggr_weight": 0.1, + "scores_thresh": 0.1, + "center_thresh": 0.1, + "nms_threshold": 0.5, + "suppres_ghost": true, + + "max_per_image": 100 + } +} diff --git a/db/__init__.py b/db/__init__.py new file mode 100755 index 0000000..e69de29 diff --git a/db/base.py b/db/base.py new file mode 100644 index 0000000..3b6d894 --- /dev/null +++ b/db/base.py @@ -0,0 +1,90 @@ +import os +import h5py +import numpy as np + +from config import system_configs + +class BASE(object): + def __init__(self): + self._split = None + self._db_inds = [] + self._image_ids = [] + + self._data = None + self._image_hdf5 = None + self._image_file = None + self._image_hdf5_file = None + + self._mean = np.zeros((3, ), dtype=np.float32) + self._std = np.ones((3, ), dtype=np.float32) + self._eig_val = np.ones((3, ), dtype=np.float32) + self._eig_vec = np.zeros((3, 3), dtype=np.float32) + + self._configs = {} + self._configs["data_aug"] = True + + self._data_rng = None + + @property + def data(self): + if self._data is None: + raise ValueError("data is not set") + return self._data + + @property + def configs(self): + return self._configs + + @property + def mean(self): + return self._mean + + @property + def std(self): + return self._std + + @property + def eig_val(self): + return self._eig_val + + @property + def eig_vec(self): + return self._eig_vec + + @property + def db_inds(self): + return self._db_inds + + @property + def split(self): + return self._split + + def update_config(self, new): + for key in new: + if key in self._configs: + self._configs[key] = new[key] + + def image_ids(self, ind): + return self._image_ids[ind] + + def image_file(self, ind): + if self._image_file is None: + raise ValueError("Image path is not initialized") + + image_id = self._image_ids[ind] + return self._image_file.format(image_id) + + def write_result(self, ind, all_bboxes, all_scores): + pass + + def evaluate(self, name): + pass + + def shuffle_inds(self, quiet=False): + if self._data_rng is None: + self._data_rng = np.random.RandomState(os.getpid()) + + if not quiet: + print("shuffling indices...") + rand_perm = self._data_rng.permutation(len(self._db_inds)) + self._db_inds = self._db_inds[rand_perm] diff --git a/db/coco.py b/db/coco.py new file mode 100755 index 0000000..3784005 --- /dev/null +++ b/db/coco.py @@ -0,0 +1,183 @@ +import sys +sys.path.insert(0, "data/coco/PythonAPI/") + +import os +import json +import numpy as np +import pickle + +from tqdm import tqdm +from db.detection import DETECTION +from config import system_configs +from pycocotools.coco import COCO +from pycocotools.cocoeval import COCOeval + +class MSCOCO(DETECTION): + def __init__(self, db_config, split): + super(MSCOCO, self).__init__(db_config) + data_dir = system_configs.data_dir + result_dir = system_configs.result_dir + cache_dir = system_configs.cache_dir + + self._split = split + self._dataset = { + "trainval": "trainval2014", + "minival": "minival2014", + "testdev": "testdev2017" + }[self._split] + + self._coco_dir = os.path.join(data_dir, "coco") + + self._label_dir = os.path.join(self._coco_dir, "annotations") + self._label_file = os.path.join(self._label_dir, "instances_{}.json") + self._label_file = self._label_file.format(self._dataset) + + self._image_dir = os.path.join(self._coco_dir, "images", self._dataset) + self._image_file = os.path.join(self._image_dir, "{}") + + self._data = "coco" + self._mean = np.array([0.40789654, 0.44719302, 0.47026115], dtype=np.float32) + self._std = np.array([0.28863828, 0.27408164, 0.27809835], dtype=np.float32) + self._eig_val = np.array([0.2141788, 0.01817699, 0.00341571], dtype=np.float32) + self._eig_vec = np.array([ + [-0.58752847, -0.69563484, 0.41340352], + [-0.5832747, 0.00994535, -0.81221408], + [-0.56089297, 0.71832671, 0.41158938] + ], dtype=np.float32) + + self._cat_ids = [ + 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, + 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, + 24, 25, 27, 28, 31, 32, 33, 34, 35, 36, + 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, + 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, + 58, 59, 60, 61, 62, 63, 64, 65, 67, 70, + 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, + 82, 84, 85, 86, 87, 88, 89, 90 + ] + self._classes = { + ind + 1: cat_id for ind, cat_id in enumerate(self._cat_ids) + } + self._coco_to_class_map = { + value: key for key, value in self._classes.items() + } + + self._cache_file = os.path.join(cache_dir, "coco_{}.pkl".format(self._dataset)) + self._load_data() + self._db_inds = np.arange(len(self._image_ids)) + + self._load_coco_data() + + def _load_data(self): + print("loading from cache file: {}".format(self._cache_file)) + if not os.path.exists(self._cache_file): + print("No cache file found...") + self._extract_data() + with open(self._cache_file, "wb") as f: + pickle.dump([self._detections, self._image_ids], f) + else: + with open(self._cache_file, "rb") as f: + self._detections, self._image_ids = pickle.load(f) + + def _load_coco_data(self): + self._coco = COCO(self._label_file) + with open(self._label_file, "r") as f: + data = json.load(f) + + coco_ids = self._coco.getImgIds() + eval_ids = { + self._coco.loadImgs(coco_id)[0]["file_name"]: coco_id + for coco_id in coco_ids + } + + self._coco_categories = data["categories"] + self._coco_eval_ids = eval_ids + + def class_name(self, cid): + cat_id = self._classes[cid] + cat = self._coco.loadCats([cat_id])[0] + return cat["name"] + + def _extract_data(self): + self._coco = COCO(self._label_file) + self._cat_ids = self._coco.getCatIds() + + coco_image_ids = self._coco.getImgIds() + + self._image_ids = [ + self._coco.loadImgs(img_id)[0]["file_name"] + for img_id in coco_image_ids + ] + self._detections = {} + for ind, (coco_image_id, image_id) in enumerate(tqdm(zip(coco_image_ids, self._image_ids))): + image = self._coco.loadImgs(coco_image_id)[0] + bboxes = [] + categories = [] + + for cat_id in self._cat_ids: + annotation_ids = self._coco.getAnnIds(imgIds=image["id"], catIds=cat_id) + annotations = self._coco.loadAnns(annotation_ids) + category = self._coco_to_class_map[cat_id] + for annotation in annotations: + bbox = np.array(annotation["bbox"]) + bbox[[2, 3]] += bbox[[0, 1]] + bboxes.append(bbox) + + categories.append(category) + + bboxes = np.array(bboxes, dtype=float) + categories = np.array(categories, dtype=float) + if bboxes.size == 0 or categories.size == 0: + self._detections[image_id] = np.zeros((0, 5), dtype=np.float32) + else: + self._detections[image_id] = np.hstack((bboxes, categories[:, None])) + + def detections(self, ind): + image_id = self._image_ids[ind] + detections = self._detections[image_id] + + return detections.astype(float).copy() + + def _to_float(self, x): + return float("{:.2f}".format(x)) + + def convert_to_coco(self, all_bboxes): + detections = [] + for image_id in all_bboxes: + coco_id = self._coco_eval_ids[image_id] + for cls_ind in all_bboxes[image_id]: + category_id = self._classes[cls_ind] + for bbox in all_bboxes[image_id][cls_ind]: + bbox[2] -= bbox[0] + bbox[3] -= bbox[1] + + score = bbox[4] + bbox = list(map(self._to_float, bbox[0:4])) + + detection = { + "image_id": coco_id, + "category_id": category_id, + "bbox": bbox, + "score": float("{:.2f}".format(score)) + } + + detections.append(detection) + return detections + + def evaluate(self, result_json, cls_ids, image_ids, gt_json=None): + if self._split == "testdev": + return None + + coco = self._coco if gt_json is None else COCO(gt_json) + + eval_ids = [self._coco_eval_ids[image_id] for image_id in image_ids] + cat_ids = [self._classes[cls_id] for cls_id in cls_ids] + + coco_dets = coco.loadRes(result_json) + coco_eval = COCOeval(coco, coco_dets, "bbox") + coco_eval.params.imgIds = eval_ids + coco_eval.params.catIds = cat_ids + coco_eval.evaluate() + coco_eval.accumulate() + coco_eval.summarize() + return coco_eval.stats[0], coco_eval.stats[12:] diff --git a/db/coco_extreme.py b/db/coco_extreme.py new file mode 100755 index 0000000..7795778 --- /dev/null +++ b/db/coco_extreme.py @@ -0,0 +1,211 @@ +import sys +sys.path.insert(0, "data/coco/PythonAPI/") + +import os +import json +import numpy as np +import pickle + +from tqdm import tqdm +from db.detection import DETECTION +from config import system_configs +from pycocotools.coco import COCO +from pycocotools.cocoeval import COCOeval + +class MSCOCOExtreme(DETECTION): + def __init__(self, db_config, split): + super(MSCOCOExtreme, self).__init__(db_config) + data_dir = system_configs.data_dir + cache_dir = system_configs.cache_dir + + self._split = split + self._dataset = { + "train": "train2017", + "val": "val2017", + "testdev": "test2017" + }[self._split] + + self._coco_dir = os.path.join(data_dir, "coco") + + self._label_dir = os.path.join(self._coco_dir, "annotations") + + if self._split == 'testdev': + self._label_file = os.path.join( + self._label_dir, "image_info_test-dev2017.json") + else: + self._label_file = os.path.join(self._label_dir, + "instances_extreme_{}.json") + self._label_file = self._label_file.format(self._dataset) + + self._image_dir = os.path.join(self._coco_dir, "images", self._dataset) + self._image_file = os.path.join(self._image_dir, "{}") + + self._data = "coco_extreme" + self._mean = np.array([0.40789654, 0.44719302, 0.47026115], + dtype=np.float32) + self._std = np.array([0.28863828, 0.27408164, 0.27809835], + dtype=np.float32) + self._eig_val = np.array([0.2141788, 0.01817699, 0.00341571], + dtype=np.float32) + self._eig_vec = np.array([ + [-0.58752847, -0.69563484, 0.41340352], + [-0.5832747, 0.00994535, -0.81221408], + [-0.56089297, 0.71832671, 0.41158938] + ], dtype=np.float32) + + self._cat_ids = [ + 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, + 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, + 24, 25, 27, 28, 31, 32, 33, 34, 35, 36, + 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, + 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, + 58, 59, 60, 61, 62, 63, 64, 65, 67, 70, + 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, + 82, 84, 85, 86, 87, 88, 89, 90 + ] + self._classes = { + ind + 1: cat_id for ind, cat_id in enumerate(self._cat_ids) + } + self._coco_to_class_map = { + value: key for key, value in self._classes.items() + } + + self._cache_file = os.path.join( + cache_dir, "coco_extreme_{}.pkl".format(self._dataset)) + self._load_data() + self._db_inds = np.arange(len(self._image_ids)) + + self._load_coco_data() + + def _load_data(self): + print("loading from cache file: {}".format(self._cache_file)) + if not os.path.exists(self._cache_file): + print("No cache file found...") + self._extract_data() + with open(self._cache_file, "wb") as f: + pickle.dump([self._detections, self._image_ids, + self._extreme_pts], f) + else: + with open(self._cache_file, "rb") as f: + self._detections, self._image_ids, \ + self._extreme_pts = pickle.load(f) + + def _load_coco_data(self): + self._coco = COCO(self._label_file) + with open(self._label_file, "r") as f: + data = json.load(f) + + coco_ids = self._coco.getImgIds() + eval_ids = { + self._coco.loadImgs(coco_id)[0]["file_name"]: coco_id + for coco_id in coco_ids + } + + self._coco_categories = data["categories"] + self._coco_eval_ids = eval_ids + + def class_name(self, cid): + cat_id = self._classes[cid] + cat = self._coco.loadCats([cat_id])[0] + return cat["name"] + + def _extract_data(self): + self._coco = COCO(self._label_file) + self._cat_ids = self._coco.getCatIds() + + coco_image_ids = self._coco.getImgIds() + + self._image_ids = [ + self._coco.loadImgs(img_id)[0]["file_name"] + for img_id in coco_image_ids + ] + self._detections = {} + self._extreme_pts = {} + for ind, (coco_image_id, image_id) in enumerate(tqdm(zip(coco_image_ids, self._image_ids))): + image = self._coco.loadImgs(coco_image_id)[0] + bboxes = [] + categories = [] + extreme_pts = [] + + for cat_id in self._cat_ids: + annotation_ids = self._coco.getAnnIds(imgIds=image["id"], catIds=cat_id) + annotations = self._coco.loadAnns(annotation_ids) + category = self._coco_to_class_map[cat_id] + for annotation in annotations: + bbox = np.array(annotation["bbox"]) + bbox[[2, 3]] += bbox[[0, 1]] + bboxes.append(bbox) + + categories.append(category) + if len(annotation['extreme_points']) == 0: + extreme_pts.append(np.zeros((4, 2), dtype=float)) + else: + extreme_pt = np.array(annotation['extreme_points']) + extreme_pts.append(extreme_pt) + + bboxes = np.array(bboxes, dtype=float) + categories = np.array(categories, dtype=float) + extreme_pts = np.array(extreme_pts, dtype=float) + if bboxes.size == 0 or categories.size == 0: + self._detections[image_id] = np.zeros((0, 5), dtype=np.float32) + self._extreme_pts[image_id] = np.zeros((0, 4, 2), + dtype=np.float32) + else: + self._detections[image_id] = np.hstack((bboxes, + categories[:, None])) + self._extreme_pts[image_id] = extreme_pts + + def detections(self, ind): + image_id = self._image_ids[ind] + detections = self._detections[image_id] + extreme_pts = self._extreme_pts[image_id] + + return detections.astype(float).copy(), \ + extreme_pts.astype(float).copy() + + def _to_float(self, x): + return float("{:.2f}".format(x)) + + def convert_to_coco(self, all_bboxes): + detections = [] + for image_id in all_bboxes: + coco_id = self._coco_eval_ids[image_id] + for cls_ind in all_bboxes[image_id]: + category_id = self._classes[cls_ind] + for bbox in all_bboxes[image_id][cls_ind]: + bbox[2] -= bbox[0] + bbox[3] -= bbox[1] + + score = bbox[4] + bbox_out = list(map(self._to_float, bbox[0:4])) + + detection = { + "image_id": coco_id, + "category_id": category_id, + "bbox": bbox_out, + "score": float("{:.2f}".format(score)) + } + if len(bbox) > 5: + extreme_points = list(map(self._to_float, bbox[5:13])) + detection["extreme_points"] = extreme_points + + detections.append(detection) + return detections + + def evaluate(self, result_json, cls_ids, image_ids, gt_json=None): + if self._split == "testdev": + return None + + coco = self._coco if gt_json is None else COCO(gt_json) + + eval_ids = [self._coco_eval_ids[image_id] for image_id in image_ids] + cat_ids = [self._classes[cls_id] for cls_id in cls_ids] + + coco_dets = coco.loadRes(result_json) + coco_eval = COCOeval(coco, coco_dets, "bbox") + coco_eval.params.imgIds = eval_ids + coco_eval.params.catIds = cat_ids + coco_eval.evaluate() + coco_eval.accumulate() + coco_eval.summarize() + return coco_eval.stats[0], coco_eval.stats[12:] diff --git a/db/datasets.py b/db/datasets.py new file mode 100644 index 0000000..1101b8f --- /dev/null +++ b/db/datasets.py @@ -0,0 +1,7 @@ +from db.coco import MSCOCO +from db.coco_extreme import MSCOCOExtreme + +datasets = { + "MSCOCO": MSCOCO, + "MSCOCOExtreme": MSCOCOExtreme +} diff --git a/db/detection.py b/db/detection.py new file mode 100644 index 0000000..b12f49e --- /dev/null +++ b/db/detection.py @@ -0,0 +1,53 @@ +import numpy as np +from db.base import BASE + +class DETECTION(BASE): + def __init__(self, db_config): + super(DETECTION, self).__init__() + + self._configs["categories"] = 80 + self._configs["rand_scales"] = [1] + self._configs["rand_scale_min"] = 0.8 + self._configs["rand_scale_max"] = 1.4 + self._configs["rand_scale_step"] = 0.2 + + self._configs["input_size"] = [511] + self._configs["output_sizes"] = [[128, 128]] + + self._configs["nms_threshold"] = 0.5 + self._configs["max_per_image"] = 100 + self._configs["top_k"] = 100 + self._configs["ae_threshold"] = 0.5 + self._configs["aggr_weight"] = 0.1 + self._configs["scores_thresh"] = 0.1 + self._configs["center_thresh"] = 0.1 + self._configs["suppres_ghost"] = False + self._configs["nms_kernel"] = 3 + + self._configs["nms_algorithm"] = "exp_soft_nms" + self._configs["weight_exp"] = 8 + self._configs["merge_bbox"] = False + + self._configs["data_aug"] = True + self._configs["lighting"] = True + + self._configs["border"] = 128 + self._configs["gaussian_bump"] = True + self._configs["gaussian_iou"] = 0.7 + self._configs["gaussian_radius"] = -1 + self._configs["rand_crop"] = False + self._configs["rand_color"] = False + self._configs["rand_pushes"] = False + self._configs["rand_samples"] = False + self._configs["special_crop"] = False + + self._configs["test_scales"] = [1] + + self.update_config(db_config) + + if self._configs["rand_scales"] is None: + self._configs["rand_scales"] = np.arange( + self._configs["rand_scale_min"], + self._configs["rand_scale_max"], + self._configs["rand_scale_step"] + ) diff --git a/demo.py b/demo.py new file mode 100755 index 0000000..cdb6a3f --- /dev/null +++ b/demo.py @@ -0,0 +1,290 @@ +#!/usr/bin/env python +import os +import json +import torch +import pprint +import argparse +import importlib +import numpy as np +import cv2 + +import matplotlib +matplotlib.use("Agg") + +from config import system_configs +from nnet.py_factory import NetworkFactory + +from config import system_configs +from utils import crop_image, normalize_ +from external.nms import soft_nms_with_points as soft_nms +from utils.color_map import colormap +from utils.visualize import vis_mask, vis_octagon, vis_ex, vis_class, vis_bbox +from dextr import Dextr + +torch.backends.cudnn.benchmark = False + +class_name = [ + '__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', + 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', + 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', + 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', + 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', + 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', + 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', + 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', + 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', + 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', + 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', + 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', + 'scissors', 'teddy bear', 'hair drier', 'toothbrush' +] + +image_ext = ['jpg', 'jpeg', 'png', 'webp'] + +def parse_args(): + parser = argparse.ArgumentParser(description="Demo CornerNet") + parser.add_argument("--demo", help="demo image path or folders", + default="", type=str) + parser.add_argument("--cfg_file", help="config file", + default='ExtremeNet', type=str) + parser.add_argument("--testiter", dest="testiter", + help="test at iteration i", + default=250000) + parser.add_argument("--show_mask", action='store_true', + help="Run Deep extreme cut to obtain accurate mask") + + args = parser.parse_args() + return args + +def _rescale_dets(detections, ratios, borders, sizes): + xs, ys = detections[..., 0:4:2], detections[..., 1:4:2] + xs /= ratios[:, 1][:, None, None] + ys /= ratios[:, 0][:, None, None] + xs -= borders[:, 2][:, None, None] + ys -= borders[:, 0][:, None, None] + np.clip(xs, 0, sizes[:, 1][:, None, None], out=xs) + np.clip(ys, 0, sizes[:, 0][:, None, None], out=ys) + +def _rescale_ex_pts(detections, ratios, borders, sizes): + xs, ys = detections[..., 5:13:2], detections[..., 6:13:2] + xs /= ratios[:, 1][:, None, None] + ys /= ratios[:, 0][:, None, None] + xs -= borders[:, 2][:, None, None] + ys -= borders[:, 0][:, None, None] + np.clip(xs, 0, sizes[:, 1][:, None, None], out=xs) + np.clip(ys, 0, sizes[:, 0][:, None, None], out=ys) + +def _box_inside(box2, box1): + inside = (box2[0] >= box1[0] and box2[1] >= box1[1] and \ + box2[2] <= box1[2] and box2[3] <= box1[3]) + return inside + +def kp_decode(nnet, images, K, kernel=3, aggr_weight=0.1, + scores_thresh=0.1, center_thresh=0.1, debug=False): + detections = nnet.test( + [images], kernel=kernel, aggr_weight=aggr_weight, + scores_thresh=scores_thresh, center_thresh=center_thresh, debug=debug) + detections = detections.data.cpu().numpy() + return detections + +if __name__ == "__main__": + args = parse_args() + cfg_file = os.path.join( + system_configs.config_dir, args.cfg_file + ".json") + print("cfg_file: {}".format(cfg_file)) + + with open(cfg_file, "r") as f: + configs = json.load(f) + + configs["system"]["snapshot_name"] = args.cfg_file + system_configs.update_config(configs["system"]) + print("system config...") + pprint.pprint(system_configs.full) + + test_iter = system_configs.max_iter if args.testiter is None \ + else args.testiter + print("loading parameters at iteration: {}".format(test_iter)) + print("building neural network...") + nnet = NetworkFactory(None) + print("loading parameters...") + nnet.load_params(test_iter) + nnet.cuda() + nnet.eval_mode() + + K = configs["db"]["top_k"] + aggr_weight = configs["db"]["aggr_weight"] + scores_thresh = configs["db"]["scores_thresh"] + center_thresh = configs["db"]["center_thresh"] + suppres_ghost = True + nms_kernel = 3 + + scales = configs["db"]["test_scales"] + weight_exp = 8 + categories = configs["db"]["categories"] + nms_threshold = configs["db"]["nms_threshold"] + max_per_image = configs["db"]["max_per_image"] + nms_algorithm = { + "nms": 0, + "linear_soft_nms": 1, + "exp_soft_nms": 2 + }["exp_soft_nms"] + if args.show_mask: + dextr = Dextr() + + + mean = np.array([0.40789654, 0.44719302, 0.47026115], dtype=np.float32) + std = np.array([0.28863828, 0.27408164, 0.27809835], dtype=np.float32) + top_bboxes = {} + + if os.path.isdir(args.demo): + image_names = [] + ls = os.listdir(args.demo) + for file_name in sorted(ls): + ext = file_name[file_name.rfind('.') + 1:].lower() + if ext in image_ext: + image_names.append(os.path.join(args.demo, file_name)) + else: + image_names = [args.demo] + + for image_id, image_name in enumerate(image_names): + image = cv2.imread(image_name) + + height, width = image.shape[0:2] + + detections = [] + + for scale in scales: + new_height = int(height * scale) + new_width = int(width * scale) + new_center = np.array([new_height // 2, new_width // 2]) + + inp_height = new_height | 127 + inp_width = new_width | 127 + + images = np.zeros((1, 3, inp_height, inp_width), dtype=np.float32) + ratios = np.zeros((1, 2), dtype=np.float32) + borders = np.zeros((1, 4), dtype=np.float32) + sizes = np.zeros((1, 2), dtype=np.float32) + + out_height, out_width = (inp_height + 1) // 4, (inp_width + 1) // 4 + height_ratio = out_height / inp_height + width_ratio = out_width / inp_width + + resized_image = cv2.resize(image, (new_width, new_height)) + resized_image, border, offset = crop_image( + resized_image, new_center, [inp_height, inp_width]) + + resized_image = resized_image / 255. + normalize_(resized_image, mean, std) + + images[0] = resized_image.transpose((2, 0, 1)) + borders[0] = border + sizes[0] = [int(height * scale), int(width * scale)] + ratios[0] = [height_ratio, width_ratio] + + images = np.concatenate((images, images[:, :, :, ::-1]), axis=0) + images = torch.from_numpy(images) + dets = kp_decode( + nnet, images, K, aggr_weight=aggr_weight, + scores_thresh=scores_thresh, center_thresh=center_thresh, + kernel=nms_kernel, debug=True) + dets = dets.reshape(2, -1, 14) + dets[1, :, [0, 2]] = out_width - dets[1, :, [2, 0]] + dets[1, :, [5, 7, 9, 11]] = out_width - dets[1, :, [5, 7, 9, 11]] + dets[1, :, [7, 8, 11, 12]] = dets[1, :, [11, 12, 7, 8]].copy() + dets = dets.reshape(1, -1, 14) + + _rescale_dets(dets, ratios, borders, sizes) + _rescale_ex_pts(dets, ratios, borders, sizes) + dets[:, :, 0:4] /= scale + dets[:, :, 5:13] /= scale + detections.append(dets) + + detections = np.concatenate(detections, axis=1) + + classes = detections[..., -1] + classes = classes[0] + detections = detections[0] + + # reject detections with negative scores + keep_inds = (detections[:, 4] > 0) + detections = detections[keep_inds] + classes = classes[keep_inds] + + top_bboxes[image_id] = {} + for j in range(categories): + keep_inds = (classes == j) + top_bboxes[image_id][j + 1] = \ + detections[keep_inds].astype(np.float32) + soft_nms(top_bboxes[image_id][j + 1], + Nt=nms_threshold, method=nms_algorithm) + + scores = np.hstack([ + top_bboxes[image_id][j][:, 4] + for j in range(1, categories + 1) + ]) + if len(scores) > max_per_image: + kth = len(scores) - max_per_image + thresh = np.partition(scores, kth)[kth] + for j in range(1, categories + 1): + keep_inds = (top_bboxes[image_id][j][:, 4] >= thresh) + top_bboxes[image_id][j] = top_bboxes[image_id][j][keep_inds] + + if suppres_ghost: + for j in range(1, categories + 1): + n = len(top_bboxes[image_id][j]) + for k in range(n): + inside_score = 0 + if top_bboxes[image_id][j][k, 4] > 0.2: + for t in range(n): + if _box_inside(top_bboxes[image_id][j][t], + top_bboxes[image_id][j][k]): + inside_score += top_bboxes[image_id][j][t, 4] + if inside_score > top_bboxes[image_id][j][k, 4] * 3: + top_bboxes[image_id][j][k, 4] /= 2 + + + if 1: # visualize + color_list = colormap(rgb=True) + mask_color_id = 0 + image = cv2.imread(image_name) + input_image = image.copy() + mask_image = image.copy() + bboxes = {} + for j in range(1, categories + 1): + keep_inds = (top_bboxes[image_id][j][:, 4] > 0.5) + cat_name = class_name[j] + for bbox in top_bboxes[image_id][j][keep_inds]: + sc = bbox[4] + ex = bbox[5:13].astype(np.int32).reshape(4, 2) + bbox = bbox[0:4].astype(np.int32) + txt = '{}{:.2f}'.format(cat_name, sc) + color_mask = color_list[mask_color_id % len(color_list), :3] + mask_color_id += 1 + image = vis_bbox(image, + (bbox[0], bbox[1], + bbox[2] - bbox[0], bbox[3] - bbox[1])) + image = vis_class(image, + (bbox[0], bbox[1] - 2), txt) + image = vis_octagon( + image, ex, color_mask) + image = vis_ex(image, ex, color_mask) + + if args.show_mask: + mask = dextr.segment(input_image[:, :, ::-1], ex) # BGR to RGB + mask = np.asfortranarray(mask.astype(np.uint8)) + mask_image = vis_bbox(mask_image, + (bbox[0], bbox[1], + bbox[2] - bbox[0], + bbox[3] - bbox[1])) + mask_image = vis_class(mask_image, + (bbox[0], bbox[1] - 2), txt) + mask_image = vis_mask(mask_image, mask, color_mask) + + if args.show_mask: + cv2.imshow('mask', mask_image) + cv2.imshow('out', image) + cv2.waitKey() + + + diff --git a/dextr b/dextr new file mode 160000 index 0000000..67ca085 --- /dev/null +++ b/dextr @@ -0,0 +1 @@ +Subproject commit 67ca085f9509eeb2b168b07294d72f7625509fa5 diff --git a/dextr.py b/dextr.py new file mode 100644 index 0000000..2ca8914 --- /dev/null +++ b/dextr.py @@ -0,0 +1,98 @@ +import os +import torch +from collections import OrderedDict +from PIL import Image +import numpy as np +from matplotlib import pyplot as plt +import sys +from torch.nn.functional import upsample +this_dir = os.path.dirname(__file__) +sys.path.insert(0, 'dextr') +import networks.deeplab_resnet as resnet +from dataloaders import helpers as helpers + + +class Dextr(object): + def __init__(self, model_path='', + gpu_id=0, flip_test=True): + if model_path == '': + model_path = os.path.join( + 'cache', 'dextr_pascal-sbd.pth') + self.pad = 50 + self.thres = 0.8 + self.device = torch.device( + "cuda:"+str(gpu_id) if torch.cuda.is_available() else "cpu") + self.flip_test = flip_test + + # Create the network and load the weights + self.net = resnet.resnet101(1, nInputChannels=4, classifier='psp') + print("Initializing weights from: {}".format(model_path)) + state_dict_checkpoint = torch.load( + model_path, map_location=lambda storage, loc: storage) + # Remove the prefix .module from the model when it is trained using DataParallel + if 'module.' in list(state_dict_checkpoint.keys())[0]: + new_state_dict = OrderedDict() + for k, v in state_dict_checkpoint.items(): + name = k[7:] # remove `module.` from multi-gpu training + new_state_dict[name] = v + else: + new_state_dict = state_dict_checkpoint + self.net.load_state_dict(new_state_dict) + self.net.eval() + self.net.to(self.device) + + def segment(self, image, extreme_points_ori): + # Crop image to the bounding box from the extreme points and resize + bbox = helpers.get_bbox(image, points=extreme_points_ori, pad=self.pad, zero_pad=True) + crop_image = helpers.crop_from_bbox(image, bbox, zero_pad=True) + resize_image = helpers.fixed_resize(crop_image, (512, 512)).astype(np.float32) + + # Generate extreme point heat map normalized to image values + extreme_points = extreme_points_ori - [np.min(extreme_points_ori[:, 0]), np.min(extreme_points_ori[:, 1])] + [self.pad, + self.pad] + extreme_points = (512 * extreme_points * [1 / crop_image.shape[1], 1 / crop_image.shape[0]]).astype(np.int) + extreme_heatmap = helpers.make_gt(resize_image, extreme_points, sigma=10) + extreme_heatmap = helpers.cstm_normalize(extreme_heatmap, 255) + + # Concatenate inputs and convert to tensor + input_dextr = np.concatenate((resize_image, extreme_heatmap[:, :, np.newaxis]), axis=2) + inputs = input_dextr.transpose((2, 0, 1))[np.newaxis, ...] + # import pdb; pdb.set_trace() + if self.flip_test: + inputs = np.concatenate([inputs, inputs[:, :, :, ::-1]], axis=0) + inputs = torch.from_numpy(inputs) + # Run a forward pass + inputs = inputs.to(self.device) + outputs = self.net.forward(inputs) + outputs = upsample(outputs, size=(512, 512), mode='bilinear', align_corners=True) + outputs = outputs.to(torch.device('cpu')) + outputs = outputs.data.numpy() + if self.flip_test: + outputs = (outputs[:1] + outputs[1:, :, :, ::-1]) / 2 + + pred = np.transpose(outputs[0, ...], (1, 2, 0)) + pred = 1 / (1 + np.exp(-pred)) + pred = np.squeeze(pred) + result = helpers.crop2fullmask(pred, bbox, im_size=image.shape[:2], zero_pad=True, relax=self.pad) > self.thres + return result + +if __name__ == '__main__': + dextr = Dextr() + # Read image and click the points + # image = np.array(Image.open('ims/dog-cat.jpg')) + image = np.array(Image.open(sys.argv[1])) + plt.ion() + plt.axis('off') + plt.imshow(image) + plt.title('Click the four extreme points of the objects\nHit enter when done (do not close the window)') + results = [] + + with torch.no_grad(): + while 1: + extreme_points_ori = np.array(plt.ginput(4, timeout=0)).astype(np.int) + result = dextr.segment(image, extreme_points_ori) + # import pdb; pdb.set_trace() + results.append(result) + # Plot the results + plt.imshow(helpers.overlay_masks(image / 255, results)) + plt.plot(extreme_points_ori[:, 0], extreme_points_ori[:, 1], 'gx') diff --git a/eval_dextr_mask.py b/eval_dextr_mask.py new file mode 100644 index 0000000..c6ad4b6 --- /dev/null +++ b/eval_dextr_mask.py @@ -0,0 +1,66 @@ +from dextr.dextr import Dextr +import pycocotools.coco as cocoapi +from pycocotools.cocoeval import COCOeval +from pycocotools import mask as COCOmask +import numpy as np +import sys +import cv2 +import json +from progress.bar import Bar +DEBUG = False +ANN_PATH = 'data/coco/annotations/instances_extreme_val2017.json' +IMG_DIR = 'data/coco/images/val2017/' + +if __name__ == '__main__': + dextr = Dextr() + coco = cocoapi.COCO(ANN_PATH) + pred_path = sys.argv[1] + out_path = pred_path[:-5] + '_segm.json' + data = json.load(open(pred_path, 'r')) + anns = data + results = [] + score_thresh = 0.2 + num_boxes = 0 + for i, ann in enumerate(anns): + if ann['score'] >= score_thresh: + num_boxes += 1 + + bar = Bar('Pred + Dextr', max=num_boxes) + for i, ann in enumerate(anns): + if ann['score'] < score_thresh: + continue + ex = np.array(ann['extreme_points'], dtype=np.int32).reshape(4, 2) + img_id = ann['image_id'] + img_info = coco.loadImgs(ids=[img_id])[0] + img_path = IMG_DIR + img_info['file_name'] + img = cv2.imread(img_path) + mask = dextr.segment(img[:, :, ::-1], ex) + mask = np.asfortranarray(mask.astype(np.uint8)) + if DEBUG: + if ann['score'] < 0.1: + continue + print(ann['score']) + img = (0.4 * img + 0.6 * mask.reshape( + mask.shape[0], mask.shape[1], 1) * 255).astype(np.uint8) + cv2.imshow('img', img) + cv2.waitKey() + encode = COCOmask.encode(mask) + if 'counts' in encode: + encode['counts'] = encode['counts'].decode("utf8") + pred = {'image_id': ann['image_id'], + 'category_id': ann['category_id'], + 'score': ann['score'], + 'segmentation': encode, + 'extreme_points': ann['extreme_points']} + results.append(pred) + Bar.suffix = '[{0}/{1}]| Total: {total:} | ETA: {eta:} |'.format( + i, num_boxes, total=bar.elapsed_td, eta=bar.eta_td) + bar.next() + bar.finish() + json.dump(results, open(out_path, 'w')) + + dets = coco.loadRes(out_path) + coco_eval = COCOeval(coco, dets, "segm") + coco_eval.evaluate() + coco_eval.accumulate() + coco_eval.summarize() diff --git a/external/.gitignore b/external/.gitignore new file mode 100644 index 0000000..f7c8c1a --- /dev/null +++ b/external/.gitignore @@ -0,0 +1,7 @@ +bbox.c +bbox.cpython-35m-x86_64-linux-gnu.so +bbox.cpython-36m-x86_64-linux-gnu.so + +nms.c +nms.cpython-35m-x86_64-linux-gnu.so +nms.cpython-36m-x86_64-linux-gnu.so diff --git a/external/Makefile b/external/Makefile new file mode 100644 index 0000000..a482398 --- /dev/null +++ b/external/Makefile @@ -0,0 +1,3 @@ +all: + python setup.py build_ext --inplace + rm -rf build diff --git a/external/__init__.py b/external/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/external/nms.pyx b/external/nms.pyx new file mode 100644 index 0000000..a70bede --- /dev/null +++ b/external/nms.pyx @@ -0,0 +1,421 @@ +# -------------------------------------------------------- +# Fast R-CNN +# Copyright (c) 2015 Microsoft +# Licensed under The MIT License [see LICENSE for details] +# Written by Ross Girshick +# -------------------------------------------------------- + +# ---------------------------------------------------------- +# Soft-NMS: Improving Object Detection With One Line of Code +# Copyright (c) University of Maryland, College Park +# Licensed under The MIT License [see LICENSE for details] +# Written by Navaneeth Bodla and Bharat Singh +# ---------------------------------------------------------- + +import numpy as np +cimport numpy as np + +cdef inline np.float32_t max(np.float32_t a, np.float32_t b): + return a if a >= b else b + +cdef inline np.float32_t min(np.float32_t a, np.float32_t b): + return a if a <= b else b + +def nms(np.ndarray[np.float32_t, ndim=2] dets, np.float thresh): + cdef np.ndarray[np.float32_t, ndim=1] x1 = dets[:, 0] + cdef np.ndarray[np.float32_t, ndim=1] y1 = dets[:, 1] + cdef np.ndarray[np.float32_t, ndim=1] x2 = dets[:, 2] + cdef np.ndarray[np.float32_t, ndim=1] y2 = dets[:, 3] + cdef np.ndarray[np.float32_t, ndim=1] scores = dets[:, 4] + + cdef np.ndarray[np.float32_t, ndim=1] areas = (x2 - x1 + 1) * (y2 - y1 + 1) + cdef np.ndarray[np.int_t, ndim=1] order = scores.argsort()[::-1] + + cdef int ndets = dets.shape[0] + cdef np.ndarray[np.int_t, ndim=1] suppressed = \ + np.zeros((ndets), dtype=np.int) + + # nominal indices + cdef int _i, _j + # sorted indices + cdef int i, j + # temp variables for box i's (the box currently under consideration) + cdef np.float32_t ix1, iy1, ix2, iy2, iarea + # variables for computing overlap with box j (lower scoring box) + cdef np.float32_t xx1, yy1, xx2, yy2 + cdef np.float32_t w, h + cdef np.float32_t inter, ovr + + keep = [] + for _i in range(ndets): + i = order[_i] + if suppressed[i] == 1: + continue + keep.append(i) + ix1 = x1[i] + iy1 = y1[i] + ix2 = x2[i] + iy2 = y2[i] + iarea = areas[i] + for _j in range(_i + 1, ndets): + j = order[_j] + if suppressed[j] == 1: + continue + xx1 = max(ix1, x1[j]) + yy1 = max(iy1, y1[j]) + xx2 = min(ix2, x2[j]) + yy2 = min(iy2, y2[j]) + w = max(0.0, xx2 - xx1 + 1) + h = max(0.0, yy2 - yy1 + 1) + inter = w * h + ovr = inter / (iarea + areas[j] - inter) + if ovr >= thresh: + suppressed[j] = 1 + + return keep + +def soft_nms(np.ndarray[float, ndim=2] boxes, float sigma=0.5, float Nt=0.3, float threshold=0.001, unsigned int method=0): + cdef unsigned int N = boxes.shape[0] + cdef float iw, ih, box_area + cdef float ua + cdef int pos = 0 + cdef float maxscore = 0 + cdef int maxpos = 0 + cdef float x1,x2,y1,y2,tx1,tx2,ty1,ty2,ts,area,weight,ov + + for i in range(N): + maxscore = boxes[i, 4] + maxpos = i + + tx1 = boxes[i,0] + ty1 = boxes[i,1] + tx2 = boxes[i,2] + ty2 = boxes[i,3] + ts = boxes[i,4] + + pos = i + 1 + # get max box + while pos < N: + if maxscore < boxes[pos, 4]: + maxscore = boxes[pos, 4] + maxpos = pos + pos = pos + 1 + + # add max box as a detection + boxes[i,0] = boxes[maxpos,0] + boxes[i,1] = boxes[maxpos,1] + boxes[i,2] = boxes[maxpos,2] + boxes[i,3] = boxes[maxpos,3] + boxes[i,4] = boxes[maxpos,4] + + # swap ith box with position of max box + boxes[maxpos,0] = tx1 + boxes[maxpos,1] = ty1 + boxes[maxpos,2] = tx2 + boxes[maxpos,3] = ty2 + boxes[maxpos,4] = ts + + tx1 = boxes[i,0] + ty1 = boxes[i,1] + tx2 = boxes[i,2] + ty2 = boxes[i,3] + ts = boxes[i,4] + + pos = i + 1 + # NMS iterations, note that N changes if detection boxes fall below threshold + while pos < N: + x1 = boxes[pos, 0] + y1 = boxes[pos, 1] + x2 = boxes[pos, 2] + y2 = boxes[pos, 3] + s = boxes[pos, 4] + + area = (x2 - x1 + 1) * (y2 - y1 + 1) + iw = (min(tx2, x2) - max(tx1, x1) + 1) + if iw > 0: + ih = (min(ty2, y2) - max(ty1, y1) + 1) + if ih > 0: + ua = float((tx2 - tx1 + 1) * (ty2 - ty1 + 1) + area - iw * ih) + ov = iw * ih / ua #iou between max box and detection box + + if method == 1: # linear + if ov > Nt: + weight = 1 - ov + else: + weight = 1 + elif method == 2: # gaussian + weight = np.exp(-(ov * ov)/sigma) + else: # original NMS + if ov > Nt: + weight = 0 + else: + weight = 1 + + boxes[pos, 4] = weight*boxes[pos, 4] + + # if box score falls below threshold, discard the box by swapping with last box + # update N + if boxes[pos, 4] < threshold: + boxes[pos,0] = boxes[N-1, 0] + boxes[pos,1] = boxes[N-1, 1] + boxes[pos,2] = boxes[N-1, 2] + boxes[pos,3] = boxes[N-1, 3] + boxes[pos,4] = boxes[N-1, 4] + N = N - 1 + pos = pos - 1 + + pos = pos + 1 + + keep = [i for i in range(N)] + return keep + +def soft_nms_with_points(np.ndarray[float, ndim=2] boxes, float sigma=0.5, float Nt=0.3, float threshold=0.001, unsigned int method=0): + cdef unsigned int N = boxes.shape[0] + cdef float iw, ih, box_area + cdef float ua + cdef int pos = 0 + cdef float maxscore = 0 + cdef int maxpos = 0 + cdef float x1,x2,y1,y2,tx1,tx2,ty1,ty2,ts,area,weight,ov + + for i in range(N): + maxscore = boxes[i, 4] + maxpos = i + + tx1 = boxes[i,0] + ty1 = boxes[i,1] + tx2 = boxes[i,2] + ty2 = boxes[i,3] + ts = boxes[i,4] + ttx = boxes[i,5] + tty = boxes[i,6] + tlx = boxes[i,7] + tly = boxes[i,8] + tbx = boxes[i,9] + tby = boxes[i,10] + trx = boxes[i,11] + try_ = boxes[i,12] + + pos = i + 1 + # get max box + while pos < N: + if maxscore < boxes[pos, 4]: + maxscore = boxes[pos, 4] + maxpos = pos + pos = pos + 1 + + # add max box as a detection + boxes[i,0] = boxes[maxpos,0] + boxes[i,1] = boxes[maxpos,1] + boxes[i,2] = boxes[maxpos,2] + boxes[i,3] = boxes[maxpos,3] + boxes[i,4] = boxes[maxpos,4] + boxes[i,5] = boxes[maxpos,5] + boxes[i,6] = boxes[maxpos,6] + boxes[i,7] = boxes[maxpos,7] + boxes[i,8] = boxes[maxpos,8] + boxes[i,9] = boxes[maxpos,9] + boxes[i,10] = boxes[maxpos,10] + boxes[i,11] = boxes[maxpos,11] + boxes[i,12] = boxes[maxpos,12] + + # swap ith box with position of max box + boxes[maxpos,0] = tx1 + boxes[maxpos,1] = ty1 + boxes[maxpos,2] = tx2 + boxes[maxpos,3] = ty2 + boxes[maxpos,4] = ts + boxes[maxpos,5] = ttx + boxes[maxpos,6] = tty + boxes[maxpos,7] = tlx + boxes[maxpos,8] = tly + boxes[maxpos,9] = tbx + boxes[maxpos,10] = tby + boxes[maxpos,11] = trx + boxes[maxpos,12] = try_ + + tx1 = boxes[i,0] + ty1 = boxes[i,1] + tx2 = boxes[i,2] + ty2 = boxes[i,3] + ts = boxes[i,4] + ttx = boxes[i,5] + tty = boxes[i,6] + tlx = boxes[i,7] + tly = boxes[i,8] + tbx = boxes[i,9] + tby = boxes[i,10] + trx = boxes[i,11] + try_ = boxes[i,12] + + pos = i + 1 + # NMS iterations, note that N changes if detection boxes fall below threshold + while pos < N: + x1 = boxes[pos, 0] + y1 = boxes[pos, 1] + x2 = boxes[pos, 2] + y2 = boxes[pos, 3] + s = boxes[pos, 4] + + area = (x2 - x1 + 1) * (y2 - y1 + 1) + iw = (min(tx2, x2) - max(tx1, x1) + 1) + if iw > 0: + ih = (min(ty2, y2) - max(ty1, y1) + 1) + if ih > 0: + ua = float((tx2 - tx1 + 1) * (ty2 - ty1 + 1) + area - iw * ih) + ov = iw * ih / ua #iou between max box and detection box + + if method == 1: # linear + if ov > Nt: + weight = 1 - ov + else: + weight = 1 + elif method == 2: # gaussian + weight = np.exp(-(ov * ov)/sigma) + else: # original NMS + if ov > Nt: + weight = 0 + else: + weight = 1 + + boxes[pos, 4] = weight*boxes[pos, 4] + + # if box score falls below threshold, discard the box by swapping with last box + # update N + if boxes[pos, 4] < threshold: + boxes[pos,0] = boxes[N-1, 0] + boxes[pos,1] = boxes[N-1, 1] + boxes[pos,2] = boxes[N-1, 2] + boxes[pos,3] = boxes[N-1, 3] + boxes[pos,4] = boxes[N-1, 4] + boxes[pos,5] = boxes[N-1, 5] + boxes[pos,6] = boxes[N-1, 6] + boxes[pos,7] = boxes[N-1, 7] + boxes[pos,8] = boxes[N-1, 8] + boxes[pos,9] = boxes[N-1, 9] + boxes[pos,10] = boxes[N-1, 10] + boxes[pos,11] = boxes[N-1, 11] + boxes[pos,12] = boxes[N-1, 12] + N = N - 1 + pos = pos - 1 + + pos = pos + 1 + + keep = [i for i in range(N)] + return keep + +def soft_nms_merge(np.ndarray[float, ndim=2] boxes, float sigma=0.5, float Nt=0.3, float threshold=0.001, unsigned int method=0, float weight_exp=6): + cdef unsigned int N = boxes.shape[0] + cdef float iw, ih, box_area + cdef float ua + cdef int pos = 0 + cdef float maxscore = 0 + cdef int maxpos = 0 + cdef float x1,x2,y1,y2,tx1,tx2,ty1,ty2,ts,area,weight,ov + cdef float mx1,mx2,my1,my2,mts,mbs,mw + + for i in range(N): + maxscore = boxes[i, 4] + maxpos = i + + tx1 = boxes[i,0] + ty1 = boxes[i,1] + tx2 = boxes[i,2] + ty2 = boxes[i,3] + ts = boxes[i,4] + + pos = i + 1 + # get max box + while pos < N: + if maxscore < boxes[pos, 4]: + maxscore = boxes[pos, 4] + maxpos = pos + pos = pos + 1 + + # add max box as a detection + boxes[i,0] = boxes[maxpos,0] + boxes[i,1] = boxes[maxpos,1] + boxes[i,2] = boxes[maxpos,2] + boxes[i,3] = boxes[maxpos,3] + boxes[i,4] = boxes[maxpos,4] + + mx1 = boxes[i, 0] * boxes[i, 5] + my1 = boxes[i, 1] * boxes[i, 5] + mx2 = boxes[i, 2] * boxes[i, 6] + my2 = boxes[i, 3] * boxes[i, 6] + mts = boxes[i, 5] + mbs = boxes[i, 6] + + # swap ith box with position of max box + boxes[maxpos,0] = tx1 + boxes[maxpos,1] = ty1 + boxes[maxpos,2] = tx2 + boxes[maxpos,3] = ty2 + boxes[maxpos,4] = ts + + tx1 = boxes[i,0] + ty1 = boxes[i,1] + tx2 = boxes[i,2] + ty2 = boxes[i,3] + ts = boxes[i,4] + + pos = i + 1 + # NMS iterations, note that N changes if detection boxes fall below threshold + while pos < N: + x1 = boxes[pos, 0] + y1 = boxes[pos, 1] + x2 = boxes[pos, 2] + y2 = boxes[pos, 3] + s = boxes[pos, 4] + + area = (x2 - x1 + 1) * (y2 - y1 + 1) + iw = (min(tx2, x2) - max(tx1, x1) + 1) + if iw > 0: + ih = (min(ty2, y2) - max(ty1, y1) + 1) + if ih > 0: + ua = float((tx2 - tx1 + 1) * (ty2 - ty1 + 1) + area - iw * ih) + ov = iw * ih / ua #iou between max box and detection box + + if method == 1: # linear + if ov > Nt: + weight = 1 - ov + else: + weight = 1 + elif method == 2: # gaussian + weight = np.exp(-(ov * ov)/sigma) + else: # original NMS + if ov > Nt: + weight = 0 + else: + weight = 1 + + mw = (1 - weight) ** weight_exp + mx1 = mx1 + boxes[pos, 0] * boxes[pos, 5] * mw + my1 = my1 + boxes[pos, 1] * boxes[pos, 5] * mw + mx2 = mx2 + boxes[pos, 2] * boxes[pos, 6] * mw + my2 = my2 + boxes[pos, 3] * boxes[pos, 6] * mw + mts = mts + boxes[pos, 5] * mw + mbs = mbs + boxes[pos, 6] * mw + + boxes[pos, 4] = weight*boxes[pos, 4] + + # if box score falls below threshold, discard the box by swapping with last box + # update N + if boxes[pos, 4] < threshold: + boxes[pos,0] = boxes[N-1, 0] + boxes[pos,1] = boxes[N-1, 1] + boxes[pos,2] = boxes[N-1, 2] + boxes[pos,3] = boxes[N-1, 3] + boxes[pos,4] = boxes[N-1, 4] + N = N - 1 + pos = pos - 1 + + pos = pos + 1 + + boxes[i, 0] = mx1 / mts + boxes[i, 1] = my1 / mts + boxes[i, 2] = mx2 / mbs + boxes[i, 3] = my2 / mbs + + keep = [i for i in range(N)] + return keep diff --git a/external/setup.py b/external/setup.py new file mode 100644 index 0000000..c4d2571 --- /dev/null +++ b/external/setup.py @@ -0,0 +1,18 @@ +import numpy +from distutils.core import setup +from distutils.extension import Extension +from Cython.Build import cythonize + +extensions = [ + Extension( + "nms", + ["nms.pyx"], + extra_compile_args=["-Wno-cpp", "-Wno-unused-function"] + ) +] + +setup( + name="coco", + ext_modules=cythonize(extensions), + include_dirs=[numpy.get_include()] +) diff --git a/models/CornerNet.py b/models/CornerNet.py new file mode 100755 index 0000000..f39b2e0 --- /dev/null +++ b/models/CornerNet.py @@ -0,0 +1,83 @@ +import torch +import torch.nn as nn + +from .py_utils import kp, AELoss, _neg_loss, convolution, residual +from .py_utils import TopPool, BottomPool, LeftPool, RightPool + +class pool(nn.Module): + def __init__(self, dim, pool1, pool2): + super(pool, self).__init__() + self.p1_conv1 = convolution(3, dim, 128) + self.p2_conv1 = convolution(3, dim, 128) + + self.p_conv1 = nn.Conv2d(128, dim, (3, 3), padding=(1, 1), bias=False) + self.p_bn1 = nn.BatchNorm2d(dim) + + self.conv1 = nn.Conv2d(dim, dim, (1, 1), bias=False) + self.bn1 = nn.BatchNorm2d(dim) + self.relu1 = nn.ReLU(inplace=True) + + self.conv2 = convolution(3, dim, dim) + + self.pool1 = pool1() + self.pool2 = pool2() + + def forward(self, x): + # pool 1 + p1_conv1 = self.p1_conv1(x) + pool1 = self.pool1(p1_conv1) + + # pool 2 + p2_conv1 = self.p2_conv1(x) + pool2 = self.pool2(p2_conv1) + + # pool 1 + pool 2 + p_conv1 = self.p_conv1(pool1 + pool2) + p_bn1 = self.p_bn1(p_conv1) + + conv1 = self.conv1(x) + bn1 = self.bn1(conv1) + relu1 = self.relu1(p_bn1 + bn1) + + conv2 = self.conv2(relu1) + return conv2 + +class tl_pool(pool): + def __init__(self, dim): + super(tl_pool, self).__init__(dim, TopPool, LeftPool) + +class br_pool(pool): + def __init__(self, dim): + super(br_pool, self).__init__(dim, BottomPool, RightPool) + +def make_tl_layer(dim): + return tl_pool(dim) + +def make_br_layer(dim): + return br_pool(dim) + +def make_pool_layer(dim): + return nn.Sequential() + +def make_hg_layer(kernel, dim0, dim1, mod, layer=convolution, **kwargs): + layers = [layer(kernel, dim0, dim1, stride=2)] + layers += [layer(kernel, dim1, dim1) for _ in range(mod - 1)] + return nn.Sequential(*layers) + +class model(kp): + def __init__(self, db): + n = 5 + dims = [256, 256, 384, 384, 384, 512] + modules = [2, 2, 2, 2, 2, 4] + out_dim = 80 + + super(model, self).__init__( + n, 2, dims, modules, out_dim, + make_tl_layer=make_tl_layer, + make_br_layer=make_br_layer, + make_pool_layer=make_pool_layer, + make_hg_layer=make_hg_layer, + kp_layer=residual, cnv_dim=256 + ) + +loss = AELoss(pull_weight=1e-1, push_weight=1e-1, focal_loss=_neg_loss) diff --git a/models/ExtremeNet.py b/models/ExtremeNet.py new file mode 100755 index 0000000..58bbf68 --- /dev/null +++ b/models/ExtremeNet.py @@ -0,0 +1,30 @@ +import torch +import torch.nn as nn + +from .py_utils import exkp, CTLoss, _neg_loss, convolution, residual + +def make_pool_layer(dim): + return nn.Sequential() + +def make_hg_layer(kernel, dim0, dim1, mod, layer=convolution, **kwargs): + layers = [layer(kernel, dim0, dim1, stride=2)] + layers += [layer(kernel, dim1, dim1) for _ in range(mod - 1)] + return nn.Sequential(*layers) + +class model(exkp): + def __init__(self, db): + n = 5 + dims = [256, 256, 384, 384, 384, 512] + modules = [2, 2, 2, 2, 2, 4] + out_dim = 80 + + super(model, self).__init__( + n, 2, dims, modules, out_dim, + make_tl_layer=None, + make_br_layer=None, + make_pool_layer=make_pool_layer, + make_hg_layer=make_hg_layer, + kp_layer=residual, cnv_dim=256 + ) + +loss = CTLoss(focal_loss=_neg_loss) diff --git a/models/__init__.py b/models/__init__.py new file mode 100755 index 0000000..e69de29 diff --git a/models/py_utils/__init__.py b/models/py_utils/__init__.py new file mode 100644 index 0000000..01bf90a --- /dev/null +++ b/models/py_utils/__init__.py @@ -0,0 +1,8 @@ +from .kp import kp, AELoss +from .exkp import exkp, CTLoss +from .kp_utils import _neg_loss + +from .utils import convolution, fully_connected, residual + +# Un-comment this line if your want to run CornerNet +# from ._cpools import TopPool, BottomPool, LeftPool, RightPool diff --git a/models/py_utils/_cpools/.gitignore b/models/py_utils/_cpools/.gitignore new file mode 100644 index 0000000..6a0882d --- /dev/null +++ b/models/py_utils/_cpools/.gitignore @@ -0,0 +1,3 @@ +build/ +cpools.egg-info/ +dist/ diff --git a/models/py_utils/_cpools/__init__.py b/models/py_utils/_cpools/__init__.py new file mode 100644 index 0000000..1b4e76f --- /dev/null +++ b/models/py_utils/_cpools/__init__.py @@ -0,0 +1,74 @@ +import torch + +from torch import nn +from torch.autograd import Function + +import top_pool, bottom_pool, left_pool, right_pool + +class TopPoolFunction(Function): + @staticmethod + def forward(ctx, input): + output = top_pool.forward(input)[0] + ctx.save_for_backward(input) + return output + + @staticmethod + def backward(ctx, grad_output): + input = ctx.saved_variables[0] + output = top_pool.backward(input, grad_output)[0] + return output + +class BottomPoolFunction(Function): + @staticmethod + def forward(ctx, input): + output = bottom_pool.forward(input)[0] + ctx.save_for_backward(input) + return output + + @staticmethod + def backward(ctx, grad_output): + input = ctx.saved_variables[0] + output = bottom_pool.backward(input, grad_output)[0] + return output + +class LeftPoolFunction(Function): + @staticmethod + def forward(ctx, input): + output = left_pool.forward(input)[0] + ctx.save_for_backward(input) + return output + + @staticmethod + def backward(ctx, grad_output): + input = ctx.saved_variables[0] + output = left_pool.backward(input, grad_output)[0] + return output + +class RightPoolFunction(Function): + @staticmethod + def forward(ctx, input): + output = right_pool.forward(input)[0] + ctx.save_for_backward(input) + return output + + @staticmethod + def backward(ctx, grad_output): + input = ctx.saved_variables[0] + output = right_pool.backward(input, grad_output)[0] + return output + +class TopPool(nn.Module): + def forward(self, x): + return TopPoolFunction.apply(x) + +class BottomPool(nn.Module): + def forward(self, x): + return BottomPoolFunction.apply(x) + +class LeftPool(nn.Module): + def forward(self, x): + return LeftPoolFunction.apply(x) + +class RightPool(nn.Module): + def forward(self, x): + return RightPoolFunction.apply(x) diff --git a/models/py_utils/_cpools/setup.py b/models/py_utils/_cpools/setup.py new file mode 100644 index 0000000..9682833 --- /dev/null +++ b/models/py_utils/_cpools/setup.py @@ -0,0 +1,15 @@ +from setuptools import setup +from torch.utils.cpp_extension import BuildExtension, CppExtension + +setup( + name="cpools", + ext_modules=[ + CppExtension("top_pool", ["src/top_pool.cpp"]), + CppExtension("bottom_pool", ["src/bottom_pool.cpp"]), + CppExtension("left_pool", ["src/left_pool.cpp"]), + CppExtension("right_pool", ["src/right_pool.cpp"]) + ], + cmdclass={ + "build_ext": BuildExtension + } +) diff --git a/models/py_utils/_cpools/src/bottom_pool.cpp b/models/py_utils/_cpools/src/bottom_pool.cpp new file mode 100644 index 0000000..bd6c65a --- /dev/null +++ b/models/py_utils/_cpools/src/bottom_pool.cpp @@ -0,0 +1,85 @@ +#include + +#include + +std::vector pool_forward( + at::Tensor input +) { + // Initialize output + at::Tensor output = at::zeros_like(input); + + // Get height + int64_t height = input.size(2); + + // Copy the last column + at::Tensor input_temp = input.select(2, 0); + at::Tensor output_temp = output.select(2, 0); + output_temp.copy_(input_temp); + + at::Tensor max_temp; + for (int64_t ind = 0; ind < height - 1; ++ind) { + input_temp = input.select(2, ind + 1); + output_temp = output.select(2, ind); + max_temp = output.select(2, ind + 1); + + at::max_out(max_temp, input_temp, output_temp); + } + + return { + output + }; +} + +std::vector pool_backward( + at::Tensor input, + at::Tensor grad_output +) { + auto output = at::zeros_like(input); + + int32_t batch = input.size(0); + int32_t channel = input.size(1); + int32_t height = input.size(2); + int32_t width = input.size(3); + + auto max_val = at::zeros(torch::CUDA(at::kFloat), {batch, channel, width}); + auto max_ind = at::zeros(torch::CUDA(at::kLong), {batch, channel, width}); + + auto input_temp = input.select(2, 0); + max_val.copy_(input_temp); + + max_ind.fill_(0); + + auto output_temp = output.select(2, 0); + auto grad_output_temp = grad_output.select(2, 0); + output_temp.copy_(grad_output_temp); + + auto un_max_ind = max_ind.unsqueeze(2); + auto gt_mask = at::zeros(torch::CUDA(at::kByte), {batch, channel, width}); + auto max_temp = at::zeros(torch::CUDA(at::kFloat), {batch, channel, width}); + for (int32_t ind = 0; ind < height - 1; ++ind) { + input_temp = input.select(2, ind + 1); + at::gt_out(gt_mask, input_temp, max_val); + + at::masked_select_out(max_temp, input_temp, gt_mask); + max_val.masked_scatter_(gt_mask, max_temp); + max_ind.masked_fill_(gt_mask, ind + 1); + + grad_output_temp = grad_output.select(2, ind + 1).unsqueeze(2); + output.scatter_add_(2, un_max_ind, grad_output_temp); + } + + return { + output + }; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def( + "forward", &pool_forward, "Bottom Pool Forward", + py::call_guard() + ); + m.def( + "backward", &pool_backward, "Bottom Pool Backward", + py::call_guard() + ); +} diff --git a/models/py_utils/_cpools/src/left_pool.cpp b/models/py_utils/_cpools/src/left_pool.cpp new file mode 100644 index 0000000..fbc5d98 --- /dev/null +++ b/models/py_utils/_cpools/src/left_pool.cpp @@ -0,0 +1,85 @@ +#include + +#include + +std::vector pool_forward( + at::Tensor input +) { + // Initialize output + at::Tensor output = at::zeros_like(input); + + // Get width + int64_t width = input.size(3); + + // Copy the last column + at::Tensor input_temp = input.select(3, width - 1); + at::Tensor output_temp = output.select(3, width - 1); + output_temp.copy_(input_temp); + + at::Tensor max_temp; + for (int64_t ind = 1; ind < width; ++ind) { + input_temp = input.select(3, width - ind - 1); + output_temp = output.select(3, width - ind); + max_temp = output.select(3, width - ind - 1); + + at::max_out(max_temp, input_temp, output_temp); + } + + return { + output + }; +} + +std::vector pool_backward( + at::Tensor input, + at::Tensor grad_output +) { + auto output = at::zeros_like(input); + + int32_t batch = input.size(0); + int32_t channel = input.size(1); + int32_t height = input.size(2); + int32_t width = input.size(3); + + auto max_val = at::zeros(torch::CUDA(at::kFloat), {batch, channel, height}); + auto max_ind = at::zeros(torch::CUDA(at::kLong), {batch, channel, height}); + + auto input_temp = input.select(3, width - 1); + max_val.copy_(input_temp); + + max_ind.fill_(width - 1); + + auto output_temp = output.select(3, width - 1); + auto grad_output_temp = grad_output.select(3, width - 1); + output_temp.copy_(grad_output_temp); + + auto un_max_ind = max_ind.unsqueeze(3); + auto gt_mask = at::zeros(torch::CUDA(at::kByte), {batch, channel, height}); + auto max_temp = at::zeros(torch::CUDA(at::kFloat), {batch, channel, height}); + for (int32_t ind = 1; ind < width; ++ind) { + input_temp = input.select(3, width - ind - 1); + at::gt_out(gt_mask, input_temp, max_val); + + at::masked_select_out(max_temp, input_temp, gt_mask); + max_val.masked_scatter_(gt_mask, max_temp); + max_ind.masked_fill_(gt_mask, width - ind - 1); + + grad_output_temp = grad_output.select(3, width - ind - 1).unsqueeze(3); + output.scatter_add_(3, un_max_ind, grad_output_temp); + } + + return { + output + }; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def( + "forward", &pool_forward, "Left Pool Forward", + py::call_guard() + ); + m.def( + "backward", &pool_backward, "Left Pool Backward", + py::call_guard() + ); +} diff --git a/models/py_utils/_cpools/src/right_pool.cpp b/models/py_utils/_cpools/src/right_pool.cpp new file mode 100644 index 0000000..36c5c85 --- /dev/null +++ b/models/py_utils/_cpools/src/right_pool.cpp @@ -0,0 +1,85 @@ +#include + +#include + +std::vector pool_forward( + at::Tensor input +) { + // Initialize output + at::Tensor output = at::zeros_like(input); + + // Get width + int64_t width = input.size(3); + + // Copy the last column + at::Tensor input_temp = input.select(3, 0); + at::Tensor output_temp = output.select(3, 0); + output_temp.copy_(input_temp); + + at::Tensor max_temp; + for (int64_t ind = 0; ind < width - 1; ++ind) { + input_temp = input.select(3, ind + 1); + output_temp = output.select(3, ind); + max_temp = output.select(3, ind + 1); + + at::max_out(max_temp, input_temp, output_temp); + } + + return { + output + }; +} + +std::vector pool_backward( + at::Tensor input, + at::Tensor grad_output +) { + at::Tensor output = at::zeros_like(input); + + int32_t batch = input.size(0); + int32_t channel = input.size(1); + int32_t height = input.size(2); + int32_t width = input.size(3); + + auto max_val = at::zeros(torch::CUDA(at::kFloat), {batch, channel, height}); + auto max_ind = at::zeros(torch::CUDA(at::kLong), {batch, channel, height}); + + auto input_temp = input.select(3, 0); + max_val.copy_(input_temp); + + max_ind.fill_(0); + + auto output_temp = output.select(3, 0); + auto grad_output_temp = grad_output.select(3, 0); + output_temp.copy_(grad_output_temp); + + auto un_max_ind = max_ind.unsqueeze(3); + auto gt_mask = at::zeros(torch::CUDA(at::kByte), {batch, channel, height}); + auto max_temp = at::zeros(torch::CUDA(at::kFloat), {batch, channel, height}); + for (int32_t ind = 0; ind < width - 1; ++ind) { + input_temp = input.select(3, ind + 1); + at::gt_out(gt_mask, input_temp, max_val); + + at::masked_select_out(max_temp, input_temp, gt_mask); + max_val.masked_scatter_(gt_mask, max_temp); + max_ind.masked_fill_(gt_mask, ind + 1); + + grad_output_temp = grad_output.select(3, ind + 1).unsqueeze(3); + output.scatter_add_(3, un_max_ind, grad_output_temp); + } + + return { + output + }; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def( + "forward", &pool_forward, "Right Pool Forward", + py::call_guard() + ); + m.def( + "backward", &pool_backward, "Right Pool Backward", + py::call_guard() + ); +} diff --git a/models/py_utils/_cpools/src/top_pool.cpp b/models/py_utils/_cpools/src/top_pool.cpp new file mode 100644 index 0000000..4ac287f --- /dev/null +++ b/models/py_utils/_cpools/src/top_pool.cpp @@ -0,0 +1,85 @@ +#include + +#include + +std::vector top_pool_forward( + at::Tensor input +) { + // Initialize output + at::Tensor output = at::zeros_like(input); + + // Get height + int64_t height = input.size(2); + + // Copy the last column + at::Tensor input_temp = input.select(2, height - 1); + at::Tensor output_temp = output.select(2, height - 1); + output_temp.copy_(input_temp); + + at::Tensor max_temp; + for (int64_t ind = 1; ind < height; ++ind) { + input_temp = input.select(2, height - ind - 1); + output_temp = output.select(2, height - ind); + max_temp = output.select(2, height - ind - 1); + + at::max_out(max_temp, input_temp, output_temp); + } + + return { + output + }; +} + +std::vector top_pool_backward( + at::Tensor input, + at::Tensor grad_output +) { + auto output = at::zeros_like(input); + + int32_t batch = input.size(0); + int32_t channel = input.size(1); + int32_t height = input.size(2); + int32_t width = input.size(3); + + auto max_val = at::zeros(torch::CUDA(at::kFloat), {batch, channel, width}); + auto max_ind = at::zeros(torch::CUDA(at::kLong), {batch, channel, width}); + + auto input_temp = input.select(2, height - 1); + max_val.copy_(input_temp); + + max_ind.fill_(height - 1); + + auto output_temp = output.select(2, height - 1); + auto grad_output_temp = grad_output.select(2, height - 1); + output_temp.copy_(grad_output_temp); + + auto un_max_ind = max_ind.unsqueeze(2); + auto gt_mask = at::zeros(torch::CUDA(at::kByte), {batch, channel, width}); + auto max_temp = at::zeros(torch::CUDA(at::kFloat), {batch, channel, width}); + for (int32_t ind = 1; ind < height; ++ind) { + input_temp = input.select(2, height - ind - 1); + at::gt_out(gt_mask, input_temp, max_val); + + at::masked_select_out(max_temp, input_temp, gt_mask); + max_val.masked_scatter_(gt_mask, max_temp); + max_ind.masked_fill_(gt_mask, height - ind - 1); + + grad_output_temp = grad_output.select(2, height - ind - 1).unsqueeze(2); + output.scatter_add_(2, un_max_ind, grad_output_temp); + } + + return { + output + }; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def( + "forward", &top_pool_forward, "Top Pool Forward", + py::call_guard() + ); + m.def( + "backward", &top_pool_backward, "Top Pool Backward", + py::call_guard() + ); +} diff --git a/models/py_utils/data_parallel.py b/models/py_utils/data_parallel.py new file mode 100644 index 0000000..cc2a5a8 --- /dev/null +++ b/models/py_utils/data_parallel.py @@ -0,0 +1,116 @@ +import torch +from torch.nn.modules import Module +from torch.nn.parallel.scatter_gather import gather +from torch.nn.parallel.replicate import replicate +from torch.nn.parallel.parallel_apply import parallel_apply + +from .scatter_gather import scatter_kwargs + +class DataParallel(Module): + r"""Implements data parallelism at the module level. + + This container parallelizes the application of the given module by + splitting the input across the specified devices by chunking in the batch + dimension. In the forward pass, the module is replicated on each device, + and each replica handles a portion of the input. During the backwards + pass, gradients from each replica are summed into the original module. + + The batch size should be larger than the number of GPUs used. It should + also be an integer multiple of the number of GPUs so that each chunk is the + same size (so that each GPU processes the same number of samples). + + See also: :ref:`cuda-nn-dataparallel-instead` + + Arbitrary positional and keyword inputs are allowed to be passed into + DataParallel EXCEPT Tensors. All variables will be scattered on dim + specified (default 0). Primitive types will be broadcasted, but all + other types will be a shallow copy and can be corrupted if written to in + the model's forward pass. + + Args: + module: module to be parallelized + device_ids: CUDA devices (default: all devices) + output_device: device location of output (default: device_ids[0]) + + Example:: + + >>> net = torch.nn.DataParallel(model, device_ids=[0, 1, 2]) + >>> output = net(input_var) + """ + + # TODO: update notes/cuda.rst when this class handles 8+ GPUs well + + def __init__(self, module, device_ids=None, output_device=None, dim=0, chunk_sizes=None): + super(DataParallel, self).__init__() + + if not torch.cuda.is_available(): + self.module = module + self.device_ids = [] + return + + if device_ids is None: + device_ids = list(range(torch.cuda.device_count())) + if output_device is None: + output_device = device_ids[0] + self.dim = dim + self.module = module + self.device_ids = device_ids + self.chunk_sizes = chunk_sizes + self.output_device = output_device + if len(self.device_ids) == 1: + self.module.cuda(device_ids[0]) + + def forward(self, *inputs, **kwargs): + if not self.device_ids: + return self.module(*inputs, **kwargs) + inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids, self.chunk_sizes) + if len(self.device_ids) == 1: + return self.module(*inputs[0], **kwargs[0]) + replicas = self.replicate(self.module, self.device_ids[:len(inputs)]) + outputs = self.parallel_apply(replicas, inputs, kwargs) + return self.gather(outputs, self.output_device) + + def replicate(self, module, device_ids): + return replicate(module, device_ids) + + def scatter(self, inputs, kwargs, device_ids, chunk_sizes): + return scatter_kwargs(inputs, kwargs, device_ids, dim=self.dim, chunk_sizes=self.chunk_sizes) + + def parallel_apply(self, replicas, inputs, kwargs): + return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)]) + + def gather(self, outputs, output_device): + return gather(outputs, output_device, dim=self.dim) + + +def data_parallel(module, inputs, device_ids=None, output_device=None, dim=0, module_kwargs=None): + r"""Evaluates module(input) in parallel across the GPUs given in device_ids. + + This is the functional version of the DataParallel module. + + Args: + module: the module to evaluate in parallel + inputs: inputs to the module + device_ids: GPU ids on which to replicate module + output_device: GPU location of the output Use -1 to indicate the CPU. + (default: device_ids[0]) + Returns: + a Variable containing the result of module(input) located on + output_device + """ + if not isinstance(inputs, tuple): + inputs = (inputs,) + + if device_ids is None: + device_ids = list(range(torch.cuda.device_count())) + + if output_device is None: + output_device = device_ids[0] + + inputs, module_kwargs = scatter_kwargs(inputs, module_kwargs, device_ids, dim) + if len(device_ids) == 1: + return module(*inputs[0], **module_kwargs[0]) + used_device_ids = device_ids[:len(inputs)] + replicas = replicate(module, used_device_ids) + outputs = parallel_apply(replicas, inputs, module_kwargs, used_device_ids) + return gather(outputs, output_device, dim) diff --git a/models/py_utils/exkp.py b/models/py_utils/exkp.py new file mode 100644 index 0000000..6f87fd3 --- /dev/null +++ b/models/py_utils/exkp.py @@ -0,0 +1,375 @@ +import numpy as np +import torch +import torch.nn as nn + +from .utils import convolution, residual +from .utils import make_layer, make_layer_revr + +from .kp_utils import _tranpose_and_gather_feat, _exct_decode +from .kp_utils import _sigmoid, _regr_loss, _neg_loss +from .kp_utils import make_kp_layer +from .kp_utils import make_pool_layer, make_unpool_layer +from .kp_utils import make_merge_layer, make_inter_layer, make_cnv_layer +from .kp_utils import _h_aggregate, _v_aggregate +from utils.debugger import Debugger + +class kp_module(nn.Module): + def __init__( + self, n, dims, modules, layer=residual, + make_up_layer=make_layer, make_low_layer=make_layer, + make_hg_layer=make_layer, make_hg_layer_revr=make_layer_revr, + make_pool_layer=make_pool_layer, make_unpool_layer=make_unpool_layer, + make_merge_layer=make_merge_layer, **kwargs + ): + super(kp_module, self).__init__() + + self.n = n + + curr_mod = modules[0] + next_mod = modules[1] + + curr_dim = dims[0] + next_dim = dims[1] + + self.up1 = make_up_layer( + 3, curr_dim, curr_dim, curr_mod, + layer=layer, **kwargs + ) + self.max1 = make_pool_layer(curr_dim) + self.low1 = make_hg_layer( + 3, curr_dim, next_dim, curr_mod, + layer=layer, **kwargs + ) + self.low2 = kp_module( + n - 1, dims[1:], modules[1:], layer=layer, + make_up_layer=make_up_layer, + make_low_layer=make_low_layer, + make_hg_layer=make_hg_layer, + make_hg_layer_revr=make_hg_layer_revr, + make_pool_layer=make_pool_layer, + make_unpool_layer=make_unpool_layer, + make_merge_layer=make_merge_layer, + **kwargs + ) if self.n > 1 else \ + make_low_layer( + 3, next_dim, next_dim, next_mod, + layer=layer, **kwargs + ) + self.low3 = make_hg_layer_revr( + 3, next_dim, curr_dim, curr_mod, + layer=layer, **kwargs + ) + self.up2 = make_unpool_layer(curr_dim) + + self.merge = make_merge_layer(curr_dim) + + def forward(self, x): + up1 = self.up1(x) + max1 = self.max1(x) + low1 = self.low1(max1) + low2 = self.low2(low1) + low3 = self.low3(low2) + up2 = self.up2(low3) + return self.merge(up1, up2) + +class exkp(nn.Module): + def __init__( + self, n, nstack, dims, modules, out_dim, pre=None, cnv_dim=256, + make_tl_layer=None, make_br_layer=None, + make_cnv_layer=make_cnv_layer, make_heat_layer=make_kp_layer, + make_tag_layer=make_kp_layer, make_regr_layer=make_kp_layer, + make_up_layer=make_layer, make_low_layer=make_layer, + make_hg_layer=make_layer, make_hg_layer_revr=make_layer_revr, + make_pool_layer=make_pool_layer, make_unpool_layer=make_unpool_layer, + make_merge_layer=make_merge_layer, make_inter_layer=make_inter_layer, + kp_layer=residual + ): + super(exkp, self).__init__() + + self.nstack = nstack + self._decode = _exct_decode + + curr_dim = dims[0] + + self.pre = nn.Sequential( + convolution(7, 3, 128, stride=2), + residual(3, 128, 256, stride=2) + ) if pre is None else pre + + self.kps = nn.ModuleList([ + kp_module( + n, dims, modules, layer=kp_layer, + make_up_layer=make_up_layer, + make_low_layer=make_low_layer, + make_hg_layer=make_hg_layer, + make_hg_layer_revr=make_hg_layer_revr, + make_pool_layer=make_pool_layer, + make_unpool_layer=make_unpool_layer, + make_merge_layer=make_merge_layer + ) for _ in range(nstack) + ]) + self.cnvs = nn.ModuleList([ + make_cnv_layer(curr_dim, cnv_dim) for _ in range(nstack) + ]) + + ## keypoint heatmaps + self.t_heats = nn.ModuleList([ + make_heat_layer(cnv_dim, curr_dim, out_dim) for _ in range(nstack) + ]) + + self.l_heats = nn.ModuleList([ + make_heat_layer(cnv_dim, curr_dim, out_dim) for _ in range(nstack) + ]) + + self.b_heats = nn.ModuleList([ + make_heat_layer(cnv_dim, curr_dim, out_dim) for _ in range(nstack) + ]) + + self.r_heats = nn.ModuleList([ + make_heat_layer(cnv_dim, curr_dim, out_dim) for _ in range(nstack) + ]) + + self.ct_heats = nn.ModuleList([ + make_heat_layer(cnv_dim, curr_dim, out_dim) for _ in range(nstack) + ]) + + for t_heat, l_heat, b_heat, r_heat, ct_heat in \ + zip(self.t_heats, self.l_heats, self.b_heats, \ + self.r_heats, self.ct_heats): + t_heat[-1].bias.data.fill_(-2.19) + l_heat[-1].bias.data.fill_(-2.19) + b_heat[-1].bias.data.fill_(-2.19) + r_heat[-1].bias.data.fill_(-2.19) + ct_heat[-1].bias.data.fill_(-2.19) + + self.inters = nn.ModuleList([ + make_inter_layer(curr_dim) for _ in range(nstack - 1) + ]) + + self.inters_ = nn.ModuleList([ + nn.Sequential( + nn.Conv2d(curr_dim, curr_dim, (1, 1), bias=False), + nn.BatchNorm2d(curr_dim) + ) for _ in range(nstack - 1) + ]) + self.cnvs_ = nn.ModuleList([ + nn.Sequential( + nn.Conv2d(cnv_dim, curr_dim, (1, 1), bias=False), + nn.BatchNorm2d(curr_dim) + ) for _ in range(nstack - 1) + ]) + + self.t_regrs = nn.ModuleList([ + make_regr_layer(cnv_dim, curr_dim, 2) for _ in range(nstack) + ]) + self.l_regrs = nn.ModuleList([ + make_regr_layer(cnv_dim, curr_dim, 2) for _ in range(nstack) + ]) + self.b_regrs = nn.ModuleList([ + make_regr_layer(cnv_dim, curr_dim, 2) for _ in range(nstack) + ]) + self.r_regrs = nn.ModuleList([ + make_regr_layer(cnv_dim, curr_dim, 2) for _ in range(nstack) + ]) + + self.relu = nn.ReLU(inplace=True) + + def _train(self, *xs): + image = xs[0] + t_inds = xs[1] + l_inds = xs[2] + b_inds = xs[3] + r_inds = xs[4] + + inter = self.pre(image) + outs = [] + + layers = zip( + self.kps, self.cnvs, + self.t_heats, self.l_heats, self.b_heats, self.r_heats, + self.ct_heats, + self.t_regrs, self.l_regrs, self.b_regrs, self.r_regrs, + ) + for ind, layer in enumerate(layers): + kp_, cnv_ = layer[0:2] + t_heat_, l_heat_, b_heat_, r_heat_ = layer[2:6] + ct_heat_ = layer[6] + t_regr_, l_regr_, b_regr_, r_regr_ = layer[7:11] + + kp = kp_(inter) + cnv = cnv_(kp) + + t_heat, l_heat = t_heat_(cnv), l_heat_(cnv) + b_heat, r_heat = b_heat_(cnv), r_heat_(cnv) + ct_heat = ct_heat_(cnv) + + t_regr, l_regr = t_regr_(cnv), l_regr_(cnv) + b_regr, r_regr = b_regr_(cnv), r_regr_(cnv) + + t_regr = _tranpose_and_gather_feat(t_regr, t_inds) + l_regr = _tranpose_and_gather_feat(l_regr, l_inds) + b_regr = _tranpose_and_gather_feat(b_regr, b_inds) + r_regr = _tranpose_and_gather_feat(r_regr, r_inds) + + outs += [t_heat, l_heat, b_heat, r_heat, ct_heat, \ + t_regr, l_regr, b_regr, r_regr] + + if ind < self.nstack - 1: + inter = self.inters_[ind](inter) + self.cnvs_[ind](cnv) + inter = self.relu(inter) + inter = self.inters[ind](inter) + return outs + + def _test(self, *xs, **kwargs): + image = xs[0] + + inter = self.pre(image) + outs = [] + + layers = zip( + self.kps, self.cnvs, + self.t_heats, self.l_heats, self.b_heats, self.r_heats, + self.ct_heats, + self.t_regrs, self.l_regrs, self.b_regrs, self.r_regrs, + ) + for ind, layer in enumerate(layers): + kp_, cnv_ = layer[0:2] + t_heat_, l_heat_, b_heat_, r_heat_ = layer[2:6] + ct_heat_ = layer[6] + t_regr_, l_regr_, b_regr_, r_regr_ = layer[7:11] + + kp = kp_(inter) + cnv = cnv_(kp) + + if ind == self.nstack - 1: + t_heat, l_heat = t_heat_(cnv), l_heat_(cnv) + b_heat, r_heat = b_heat_(cnv), r_heat_(cnv) + ct_heat = ct_heat_(cnv) + + t_regr, l_regr = t_regr_(cnv), l_regr_(cnv) + b_regr, r_regr = b_regr_(cnv), r_regr_(cnv) + + outs += [t_heat, l_heat, b_heat, r_heat, ct_heat, + t_regr, l_regr, b_regr, r_regr] + + if ind < self.nstack - 1: + inter = self.inters_[ind](inter) + self.cnvs_[ind](cnv) + inter = self.relu(inter) + inter = self.inters[ind](inter) + if kwargs['debug']: + _debug(image, t_heat, l_heat, b_heat, r_heat, ct_heat) + del kwargs['debug'] + return self._decode(*outs[-9:], **kwargs) + + def forward(self, *xs, **kwargs): + if len(xs) > 1: + return self._train(*xs, **kwargs) + return self._test(*xs, **kwargs) + +class CTLoss(nn.Module): + def __init__(self, regr_weight=1, focal_loss=_neg_loss): + super(CTLoss, self).__init__() + + self.regr_weight = regr_weight + self.focal_loss = focal_loss + self.regr_loss = _regr_loss + + def forward(self, outs, targets): + stride = 9 + + t_heats = outs[0::stride] + l_heats = outs[1::stride] + b_heats = outs[2::stride] + r_heats = outs[3::stride] + ct_heats = outs[4::stride] + t_regrs = outs[5::stride] + l_regrs = outs[6::stride] + b_regrs = outs[7::stride] + r_regrs = outs[8::stride] + + gt_t_heat = targets[0] + gt_l_heat = targets[1] + gt_b_heat = targets[2] + gt_r_heat = targets[3] + gt_ct_heat = targets[4] + gt_mask = targets[5] + gt_t_regr = targets[6] + gt_l_regr = targets[7] + gt_b_regr = targets[8] + gt_r_regr = targets[9] + + # focal loss + focal_loss = 0 + + t_heats = [_sigmoid(t) for t in t_heats] + l_heats = [_sigmoid(l) for l in l_heats] + b_heats = [_sigmoid(b) for b in b_heats] + r_heats = [_sigmoid(r) for r in r_heats] + ct_heats = [_sigmoid(ct) for ct in ct_heats] + + focal_loss += self.focal_loss(t_heats, gt_t_heat) + focal_loss += self.focal_loss(l_heats, gt_l_heat) + focal_loss += self.focal_loss(b_heats, gt_b_heat) + focal_loss += self.focal_loss(r_heats, gt_r_heat) + focal_loss += self.focal_loss(ct_heats, gt_ct_heat) + + # regression loss + regr_loss = 0 + for t_regr, l_regr, b_regr, r_regr in \ + zip(t_regrs, l_regrs, b_regrs, r_regrs): + regr_loss += self.regr_loss(t_regr, gt_t_regr, gt_mask) + regr_loss += self.regr_loss(l_regr, gt_l_regr, gt_mask) + regr_loss += self.regr_loss(b_regr, gt_b_regr, gt_mask) + regr_loss += self.regr_loss(r_regr, gt_r_regr, gt_mask) + regr_loss = self.regr_weight * regr_loss + + loss = (focal_loss + regr_loss) / len(t_heats) + return loss.unsqueeze(0) + +def _debug(image, t_heat, l_heat, b_heat, r_heat, ct_heat): + debugger = Debugger(num_classes=80) + k = 0 + + t_heat = torch.sigmoid(t_heat) + l_heat = torch.sigmoid(l_heat) + b_heat = torch.sigmoid(b_heat) + r_heat = torch.sigmoid(r_heat) + + + aggr_weight = 0.1 + t_heat = _h_aggregate(t_heat, aggr_weight=aggr_weight) + l_heat = _v_aggregate(l_heat, aggr_weight=aggr_weight) + b_heat = _h_aggregate(b_heat, aggr_weight=aggr_weight) + r_heat = _v_aggregate(r_heat, aggr_weight=aggr_weight) + t_heat[t_heat > 1] = 1 + l_heat[l_heat > 1] = 1 + b_heat[b_heat > 1] = 1 + r_heat[r_heat > 1] = 1 + + + ct_heat = torch.sigmoid(ct_heat) + + t_hm = debugger.gen_colormap(t_heat[k].cpu().data.numpy()) + l_hm = debugger.gen_colormap(l_heat[k].cpu().data.numpy()) + b_hm = debugger.gen_colormap(b_heat[k].cpu().data.numpy()) + r_hm = debugger.gen_colormap(r_heat[k].cpu().data.numpy()) + ct_hm = debugger.gen_colormap(ct_heat[k].cpu().data.numpy()) + + hms = np.maximum(np.maximum(t_hm, l_hm), + np.maximum(b_hm, r_hm)) + # debugger.add_img(hms, 'hms') + if image is not None: + mean = np.array([0.40789654, 0.44719302, 0.47026115], + dtype=np.float32).reshape(3, 1, 1) + std = np.array([0.28863828, 0.27408164, 0.27809835], + dtype=np.float32).reshape(3, 1, 1) + img = (image[k].cpu().data.numpy() * std + mean) * 255 + img = img.astype(np.uint8).transpose(1, 2, 0) + debugger.add_img(img, 'img') + # debugger.add_blend_img(img, t_hm, 't_hm') + # debugger.add_blend_img(img, l_hm, 'l_hm') + # debugger.add_blend_img(img, b_hm, 'b_hm') + # debugger.add_blend_img(img, r_hm, 'r_hm') + debugger.add_blend_img(img, hms, 'extreme') + debugger.add_blend_img(img, ct_hm, 'center') + debugger.show_all_imgs(pause=False) diff --git a/models/py_utils/kp.py b/models/py_utils/kp.py new file mode 100644 index 0000000..c18ed95 --- /dev/null +++ b/models/py_utils/kp.py @@ -0,0 +1,310 @@ +import numpy as np +import torch +import torch.nn as nn + +from .utils import convolution, residual +from .utils import make_layer, make_layer_revr + +from .kp_utils import _tranpose_and_gather_feat, _decode +from .kp_utils import _sigmoid, _ae_loss, _regr_loss, _neg_loss +from .kp_utils import make_tl_layer, make_br_layer, make_kp_layer +from .kp_utils import make_pool_layer, make_unpool_layer +from .kp_utils import make_merge_layer, make_inter_layer, make_cnv_layer + +class kp_module(nn.Module): + def __init__( + self, n, dims, modules, layer=residual, + make_up_layer=make_layer, make_low_layer=make_layer, + make_hg_layer=make_layer, make_hg_layer_revr=make_layer_revr, + make_pool_layer=make_pool_layer, make_unpool_layer=make_unpool_layer, + make_merge_layer=make_merge_layer, **kwargs + ): + super(kp_module, self).__init__() + + self.n = n + + curr_mod = modules[0] + next_mod = modules[1] + + curr_dim = dims[0] + next_dim = dims[1] + + self.up1 = make_up_layer( + 3, curr_dim, curr_dim, curr_mod, + layer=layer, **kwargs + ) + self.max1 = make_pool_layer(curr_dim) + self.low1 = make_hg_layer( + 3, curr_dim, next_dim, curr_mod, + layer=layer, **kwargs + ) + self.low2 = kp_module( + n - 1, dims[1:], modules[1:], layer=layer, + make_up_layer=make_up_layer, + make_low_layer=make_low_layer, + make_hg_layer=make_hg_layer, + make_hg_layer_revr=make_hg_layer_revr, + make_pool_layer=make_pool_layer, + make_unpool_layer=make_unpool_layer, + make_merge_layer=make_merge_layer, + **kwargs + ) if self.n > 1 else \ + make_low_layer( + 3, next_dim, next_dim, next_mod, + layer=layer, **kwargs + ) + self.low3 = make_hg_layer_revr( + 3, next_dim, curr_dim, curr_mod, + layer=layer, **kwargs + ) + self.up2 = make_unpool_layer(curr_dim) + + self.merge = make_merge_layer(curr_dim) + + def forward(self, x): + up1 = self.up1(x) + max1 = self.max1(x) + low1 = self.low1(max1) + low2 = self.low2(low1) + low3 = self.low3(low2) + up2 = self.up2(low3) + return self.merge(up1, up2) + +class kp(nn.Module): + def __init__( + self, n, nstack, dims, modules, out_dim, pre=None, cnv_dim=256, + make_tl_layer=make_tl_layer, make_br_layer=make_br_layer, + make_cnv_layer=make_cnv_layer, make_heat_layer=make_kp_layer, + make_tag_layer=make_kp_layer, make_regr_layer=make_kp_layer, + make_up_layer=make_layer, make_low_layer=make_layer, + make_hg_layer=make_layer, make_hg_layer_revr=make_layer_revr, + make_pool_layer=make_pool_layer, make_unpool_layer=make_unpool_layer, + make_merge_layer=make_merge_layer, make_inter_layer=make_inter_layer, + kp_layer=residual + ): + super(kp, self).__init__() + + self.nstack = nstack + self._decode = _decode + + curr_dim = dims[0] + + self.pre = nn.Sequential( + convolution(7, 3, 128, stride=2), + residual(3, 128, 256, stride=2) + ) if pre is None else pre + + self.kps = nn.ModuleList([ + kp_module( + n, dims, modules, layer=kp_layer, + make_up_layer=make_up_layer, + make_low_layer=make_low_layer, + make_hg_layer=make_hg_layer, + make_hg_layer_revr=make_hg_layer_revr, + make_pool_layer=make_pool_layer, + make_unpool_layer=make_unpool_layer, + make_merge_layer=make_merge_layer + ) for _ in range(nstack) + ]) + self.cnvs = nn.ModuleList([ + make_cnv_layer(curr_dim, cnv_dim) for _ in range(nstack) + ]) + + self.tl_cnvs = nn.ModuleList([ + make_tl_layer(cnv_dim) for _ in range(nstack) + ]) + self.br_cnvs = nn.ModuleList([ + make_br_layer(cnv_dim) for _ in range(nstack) + ]) + + ## keypoint heatmaps + self.tl_heats = nn.ModuleList([ + make_heat_layer(cnv_dim, curr_dim, out_dim) for _ in range(nstack) + ]) + self.br_heats = nn.ModuleList([ + make_heat_layer(cnv_dim, curr_dim, out_dim) for _ in range(nstack) + ]) + + ## tags + self.tl_tags = nn.ModuleList([ + make_tag_layer(cnv_dim, curr_dim, 1) for _ in range(nstack) + ]) + self.br_tags = nn.ModuleList([ + make_tag_layer(cnv_dim, curr_dim, 1) for _ in range(nstack) + ]) + + for tl_heat, br_heat in zip(self.tl_heats, self.br_heats): + tl_heat[-1].bias.data.fill_(-2.19) + br_heat[-1].bias.data.fill_(-2.19) + + self.inters = nn.ModuleList([ + make_inter_layer(curr_dim) for _ in range(nstack - 1) + ]) + + self.inters_ = nn.ModuleList([ + nn.Sequential( + nn.Conv2d(curr_dim, curr_dim, (1, 1), bias=False), + nn.BatchNorm2d(curr_dim) + ) for _ in range(nstack - 1) + ]) + self.cnvs_ = nn.ModuleList([ + nn.Sequential( + nn.Conv2d(cnv_dim, curr_dim, (1, 1), bias=False), + nn.BatchNorm2d(curr_dim) + ) for _ in range(nstack - 1) + ]) + + self.tl_regrs = nn.ModuleList([ + make_regr_layer(cnv_dim, curr_dim, 2) for _ in range(nstack) + ]) + self.br_regrs = nn.ModuleList([ + make_regr_layer(cnv_dim, curr_dim, 2) for _ in range(nstack) + ]) + + self.relu = nn.ReLU(inplace=True) + + def _train(self, *xs): + image = xs[0] + tl_inds = xs[1] + br_inds = xs[2] + + inter = self.pre(image) + outs = [] + + layers = zip( + self.kps, self.cnvs, + self.tl_cnvs, self.br_cnvs, + self.tl_heats, self.br_heats, + self.tl_tags, self.br_tags, + self.tl_regrs, self.br_regrs + ) + for ind, layer in enumerate(layers): + kp_, cnv_ = layer[0:2] + tl_cnv_, br_cnv_ = layer[2:4] + tl_heat_, br_heat_ = layer[4:6] + tl_tag_, br_tag_ = layer[6:8] + tl_regr_, br_regr_ = layer[8:10] + + kp = kp_(inter) + cnv = cnv_(kp) + + tl_cnv = tl_cnv_(cnv) + br_cnv = br_cnv_(cnv) + + tl_heat, br_heat = tl_heat_(tl_cnv), br_heat_(br_cnv) + tl_tag, br_tag = tl_tag_(tl_cnv), br_tag_(br_cnv) + tl_regr, br_regr = tl_regr_(tl_cnv), br_regr_(br_cnv) + + tl_tag = _tranpose_and_gather_feat(tl_tag, tl_inds) + br_tag = _tranpose_and_gather_feat(br_tag, br_inds) + tl_regr = _tranpose_and_gather_feat(tl_regr, tl_inds) + br_regr = _tranpose_and_gather_feat(br_regr, br_inds) + + outs += [tl_heat, br_heat, tl_tag, br_tag, tl_regr, br_regr] + + if ind < self.nstack - 1: + inter = self.inters_[ind](inter) + self.cnvs_[ind](cnv) + inter = self.relu(inter) + inter = self.inters[ind](inter) + return outs + + def _test(self, *xs, **kwargs): + image = xs[0] + + inter = self.pre(image) + outs = [] + + layers = zip( + self.kps, self.cnvs, + self.tl_cnvs, self.br_cnvs, + self.tl_heats, self.br_heats, + self.tl_tags, self.br_tags, + self.tl_regrs, self.br_regrs + ) + for ind, layer in enumerate(layers): + kp_, cnv_ = layer[0:2] + tl_cnv_, br_cnv_ = layer[2:4] + tl_heat_, br_heat_ = layer[4:6] + tl_tag_, br_tag_ = layer[6:8] + tl_regr_, br_regr_ = layer[8:10] + + kp = kp_(inter) + cnv = cnv_(kp) + + if ind == self.nstack - 1: + tl_cnv = tl_cnv_(cnv) + br_cnv = br_cnv_(cnv) + + tl_heat, br_heat = tl_heat_(tl_cnv), br_heat_(br_cnv) + tl_tag, br_tag = tl_tag_(tl_cnv), br_tag_(br_cnv) + tl_regr, br_regr = tl_regr_(tl_cnv), br_regr_(br_cnv) + + outs += [tl_heat, br_heat, tl_tag, br_tag, tl_regr, br_regr] + + if ind < self.nstack - 1: + inter = self.inters_[ind](inter) + self.cnvs_[ind](cnv) + inter = self.relu(inter) + inter = self.inters[ind](inter) + + return self._decode(*outs[-6:], **kwargs) + + def forward(self, *xs, **kwargs): + if len(xs) > 1: + return self._train(*xs, **kwargs) + return self._test(*xs, **kwargs) + +class AELoss(nn.Module): + def __init__(self, pull_weight=1, push_weight=1, regr_weight=1, focal_loss=_neg_loss): + super(AELoss, self).__init__() + + self.pull_weight = pull_weight + self.push_weight = push_weight + self.regr_weight = regr_weight + self.focal_loss = focal_loss + self.ae_loss = _ae_loss + self.regr_loss = _regr_loss + + def forward(self, outs, targets): + stride = 6 + + tl_heats = outs[0::stride] + br_heats = outs[1::stride] + tl_tags = outs[2::stride] + br_tags = outs[3::stride] + tl_regrs = outs[4::stride] + br_regrs = outs[5::stride] + + gt_tl_heat = targets[0] + gt_br_heat = targets[1] + gt_mask = targets[2] + gt_tl_regr = targets[3] + gt_br_regr = targets[4] + + # focal loss + focal_loss = 0 + + tl_heats = [_sigmoid(t) for t in tl_heats] + br_heats = [_sigmoid(b) for b in br_heats] + + focal_loss += self.focal_loss(tl_heats, gt_tl_heat) + focal_loss += self.focal_loss(br_heats, gt_br_heat) + + # tag loss + pull_loss = 0 + push_loss = 0 + + for tl_tag, br_tag in zip(tl_tags, br_tags): + pull, push = self.ae_loss(tl_tag, br_tag, gt_mask) + pull_loss += pull + push_loss += push + pull_loss = self.pull_weight * pull_loss + push_loss = self.push_weight * push_loss + + regr_loss = 0 + for tl_regr, br_regr in zip(tl_regrs, br_regrs): + regr_loss += self.regr_loss(tl_regr, gt_tl_regr, gt_mask) + regr_loss += self.regr_loss(br_regr, gt_br_regr, gt_mask) + regr_loss = self.regr_weight * regr_loss + + loss = (focal_loss + pull_loss + push_loss + regr_loss) / len(tl_heats) + return loss.unsqueeze(0) diff --git a/models/py_utils/kp_utils.py b/models/py_utils/kp_utils.py new file mode 100644 index 0000000..c739ad9 --- /dev/null +++ b/models/py_utils/kp_utils.py @@ -0,0 +1,505 @@ +import torch +import torch.nn as nn + +from .utils import convolution, residual + +class MergeUp(nn.Module): + def forward(self, up1, up2): + return up1 + up2 + +def make_merge_layer(dim): + return MergeUp() + +def make_tl_layer(dim): + return None + +def make_br_layer(dim): + return None + +def make_pool_layer(dim): + return nn.MaxPool2d(kernel_size=2, stride=2) + +def make_unpool_layer(dim): + return nn.Upsample(scale_factor=2) + +def make_kp_layer(cnv_dim, curr_dim, out_dim): + return nn.Sequential( + convolution(3, cnv_dim, curr_dim, with_bn=False), + nn.Conv2d(curr_dim, out_dim, (1, 1)) + ) + +def make_inter_layer(dim): + return residual(3, dim, dim) + +def make_cnv_layer(inp_dim, out_dim): + return convolution(3, inp_dim, out_dim) + +def _gather_feat(feat, ind, mask=None): + dim = feat.size(2) + ind = ind.unsqueeze(2).expand(ind.size(0), ind.size(1), dim) + feat = feat.gather(1, ind) + if mask is not None: + mask = mask.unsqueeze(2).expand_as(feat) + feat = feat[mask] + feat = feat.view(-1, dim) + return feat + +def _nms(heat, kernel=1): + pad = (kernel - 1) // 2 + + hmax = nn.functional.max_pool2d(heat, (kernel, kernel), stride=1, padding=pad) + keep = (hmax == heat).float() + return heat * keep + +def _left_aggregate(heat): + ''' + heat: batchsize x channels x h x w + ''' + shape = heat.shape + heat = heat.reshape(-1, heat.shape[3]) + ret = heat.clone() + for i in range(1, heat.shape[1]): + inds = (heat[:, i] >= heat[:, i -1]) + ret[:, i] += ret[:, i - 1] * inds.float() + return (ret - heat).reshape(shape) + +def _right_aggregate(heat): + ''' + heat: batchsize x channels x h x w + ''' + shape = heat.shape + heat = heat.reshape(-1, heat.shape[3]) + ret = heat.clone() + for i in range(heat.shape[1] - 2, -1, -1): + inds = (heat[:, i] >= heat[:, i +1]) + ret[:, i] += ret[:, i + 1] * inds.float() + return (ret - heat).reshape(shape) + +def _top_aggregate(heat): + ''' + heat: batchsize x channels x h x w + ''' + heat = heat.transpose(3, 2) + shape = heat.shape + heat = heat.reshape(-1, heat.shape[3]) + ret = heat.clone() + for i in range(1, heat.shape[1]): + inds = (heat[:, i] >= heat[:, i - 1]) + ret[:, i] += ret[:, i - 1] * inds.float() + return (ret - heat).reshape(shape).transpose(3, 2) + +def _bottom_aggregate(heat): + ''' + heat: batchsize x channels x h x w + ''' + heat = heat.transpose(3, 2) + shape = heat.shape + heat = heat.reshape(-1, heat.shape[3]) + ret = heat.clone() + for i in range(heat.shape[1] - 2, -1, -1): + inds = (heat[:, i] >= heat[:, i + 1]) + ret[:, i] += ret[:, i + 1] * inds.float() + return (ret - heat).reshape(shape).transpose(3, 2) + +def _h_aggregate(heat, aggr_weight=0.1): + return aggr_weight * _left_aggregate(heat) + \ + aggr_weight * _right_aggregate(heat) + heat + +def _v_aggregate(heat, aggr_weight=0.1): + return aggr_weight * _top_aggregate(heat) + \ + aggr_weight * _bottom_aggregate(heat) + heat + +def _tranpose_and_gather_feat(feat, ind): + feat = feat.permute(0, 2, 3, 1).contiguous() + feat = feat.view(feat.size(0), -1, feat.size(3)) + feat = _gather_feat(feat, ind) + return feat + +def _filter(heat, direction, val=0.1): + num_channels = heat.shape[1] + if direction == 'v': + kernel = torch.zeros((num_channels, num_channels, 3, 1)) + for i in range(num_channels): + kernel[i, i, 0, 0] = val + kernel[i, i, 1, 0] = 1 + kernel[i, i, 2, 0] = val + padding = (1, 0) + elif direction == 'h': + kernel = torch.zeros((num_channels, num_channels, 1, 3)) + for i in range(num_channels): + kernel[i, i, 0, 0] = val + kernel[i, i, 0, 1] = 1 + kernel[i, i, 0, 2] = val + padding = (0, 1) + else: + assert 0, direction + + heat = nn.functional.conv2d(heat, kernel.cuda(), padding=padding) + # heat[heat > 1] = 1 + return heat + +def _topk(scores, K=20): + batch, cat, height, width = scores.size() + + topk_scores, topk_inds = torch.topk(scores.view(batch, -1), K) + + topk_clses = (topk_inds / (height * width)).int() + + topk_inds = topk_inds % (height * width) + topk_ys = (topk_inds / width).int().float() + topk_xs = (topk_inds % width).int().float() + return topk_scores, topk_inds, topk_clses, topk_ys, topk_xs + +def _decode( + tl_heat, br_heat, tl_tag, br_tag, tl_regr, br_regr, + K=100, kernel=1, ae_threshold=1, num_dets=1000 +): + batch, cat, height, width = tl_heat.size() + + tl_heat = torch.sigmoid(tl_heat) + br_heat = torch.sigmoid(br_heat) + + # perform nms on heatmaps + tl_heat = _nms(tl_heat, kernel=kernel) + br_heat = _nms(br_heat, kernel=kernel) + + tl_scores, tl_inds, tl_clses, tl_ys, tl_xs = _topk(tl_heat, K=K) + br_scores, br_inds, br_clses, br_ys, br_xs = _topk(br_heat, K=K) + + tl_ys = tl_ys.view(batch, K, 1).expand(batch, K, K) + tl_xs = tl_xs.view(batch, K, 1).expand(batch, K, K) + br_ys = br_ys.view(batch, 1, K).expand(batch, K, K) + br_xs = br_xs.view(batch, 1, K).expand(batch, K, K) + + if tl_regr is not None and br_regr is not None: + tl_regr = _tranpose_and_gather_feat(tl_regr, tl_inds) + tl_regr = tl_regr.view(batch, K, 1, 2) + br_regr = _tranpose_and_gather_feat(br_regr, br_inds) + br_regr = br_regr.view(batch, 1, K, 2) + + tl_xs = tl_xs + tl_regr[..., 0] + tl_ys = tl_ys + tl_regr[..., 1] + br_xs = br_xs + br_regr[..., 0] + br_ys = br_ys + br_regr[..., 1] + + # all possible boxes based on top k corners (ignoring class) + bboxes = torch.stack((tl_xs, tl_ys, br_xs, br_ys), dim=3) + + tl_tag = _tranpose_and_gather_feat(tl_tag, tl_inds) + tl_tag = tl_tag.view(batch, K, 1) + br_tag = _tranpose_and_gather_feat(br_tag, br_inds) + br_tag = br_tag.view(batch, 1, K) + dists = torch.abs(tl_tag - br_tag) + + tl_scores = tl_scores.view(batch, K, 1).expand(batch, K, K) + br_scores = br_scores.view(batch, 1, K).expand(batch, K, K) + scores = (tl_scores + br_scores) / 2 + + # reject boxes based on classes + tl_clses = tl_clses.view(batch, K, 1).expand(batch, K, K) + br_clses = br_clses.view(batch, 1, K).expand(batch, K, K) + cls_inds = (tl_clses != br_clses) + + # reject boxes based on distances + dist_inds = (dists > ae_threshold) + + # reject boxes based on widths and heights + width_inds = (br_xs < tl_xs) + height_inds = (br_ys < tl_ys) + + scores[cls_inds] = -1 + scores[dist_inds] = -1 + scores[width_inds] = -1 + scores[height_inds] = -1 + + scores = scores.view(batch, -1) + scores, inds = torch.topk(scores, num_dets) + scores = scores.unsqueeze(2) + + bboxes = bboxes.view(batch, -1, 4) + bboxes = _gather_feat(bboxes, inds) + + clses = tl_clses.contiguous().view(batch, -1, 1) + clses = _gather_feat(clses, inds).float() + + tl_scores = tl_scores.contiguous().view(batch, -1, 1) + tl_scores = _gather_feat(tl_scores, inds).float() + br_scores = br_scores.contiguous().view(batch, -1, 1) + br_scores = _gather_feat(br_scores, inds).float() + + detections = torch.cat([bboxes, scores, tl_scores, br_scores, clses], dim=2) + return detections + +def _exct_decode( + t_heat, l_heat, b_heat, r_heat, ct_heat, + t_regr, l_regr, b_regr, r_regr, + K=40, kernel=3, + aggr_weight=0.1, scores_thresh=0.1, center_thresh=0.1,num_dets=1000 +): + batch, cat, height, width = t_heat.size() + + ''' + filter_kernel = 0.1 + t_heat = _filter(t_heat, direction='h', val=filter_kernel) + l_heat = _filter(l_heat, direction='v', val=filter_kernel) + b_heat = _filter(b_heat, direction='h', val=filter_kernel) + r_heat = _filter(r_heat, direction='v', val=filter_kernel) + ''' + + t_heat = torch.sigmoid(t_heat) + l_heat = torch.sigmoid(l_heat) + b_heat = torch.sigmoid(b_heat) + r_heat = torch.sigmoid(r_heat) + ct_heat = torch.sigmoid(ct_heat) + + if aggr_weight > 0: + t_heat = _h_aggregate(t_heat, aggr_weight=aggr_weight) + l_heat = _v_aggregate(l_heat, aggr_weight=aggr_weight) + b_heat = _h_aggregate(b_heat, aggr_weight=aggr_weight) + r_heat = _v_aggregate(r_heat, aggr_weight=aggr_weight) + + + # perform nms on heatmaps + t_heat = _nms(t_heat, kernel=kernel) + l_heat = _nms(l_heat, kernel=kernel) + b_heat = _nms(b_heat, kernel=kernel) + r_heat = _nms(r_heat, kernel=kernel) + + t_heat[t_heat > 1] = 1 + l_heat[l_heat > 1] = 1 + b_heat[b_heat > 1] = 1 + r_heat[r_heat > 1] = 1 + + t_scores, t_inds, t_clses, t_ys, t_xs = _topk(t_heat, K=K) + l_scores, l_inds, l_clses, l_ys, l_xs = _topk(l_heat, K=K) + b_scores, b_inds, b_clses, b_ys, b_xs = _topk(b_heat, K=K) + r_scores, r_inds, r_clses, r_ys, r_xs = _topk(r_heat, K=K) + + t_ys = t_ys.view(batch, K, 1, 1, 1).expand(batch, K, K, K, K) + t_xs = t_xs.view(batch, K, 1, 1, 1).expand(batch, K, K, K, K) + l_ys = l_ys.view(batch, 1, K, 1, 1).expand(batch, K, K, K, K) + l_xs = l_xs.view(batch, 1, K, 1, 1).expand(batch, K, K, K, K) + b_ys = b_ys.view(batch, 1, 1, K, 1).expand(batch, K, K, K, K) + b_xs = b_xs.view(batch, 1, 1, K, 1).expand(batch, K, K, K, K) + r_ys = r_ys.view(batch, 1, 1, 1, K).expand(batch, K, K, K, K) + r_xs = r_xs.view(batch, 1, 1, 1, K).expand(batch, K, K, K, K) + + t_clses = t_clses.view(batch, K, 1, 1, 1).expand(batch, K, K, K, K) + l_clses = l_clses.view(batch, 1, K, 1, 1).expand(batch, K, K, K, K) + b_clses = b_clses.view(batch, 1, 1, K, 1).expand(batch, K, K, K, K) + r_clses = r_clses.view(batch, 1, 1, 1, K).expand(batch, K, K, K, K) + box_ct_xs = ((l_xs + r_xs + 0.5) / 2).long() + box_ct_ys = ((t_ys + b_ys + 0.5) / 2).long() + ct_inds = t_clses.long() * (height * width) + box_ct_ys * width + box_ct_xs + ct_inds = ct_inds.view(batch, -1) + ct_heat = ct_heat.view(batch, -1, 1) + ct_scores = _gather_feat(ct_heat, ct_inds) + + t_scores = t_scores.view(batch, K, 1, 1, 1).expand(batch, K, K, K, K) + l_scores = l_scores.view(batch, 1, K, 1, 1).expand(batch, K, K, K, K) + b_scores = b_scores.view(batch, 1, 1, K, 1).expand(batch, K, K, K, K) + r_scores = r_scores.view(batch, 1, 1, 1, K).expand(batch, K, K, K, K) + ct_scores = ct_scores.view(batch, K, K, K, K) + scores = (t_scores + l_scores + b_scores + r_scores + 2 * ct_scores) / 6 + + # reject boxes based on classes + cls_inds = (t_clses != l_clses) + (t_clses != b_clses) + \ + (t_clses != r_clses) + cls_inds = (cls_inds > 0) + + top_inds = (t_ys > l_ys) + (t_ys > b_ys) + (t_ys > r_ys) + top_inds = (top_inds > 0) + left_inds = (l_xs > t_xs) + (l_xs > b_xs) + (l_xs > r_xs) + left_inds = (left_inds > 0) + bottom_inds = (b_ys < t_ys) + (b_ys < l_ys) + (b_ys < r_ys) + bottom_inds = (bottom_inds > 0) + right_inds = (r_xs < t_xs) + (r_xs < l_xs) + (r_xs < b_xs) + right_inds = (right_inds > 0) + + sc_inds = (t_scores < scores_thresh) + (l_scores < scores_thresh) + \ + (b_scores < scores_thresh) + (r_scores < scores_thresh) + \ + (ct_scores < center_thresh) + sc_inds = (sc_inds > 0) + + ''' + scores[sc_inds] = -1 + scores[cls_inds] = -1 + scores[top_inds] = -1 + scores[left_inds] = -1 + scores[bottom_inds] = -1 + scores[right_inds] = -1 + ''' + scores = scores - sc_inds.float() + scores = scores - cls_inds.float() + scores = scores - top_inds.float() + scores = scores - left_inds.float() + scores = scores - bottom_inds.float() + scores = scores - right_inds.float() + + + scores = scores.view(batch, -1) + scores, inds = torch.topk(scores, num_dets) + scores = scores.unsqueeze(2) + + if t_regr is not None and l_regr is not None \ + and b_regr is not None and r_regr is not None: + t_regr = _tranpose_and_gather_feat(t_regr, t_inds) + t_regr = t_regr.view(batch, K, 1, 1, 1, 2) + l_regr = _tranpose_and_gather_feat(l_regr, l_inds) + l_regr = l_regr.view(batch, 1, K, 1, 1, 2) + b_regr = _tranpose_and_gather_feat(b_regr, b_inds) + b_regr = b_regr.view(batch, 1, 1, K, 1, 2) + r_regr = _tranpose_and_gather_feat(r_regr, r_inds) + r_regr = r_regr.view(batch, 1, 1, 1, K, 2) + + t_xs = t_xs + t_regr[..., 0] + t_ys = t_ys + t_regr[..., 1] + l_xs = l_xs + l_regr[..., 0] + l_ys = l_ys + l_regr[..., 1] + b_xs = b_xs + b_regr[..., 0] + b_ys = b_ys + b_regr[..., 1] + r_xs = r_xs + r_regr[..., 0] + r_ys = r_ys + r_regr[..., 1] + else: + t_xs = t_xs + 0.5 + t_ys = t_ys + 0.5 + l_xs = l_xs + 0.5 + l_ys = l_ys + 0.5 + b_xs = b_xs + 0.5 + b_ys = b_ys + 0.5 + r_xs = r_xs + 0.5 + r_ys = r_ys + 0.5 + + bboxes = torch.stack((l_xs, t_ys, r_xs, b_ys), dim=5) + bboxes = bboxes.view(batch, -1, 4) + bboxes = _gather_feat(bboxes, inds) + + clses = t_clses.contiguous().view(batch, -1, 1) + clses = _gather_feat(clses, inds).float() + + t_xs = t_xs.contiguous().view(batch, -1, 1) + t_xs = _gather_feat(t_xs, inds).float() + t_ys = t_ys.contiguous().view(batch, -1, 1) + t_ys = _gather_feat(t_ys, inds).float() + l_xs = l_xs.contiguous().view(batch, -1, 1) + l_xs = _gather_feat(l_xs, inds).float() + l_ys = l_ys.contiguous().view(batch, -1, 1) + l_ys = _gather_feat(l_ys, inds).float() + b_xs = b_xs.contiguous().view(batch, -1, 1) + b_xs = _gather_feat(b_xs, inds).float() + b_ys = b_ys.contiguous().view(batch, -1, 1) + b_ys = _gather_feat(b_ys, inds).float() + r_xs = r_xs.contiguous().view(batch, -1, 1) + r_xs = _gather_feat(r_xs, inds).float() + r_ys = r_ys.contiguous().view(batch, -1, 1) + r_ys = _gather_feat(r_ys, inds).float() + + + detections = torch.cat([bboxes, scores, t_xs, t_ys, l_xs, l_ys, + b_xs, b_ys, r_xs, r_ys, clses], dim=2) + + return detections + +''' +# Faster but costs more memory +def _neg_loss(preds, gt): + pos_inds = gt.eq(1).float() + neg_inds = gt.lt(1).float() + + neg_weights = torch.pow(1 - gt, 4) + + loss = 0 + for pred in preds: + pos_loss = torch.log(pred) * torch.pow(1 - pred, 2) * pos_inds + neg_loss = torch.log(1 - pred) * torch.pow(pred, 2) * \ + neg_weights * neg_inds + + num_pos = pos_inds.float().sum() + pos_loss = pos_loss.sum() + neg_loss = neg_loss.sum() + + if num_pos == 0: + loss = loss - neg_loss + else: + loss = loss - (pos_loss + neg_loss) / num_pos + return loss +''' + +def _neg_loss(preds, gt): + pos_inds = gt.eq(1) + neg_inds = gt.lt(1) + + neg_weights = torch.pow(1 - gt[neg_inds], 4) + + loss = 0 + for pred in preds: + pos_pred = pred[pos_inds] + neg_pred = pred[neg_inds] + + pos_loss = torch.log(pos_pred) * torch.pow(1 - pos_pred, 2) + neg_loss = torch.log(1 - neg_pred) * torch.pow(neg_pred, 2) * neg_weights + + num_pos = pos_inds.float().sum() + pos_loss = pos_loss.sum() + neg_loss = neg_loss.sum() + + if pos_pred.nelement() == 0: + loss = loss - neg_loss + else: + loss = loss - (pos_loss + neg_loss) / num_pos + return loss + + +def _sigmoid(x): + x = torch.clamp(x.sigmoid_(), min=1e-4, max=1-1e-4) + return x + +def _ae_loss(tag0, tag1, mask): + num = mask.sum(dim=1, keepdim=True).float() + tag0 = tag0.squeeze() + tag1 = tag1.squeeze() + + tag_mean = (tag0 + tag1) / 2 + + tag0 = torch.pow(tag0 - tag_mean, 2) / (num + 1e-4) + tag0 = tag0[mask].sum() + tag1 = torch.pow(tag1 - tag_mean, 2) / (num + 1e-4) + tag1 = tag1[mask].sum() + pull = tag0 + tag1 + + mask = mask.unsqueeze(1) + mask.unsqueeze(2) + mask = mask.eq(2) + num = num.unsqueeze(2) + num2 = (num - 1) * num + dist = tag_mean.unsqueeze(1) - tag_mean.unsqueeze(2) + dist = 1 - torch.abs(dist) + dist = nn.functional.relu(dist, inplace=True) + dist = dist - 1 / (num + 1e-4) + dist = dist / (num2 + 1e-4) + dist = dist[mask] + push = dist.sum() + return pull, push + +''' +def _regr_loss(regr, gt_regr, mask): + num = mask.float().sum() + mask = mask.unsqueeze(2).expand_as(gt_regr).float() + + regr = regr * mask + gt_regr = gt_regr * mask + + regr_loss = nn.functional.smooth_l1_loss(regr, gt_regr, size_average=False) + regr_loss = regr_loss / (num + 1e-4) + return regr_loss +''' + +def _regr_loss(regr, gt_regr, mask): + num = mask.float().sum() + mask = mask.unsqueeze(2).expand_as(gt_regr) + + regr = regr[mask] + gt_regr = gt_regr[mask] + + regr_loss = nn.functional.smooth_l1_loss(regr, gt_regr, size_average=False) + regr_loss = regr_loss / (num + 1e-4) + return regr_loss diff --git a/models/py_utils/scatter_gather.py b/models/py_utils/scatter_gather.py new file mode 100644 index 0000000..9a46058 --- /dev/null +++ b/models/py_utils/scatter_gather.py @@ -0,0 +1,38 @@ +import torch +from torch.autograd import Variable +from torch.nn.parallel._functions import Scatter, Gather + + +def scatter(inputs, target_gpus, dim=0, chunk_sizes=None): + r""" + Slices variables into approximately equal chunks and + distributes them across given GPUs. Duplicates + references to objects that are not variables. Does not + support Tensors. + """ + def scatter_map(obj): + if isinstance(obj, Variable): + return Scatter.apply(target_gpus, chunk_sizes, dim, obj) + assert not torch.is_tensor(obj), "Tensors not supported in scatter." + if isinstance(obj, tuple): + return list(zip(*map(scatter_map, obj))) + if isinstance(obj, list): + return list(map(list, zip(*map(scatter_map, obj)))) + if isinstance(obj, dict): + return list(map(type(obj), zip(*map(scatter_map, obj.items())))) + return [obj for targets in target_gpus] + + return scatter_map(inputs) + + +def scatter_kwargs(inputs, kwargs, target_gpus, dim=0, chunk_sizes=None): + r"""Scatter with support for kwargs dictionary""" + inputs = scatter(inputs, target_gpus, dim, chunk_sizes) if inputs else [] + kwargs = scatter(kwargs, target_gpus, dim, chunk_sizes) if kwargs else [] + if len(inputs) < len(kwargs): + inputs.extend([() for _ in range(len(kwargs) - len(inputs))]) + elif len(kwargs) < len(inputs): + kwargs.extend([{} for _ in range(len(inputs) - len(kwargs))]) + inputs = tuple(inputs) + kwargs = tuple(kwargs) + return inputs, kwargs diff --git a/models/py_utils/utils.py b/models/py_utils/utils.py new file mode 100644 index 0000000..b4d191d --- /dev/null +++ b/models/py_utils/utils.py @@ -0,0 +1,74 @@ +import torch +import torch.nn as nn + +class convolution(nn.Module): + def __init__(self, k, inp_dim, out_dim, stride=1, with_bn=True): + super(convolution, self).__init__() + + pad = (k - 1) // 2 + self.conv = nn.Conv2d(inp_dim, out_dim, (k, k), padding=(pad, pad), stride=(stride, stride), bias=not with_bn) + self.bn = nn.BatchNorm2d(out_dim) if with_bn else nn.Sequential() + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + conv = self.conv(x) + bn = self.bn(conv) + relu = self.relu(bn) + return relu + +class fully_connected(nn.Module): + def __init__(self, inp_dim, out_dim, with_bn=True): + super(fully_connected, self).__init__() + self.with_bn = with_bn + + self.linear = nn.Linear(inp_dim, out_dim) + if self.with_bn: + self.bn = nn.BatchNorm1d(out_dim) + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + linear = self.linear(x) + bn = self.bn(linear) if self.with_bn else linear + relu = self.relu(bn) + return relu + +class residual(nn.Module): + def __init__(self, k, inp_dim, out_dim, stride=1, with_bn=True): + super(residual, self).__init__() + + self.conv1 = nn.Conv2d(inp_dim, out_dim, (3, 3), padding=(1, 1), stride=(stride, stride), bias=False) + self.bn1 = nn.BatchNorm2d(out_dim) + self.relu1 = nn.ReLU(inplace=True) + + self.conv2 = nn.Conv2d(out_dim, out_dim, (3, 3), padding=(1, 1), bias=False) + self.bn2 = nn.BatchNorm2d(out_dim) + + self.skip = nn.Sequential( + nn.Conv2d(inp_dim, out_dim, (1, 1), stride=(stride, stride), bias=False), + nn.BatchNorm2d(out_dim) + ) if stride != 1 or inp_dim != out_dim else nn.Sequential() + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + conv1 = self.conv1(x) + bn1 = self.bn1(conv1) + relu1 = self.relu1(bn1) + + conv2 = self.conv2(relu1) + bn2 = self.bn2(conv2) + + skip = self.skip(x) + return self.relu(bn2 + skip) + +def make_layer(k, inp_dim, out_dim, modules, layer=convolution, **kwargs): + layers = [layer(k, inp_dim, out_dim, **kwargs)] + for _ in range(1, modules): + layers.append(layer(k, out_dim, out_dim, **kwargs)) + return nn.Sequential(*layers) + +def make_layer_revr(k, inp_dim, out_dim, modules, layer=convolution, **kwargs): + layers = [] + for _ in range(modules - 1): + layers.append(layer(k, inp_dim, inp_dim, **kwargs)) + layers.append(layer(k, inp_dim, out_dim, **kwargs)) + return nn.Sequential(*layers) diff --git a/nnet/__init__.py b/nnet/__init__.py new file mode 100755 index 0000000..e69de29 diff --git a/nnet/py_factory.py b/nnet/py_factory.py new file mode 100755 index 0000000..87a1e41 --- /dev/null +++ b/nnet/py_factory.py @@ -0,0 +1,124 @@ +import os +import torch +import importlib +import torch.nn as nn + +from config import system_configs +from models.py_utils.data_parallel import DataParallel + +torch.manual_seed(317) + +class Network(nn.Module): + def __init__(self, model, loss): + super(Network, self).__init__() + + self.model = model + self.loss = loss + + def forward(self, xs, ys, **kwargs): + preds = self.model(*xs, **kwargs) + loss = self.loss(preds, ys, **kwargs) + return loss + +# for model backward compatibility +# previously model was wrapped by DataParallel module +class DummyModule(nn.Module): + def __init__(self, model): + super(DummyModule, self).__init__() + self.module = model + + def forward(self, *xs, **kwargs): + return self.module(*xs, **kwargs) + +class NetworkFactory(object): + def __init__(self, db): + super(NetworkFactory, self).__init__() + + module_file = "models.{}".format(system_configs.snapshot_name) + print("module_file: {}".format(module_file)) + nnet_module = importlib.import_module(module_file) + + self.model = DummyModule(nnet_module.model(db)) + self.loss = nnet_module.loss + self.network = Network(self.model, self.loss) + self.network = DataParallel(self.network, chunk_sizes=system_configs.chunk_sizes) + + total_params = 0 + for params in self.model.parameters(): + num_params = 1 + for x in params.size(): + num_params *= x + total_params += num_params + print("total parameters: {}".format(total_params)) + + if system_configs.opt_algo == "adam": + self.optimizer = torch.optim.Adam( + filter(lambda p: p.requires_grad, self.model.parameters()) + ) + elif system_configs.opt_algo == "sgd": + self.optimizer = torch.optim.SGD( + filter(lambda p: p.requires_grad, self.model.parameters()), + lr=system_configs.learning_rate, + momentum=0.9, weight_decay=0.0001 + ) + else: + raise ValueError("unknown optimizer") + + def cuda(self): + self.model.cuda() + + def train_mode(self): + self.network.train() + + def eval_mode(self): + self.network.eval() + + def train(self, xs, ys, **kwargs): + xs = [x.cuda(non_blocking=True) for x in xs] + ys = [y.cuda(non_blocking=True) for y in ys] + + self.optimizer.zero_grad() + loss = self.network(xs, ys) + loss = loss.mean() + loss.backward() + self.optimizer.step() + return loss + + def validate(self, xs, ys, **kwargs): + with torch.no_grad(): + xs = [x.cuda(non_blocking=True) for x in xs] + ys = [y.cuda(non_blocking=True) for y in ys] + + loss = self.network(xs, ys) + loss = loss.mean() + return loss + + def test(self, xs, **kwargs): + with torch.no_grad(): + xs = [x.cuda(non_blocking=True) for x in xs] + return self.model(*xs, **kwargs) + + def set_lr(self, lr): + print("setting learning rate to: {}".format(lr)) + for param_group in self.optimizer.param_groups: + param_group["lr"] = lr + + def load_pretrained_params(self, pretrained_model): + print("loading from {}".format(pretrained_model)) + with open(pretrained_model, "rb") as f: + params = torch.load(f) + self.model.load_state_dict(params, strict=False) + + def load_params(self, iteration): + cache_file = system_configs.snapshot_file.format(iteration) + print("loading model from {}".format(cache_file)) + with open(cache_file, "rb") as f: + params = torch.load(f) + self.model.load_state_dict(params) + + def save_params(self, iteration): + cache_file = system_configs.snapshot_file.format(iteration) + print("saving model to {}".format(cache_file)) + with open(cache_file, "wb") as f: + params = self.model.state_dict() + torch.save(params, f) diff --git a/readme/center.png b/readme/center.png new file mode 100644 index 0000000..18d050a Binary files /dev/null and b/readme/center.png differ diff --git a/readme/extreme.png b/readme/extreme.png new file mode 100644 index 0000000..fa75f66 Binary files /dev/null and b/readme/extreme.png differ diff --git a/readme/mask.png b/readme/mask.png new file mode 100644 index 0000000..812fa27 Binary files /dev/null and b/readme/mask.png differ diff --git a/readme/octagon.png b/readme/octagon.png new file mode 100644 index 0000000..1148624 Binary files /dev/null and b/readme/octagon.png differ diff --git a/readme/teaser.png b/readme/teaser.png new file mode 100644 index 0000000..1938a93 Binary files /dev/null and b/readme/teaser.png differ diff --git a/sample/__init__.py b/sample/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/sample/coco.py b/sample/coco.py new file mode 100755 index 0000000..e372268 --- /dev/null +++ b/sample/coco.py @@ -0,0 +1,182 @@ +import cv2 +import math +import numpy as np +import torch +import random +import string + +from config import system_configs +from utils import crop_image, normalize_, color_jittering_, lighting_ +from .utils import random_crop, draw_gaussian, gaussian_radius + +def _full_image_crop(image, detections): + detections = detections.copy() + height, width = image.shape[0:2] + + max_hw = max(height, width) + center = [height // 2, width // 2] + size = [max_hw, max_hw] + + image, border, offset = crop_image(image, center, size) + detections[:, 0:4:2] += border[2] + detections[:, 1:4:2] += border[0] + return image, detections + +def _resize_image(image, detections, size): + detections = detections.copy() + height, width = image.shape[0:2] + new_height, new_width = size + + image = cv2.resize(image, (new_width, new_height)) + + height_ratio = new_height / height + width_ratio = new_width / width + detections[:, 0:4:2] *= width_ratio + detections[:, 1:4:2] *= height_ratio + return image, detections + +def _clip_detections(image, detections): + detections = detections.copy() + height, width = image.shape[0:2] + + detections[:, 0:4:2] = np.clip(detections[:, 0:4:2], 0, width - 1) + detections[:, 1:4:2] = np.clip(detections[:, 1:4:2], 0, height - 1) + keep_inds = ((detections[:, 2] - detections[:, 0]) > 0) & \ + ((detections[:, 3] - detections[:, 1]) > 0) + detections = detections[keep_inds] + return detections + +def kp_detection(db, k_ind, data_aug, debug): + data_rng = system_configs.data_rng + batch_size = system_configs.batch_size + + categories = db.configs["categories"] + input_size = db.configs["input_size"] + output_size = db.configs["output_sizes"][0] + + border = db.configs["border"] + lighting = db.configs["lighting"] + rand_crop = db.configs["rand_crop"] + rand_color = db.configs["rand_color"] + rand_scales = db.configs["rand_scales"] + gaussian_bump = db.configs["gaussian_bump"] + gaussian_iou = db.configs["gaussian_iou"] + gaussian_rad = db.configs["gaussian_radius"] + + max_tag_len = 128 + + # allocating memory + images = np.zeros((batch_size, 3, input_size[0], input_size[1]), dtype=np.float32) + tl_heatmaps = np.zeros((batch_size, categories, output_size[0], output_size[1]), dtype=np.float32) + br_heatmaps = np.zeros((batch_size, categories, output_size[0], output_size[1]), dtype=np.float32) + tl_regrs = np.zeros((batch_size, max_tag_len, 2), dtype=np.float32) + br_regrs = np.zeros((batch_size, max_tag_len, 2), dtype=np.float32) + tl_tags = np.zeros((batch_size, max_tag_len), dtype=np.int64) + br_tags = np.zeros((batch_size, max_tag_len), dtype=np.int64) + tag_masks = np.zeros((batch_size, max_tag_len), dtype=np.uint8) + tag_lens = np.zeros((batch_size, ), dtype=np.int32) + + db_size = db.db_inds.size + for b_ind in range(batch_size): + if not debug and k_ind == 0: + db.shuffle_inds() + + db_ind = db.db_inds[k_ind] + k_ind = (k_ind + 1) % db_size + + # reading image + image_file = db.image_file(db_ind) + image = cv2.imread(image_file) + + # reading detections + detections = db.detections(db_ind) + + # cropping an image randomly + if not debug and rand_crop: + image, detections = random_crop(image, detections, rand_scales, input_size, border=border) + else: + image, detections = _full_image_crop(image, detections) + + image, detections = _resize_image(image, detections, input_size) + detections = _clip_detections(image, detections) + + width_ratio = output_size[1] / input_size[1] + height_ratio = output_size[0] / input_size[0] + + # flipping an image randomly + if not debug and np.random.uniform() > 0.5: + image[:] = image[:, ::-1, :] + width = image.shape[1] + detections[:, [0, 2]] = width - detections[:, [2, 0]] - 1 + + if not debug: + image = image.astype(np.float32) / 255. + if rand_color: + color_jittering_(data_rng, image) + if lighting: + lighting_(data_rng, image, 0.1, db.eig_val, db.eig_vec) + normalize_(image, db.mean, db.std) + images[b_ind] = image.transpose((2, 0, 1)) + + for ind, detection in enumerate(detections): + category = int(detection[-1]) - 1 + + xtl, ytl = detection[0], detection[1] + xbr, ybr = detection[2], detection[3] + + fxtl = (xtl * width_ratio) + fytl = (ytl * height_ratio) + fxbr = (xbr * width_ratio) + fybr = (ybr * height_ratio) + + xtl = int(fxtl) + ytl = int(fytl) + xbr = int(fxbr) + ybr = int(fybr) + + if gaussian_bump: + width = detection[2] - detection[0] + height = detection[3] - detection[1] + + width = math.ceil(width * width_ratio) + height = math.ceil(height * height_ratio) + + if gaussian_rad == -1: + radius = gaussian_radius((height, width), gaussian_iou) + radius = max(0, int(radius)) + else: + radius = gaussian_rad + + draw_gaussian(tl_heatmaps[b_ind, category], [xtl, ytl], radius) + draw_gaussian(br_heatmaps[b_ind, category], [xbr, ybr], radius) + else: + tl_heatmaps[b_ind, category, ytl, xtl] = 1 + br_heatmaps[b_ind, category, ybr, xbr] = 1 + + tag_ind = tag_lens[b_ind] + tl_regrs[b_ind, tag_ind, :] = [fxtl - xtl, fytl - ytl] + br_regrs[b_ind, tag_ind, :] = [fxbr - xbr, fybr - ybr] + tl_tags[b_ind, tag_ind] = ytl * output_size[1] + xtl + br_tags[b_ind, tag_ind] = ybr * output_size[1] + xbr + tag_lens[b_ind] += 1 + + for b_ind in range(batch_size): + tag_len = tag_lens[b_ind] + tag_masks[b_ind, :tag_len] = 1 + + images = torch.from_numpy(images) + tl_heatmaps = torch.from_numpy(tl_heatmaps) + br_heatmaps = torch.from_numpy(br_heatmaps) + tl_regrs = torch.from_numpy(tl_regrs) + br_regrs = torch.from_numpy(br_regrs) + tl_tags = torch.from_numpy(tl_tags) + br_tags = torch.from_numpy(br_tags) + tag_masks = torch.from_numpy(tag_masks) + + return { + "xs": [images, tl_tags, br_tags], + "ys": [tl_heatmaps, br_heatmaps, tag_masks, tl_regrs, br_regrs] + }, k_ind + +def sample_data(db, k_ind, data_aug=True, debug=False): + return globals()[system_configs.sampling_function](db, k_ind, data_aug, debug) diff --git a/sample/coco_extreme.py b/sample/coco_extreme.py new file mode 100755 index 0000000..1dba5a8 --- /dev/null +++ b/sample/coco_extreme.py @@ -0,0 +1,245 @@ +import cv2 +import math +import numpy as np +import torch +import random +import string + +from config import system_configs +from utils import crop_image, normalize_, color_jittering_, lighting_ +from .utils import random_crop_pts, draw_gaussian, gaussian_radius +from utils.debugger import Debugger + +def _resize_image_pts(image, detections, extreme_pts, size): + detections = detections.copy() + height, width = image.shape[0:2] + new_height, new_width = size + + image = cv2.resize(image, (new_width, new_height)) + + height_ratio = new_height / height + width_ratio = new_width / width + detections[:, 0:4:2] *= width_ratio + detections[:, 1:4:2] *= height_ratio + extreme_pts[:, :, 0] *= width_ratio + extreme_pts[:, :, 1] *= height_ratio + return image, detections, extreme_pts + +def _clip_detections_pts(image, detections, extreme_pts): + detections = detections.copy() + height, width = image.shape[0:2] + + detections[:, 0:4:2] = np.clip(detections[:, 0:4:2], 0, width - 1) + detections[:, 1:4:2] = np.clip(detections[:, 1:4:2], 0, height - 1) + extreme_pts[:, :, 0] = np.clip(extreme_pts[:, :, 0], 0, width - 1) + extreme_pts[:, :, 1] = np.clip(extreme_pts[:, :, 1], 0, height - 1) + keep_inds = ((detections[:, 2] - detections[:, 0]) > 0) & \ + ((detections[:, 3] - detections[:, 1]) > 0) + detections = detections[keep_inds] + extreme_pts = extreme_pts[keep_inds] + return detections, extreme_pts + +def kp_detection(db, k_ind, data_aug, debug): + data_rng = system_configs.data_rng + batch_size = system_configs.batch_size + + categories = db.configs["categories"] + input_size = db.configs["input_size"] + output_size = db.configs["output_sizes"][0] + + border = db.configs["border"] + lighting = db.configs["lighting"] + rand_crop = db.configs["rand_crop"] + rand_color = db.configs["rand_color"] + rand_scales = db.configs["rand_scales"] + gaussian_bump = db.configs["gaussian_bump"] + gaussian_iou = db.configs["gaussian_iou"] + gaussian_rad = db.configs["gaussian_radius"] + + max_tag_len = 128 + + # allocating memory + images = np.zeros((batch_size, 3, input_size[0], input_size[1]), dtype=np.float32) + t_heatmaps = np.zeros((batch_size, categories, output_size[0], output_size[1]), dtype=np.float32) + l_heatmaps = np.zeros((batch_size, categories, output_size[0], output_size[1]), dtype=np.float32) + b_heatmaps = np.zeros((batch_size, categories, output_size[0], output_size[1]), dtype=np.float32) + r_heatmaps = np.zeros((batch_size, categories, output_size[0], output_size[1]), dtype=np.float32) + ct_heatmaps = np.zeros((batch_size, categories, output_size[0], output_size[1]), dtype=np.float32) + t_regrs = np.zeros((batch_size, max_tag_len, 2), dtype=np.float32) + l_regrs = np.zeros((batch_size, max_tag_len, 2), dtype=np.float32) + b_regrs = np.zeros((batch_size, max_tag_len, 2), dtype=np.float32) + r_regrs = np.zeros((batch_size, max_tag_len, 2), dtype=np.float32) + t_tags = np.zeros((batch_size, max_tag_len), dtype=np.int64) + l_tags = np.zeros((batch_size, max_tag_len), dtype=np.int64) + b_tags = np.zeros((batch_size, max_tag_len), dtype=np.int64) + r_tags = np.zeros((batch_size, max_tag_len), dtype=np.int64) + ct_tags = np.zeros((batch_size, max_tag_len), dtype=np.int64) + tag_masks = np.zeros((batch_size, max_tag_len), dtype=np.uint8) + tag_lens = np.zeros((batch_size, ), dtype=np.int32) + + db_size = db.db_inds.size + for b_ind in range(batch_size): + if not debug and k_ind == 0: + db.shuffle_inds() + + db_ind = db.db_inds[k_ind] + k_ind = (k_ind + 1) % db_size + + # reading image + image_file = db.image_file(db_ind) + image = cv2.imread(image_file) + + # reading detections + detections, extreme_pts = db.detections(db_ind) + + # cropping an image randomly + if rand_crop: + image, detections, extreme_pts = random_crop_pts( + image, detections, extreme_pts, + rand_scales, input_size, border=border) + else: + assert 0 + # image, detections = _full_image_crop(image, detections) + + image, detections, extreme_pts = _resize_image_pts( + image, detections, extreme_pts, input_size) + detections, extreme_pts = _clip_detections_pts( + image, detections, extreme_pts) + + width_ratio = output_size[1] / input_size[1] + height_ratio = output_size[0] / input_size[0] + + # flipping an image randomly + if np.random.uniform() > 0.5: + image[:] = image[:, ::-1, :] + width = image.shape[1] + detections[:, [0, 2]] = width - detections[:, [2, 0]] - 1 + extreme_pts[:, :, 0] = width - extreme_pts[:, :, 0] - 1 + extreme_pts[:, 1, :], extreme_pts[:, 3, :] = \ + extreme_pts[:, 3, :].copy(), extreme_pts[:, 1, :].copy() + + image = image.astype(np.float32) / 255. + if not debug: + if rand_color: + color_jittering_(data_rng, image) + if lighting: + lighting_(data_rng, image, 0.1, db.eig_val, db.eig_vec) + normalize_(image, db.mean, db.std) + images[b_ind] = image.transpose((2, 0, 1)) + + for ind, detection in enumerate(detections): + category = int(detection[-1]) - 1 + extreme_pt = extreme_pts[ind] + + xt, yt = extreme_pt[0, 0], extreme_pt[0, 1] + xl, yl = extreme_pt[1, 0], extreme_pt[1, 1] + xb, yb = extreme_pt[2, 0], extreme_pt[2, 1] + xr, yr = extreme_pt[3, 0], extreme_pt[3, 1] + xct = (xl + xr) / 2 + yct = (yt + yb) / 2 + + fxt = (xt * width_ratio) + fyt = (yt * height_ratio) + fxl = (xl * width_ratio) + fyl = (yl * height_ratio) + fxb = (xb * width_ratio) + fyb = (yb * height_ratio) + fxr = (xr * width_ratio) + fyr = (yr * height_ratio) + fxct = (xct * width_ratio) + fyct = (yct * height_ratio) + + xt = int(fxt) + yt = int(fyt) + xl = int(fxl) + yl = int(fyl) + xb = int(fxb) + yb = int(fyb) + xr = int(fxr) + yr = int(fyr) + xct = int(fxct) + yct = int(fyct) + + if gaussian_bump: + width = detection[2] - detection[0] + height = detection[3] - detection[1] + + width = math.ceil(width * width_ratio) + height = math.ceil(height * height_ratio) + + if gaussian_rad == -1: + radius = gaussian_radius((height, width), gaussian_iou) + radius = max(0, int(radius)) + else: + radius = gaussian_rad + draw_gaussian(t_heatmaps[b_ind, category], [xt, yt], radius) + draw_gaussian(l_heatmaps[b_ind, category], [xl, yl], radius) + draw_gaussian(b_heatmaps[b_ind, category], [xb, yb], radius) + draw_gaussian(r_heatmaps[b_ind, category], [xr, yr], radius) + draw_gaussian(ct_heatmaps[b_ind, category], [xct, yct], radius) + else: + t_heatmaps[b_ind, category, yt, xt] = 1 + l_heatmaps[b_ind, category, yl, xl] = 1 + b_heatmaps[b_ind, category, yb, xb] = 1 + r_heatmaps[b_ind, category, yr, xr] = 1 + + tag_ind = tag_lens[b_ind] + t_regrs[b_ind, tag_ind, :] = [fxt - xt, fyt - yt] + l_regrs[b_ind, tag_ind, :] = [fxl - xl, fyl - yl] + b_regrs[b_ind, tag_ind, :] = [fxb - xb, fyb - yb] + r_regrs[b_ind, tag_ind, :] = [fxr - xr, fyr - yr] + t_tags[b_ind, tag_ind] = yt * output_size[1] + xt + l_tags[b_ind, tag_ind] = yl * output_size[1] + xl + b_tags[b_ind, tag_ind] = yb * output_size[1] + xb + r_tags[b_ind, tag_ind] = yr * output_size[1] + xr + ct_tags[b_ind, tag_ind] = yct * output_size[1] + xct + tag_lens[b_ind] += 1 + + for b_ind in range(batch_size): + tag_len = tag_lens[b_ind] + tag_masks[b_ind, :tag_len] = 1 + + if debug: + debugger = Debugger(num_classes=80) + t_hm = debugger.gen_colormap(t_heatmaps[0]) + l_hm = debugger.gen_colormap(l_heatmaps[0]) + b_hm = debugger.gen_colormap(b_heatmaps[0]) + r_hm = debugger.gen_colormap(r_heatmaps[0]) + ct_hm = debugger.gen_colormap(ct_heatmaps[0]) + img = images[0] * db.std.reshape(3, 1, 1) + db.mean.reshape(3, 1, 1) + img = (img * 255).astype(np.uint8).transpose(1, 2, 0) + debugger.add_blend_img(img, t_hm, 't_hm') + debugger.add_blend_img(img, l_hm, 'l_hm') + debugger.add_blend_img(img, b_hm, 'b_hm') + debugger.add_blend_img(img, r_hm, 'r_hm') + debugger.add_blend_img( + img, np.maximum(np.maximum(t_hm, l_hm), + np.maximum(b_hm, r_hm)), 'extreme') + debugger.add_blend_img(img, ct_hm, 'center') + debugger.show_all_imgs(pause=True) + + images = torch.from_numpy(images) + t_heatmaps = torch.from_numpy(t_heatmaps) + l_heatmaps = torch.from_numpy(l_heatmaps) + b_heatmaps = torch.from_numpy(b_heatmaps) + r_heatmaps = torch.from_numpy(r_heatmaps) + ct_heatmaps = torch.from_numpy(ct_heatmaps) + t_regrs = torch.from_numpy(t_regrs) + l_regrs = torch.from_numpy(l_regrs) + b_regrs = torch.from_numpy(b_regrs) + r_regrs = torch.from_numpy(r_regrs) + t_tags = torch.from_numpy(t_tags) + l_tags = torch.from_numpy(l_tags) + b_tags = torch.from_numpy(b_tags) + r_tags = torch.from_numpy(r_tags) + ct_tags = torch.from_numpy(ct_tags) + tag_masks = torch.from_numpy(tag_masks) + + return { + "xs": [images, t_tags, l_tags, b_tags, r_tags, ct_tags], + "ys": [t_heatmaps, l_heatmaps, b_heatmaps, r_heatmaps, ct_heatmaps, + tag_masks, t_regrs, l_regrs, b_regrs, r_regrs] + }, k_ind + +def sample_data(db, k_ind, data_aug=True, debug=False): + return globals()[system_configs.sampling_function](db, k_ind, data_aug, debug) diff --git a/sample/utils.py b/sample/utils.py new file mode 100644 index 0000000..1f5eb14 --- /dev/null +++ b/sample/utils.py @@ -0,0 +1,134 @@ +import cv2 +import numpy as np + +def gaussian2D(shape, sigma=1): + m, n = [(ss - 1.) / 2. for ss in shape] + y, x = np.ogrid[-m:m+1,-n:n+1] + + h = np.exp(-(x * x + y * y) / (2 * sigma * sigma)) + h[h < np.finfo(h.dtype).eps * h.max()] = 0 + return h + +def draw_gaussian(heatmap, center, radius, k=1): + diameter = 2 * radius + 1 + gaussian = gaussian2D((diameter, diameter), sigma=diameter / 6) + + x, y = center + + height, width = heatmap.shape[0:2] + + left, right = min(x, radius), min(width - x, radius + 1) + top, bottom = min(y, radius), min(height - y, radius + 1) + + masked_heatmap = heatmap[y - top:y + bottom, x - left:x + right] + masked_gaussian = gaussian[radius - top:radius + bottom, radius - left:radius + right] + np.maximum(masked_heatmap, masked_gaussian * k, out=masked_heatmap) + +def gaussian_radius(det_size, min_overlap): + height, width = det_size + + a1 = 1 + b1 = (height + width) + c1 = width * height * (1 - min_overlap) / (1 + min_overlap) + sq1 = np.sqrt(b1 ** 2 - 4 * a1 * c1) + r1 = (b1 + sq1) / 2 + + a2 = 4 + b2 = 2 * (height + width) + c2 = (1 - min_overlap) * width * height + sq2 = np.sqrt(b2 ** 2 - 4 * a2 * c2) + r2 = (b2 + sq2) / 2 + + a3 = 4 * min_overlap + b3 = -2 * min_overlap * (height + width) + c3 = (min_overlap - 1) * width * height + sq3 = np.sqrt(b3 ** 2 - 4 * a3 * c3) + r3 = (b3 + sq3) / 2 + return min(r1, r2, r3) + +def _get_border(border, size): + i = 1 + while size - border // i <= border // i: + i *= 2 + return border // i + +def random_crop(image, detections, random_scales, view_size, border=64): + view_height, view_width = view_size + image_height, image_width = image.shape[0:2] + + scale = np.random.choice(random_scales) + height = int(view_height * scale) + width = int(view_width * scale) + + cropped_image = np.zeros((height, width, 3), dtype=image.dtype) + + w_border = _get_border(border, image_width) + h_border = _get_border(border, image_height) + + ctx = np.random.randint(low=w_border, high=image_width - w_border) + cty = np.random.randint(low=h_border, high=image_height - h_border) + + x0, x1 = max(ctx - width // 2, 0), min(ctx + width // 2, image_width) + y0, y1 = max(cty - height // 2, 0), min(cty + height // 2, image_height) + + left_w, right_w = ctx - x0, x1 - ctx + top_h, bottom_h = cty - y0, y1 - cty + + # crop image + cropped_ctx, cropped_cty = width // 2, height // 2 + x_slice = slice(cropped_ctx - left_w, cropped_ctx + right_w) + y_slice = slice(cropped_cty - top_h, cropped_cty + bottom_h) + cropped_image[y_slice, x_slice, :] = image[y0:y1, x0:x1, :] + + # crop detections + cropped_detections = detections.copy() + cropped_detections[:, 0:4:2] -= x0 + cropped_detections[:, 1:4:2] -= y0 + cropped_detections[:, 0:4:2] += cropped_ctx - left_w + cropped_detections[:, 1:4:2] += cropped_cty - top_h + + return cropped_image, cropped_detections + +def random_crop_pts(image, detections, extreme_pts, + random_scales, view_size, border=64): + view_height, view_width = view_size + image_height, image_width = image.shape[0:2] + + scale = np.random.choice(random_scales) + height = int(view_height * scale) + width = int(view_width * scale) + + cropped_image = np.zeros((height, width, 3), dtype=image.dtype) + + w_border = _get_border(border, image_width) + h_border = _get_border(border, image_height) + + ctx = np.random.randint(low=w_border, high=image_width - w_border) + cty = np.random.randint(low=h_border, high=image_height - h_border) + + x0, x1 = max(ctx - width // 2, 0), min(ctx + width // 2, image_width) + y0, y1 = max(cty - height // 2, 0), min(cty + height // 2, image_height) + + left_w, right_w = ctx - x0, x1 - ctx + top_h, bottom_h = cty - y0, y1 - cty + + # crop image + cropped_ctx, cropped_cty = width // 2, height // 2 + x_slice = slice(cropped_ctx - left_w, cropped_ctx + right_w) + y_slice = slice(cropped_cty - top_h, cropped_cty + bottom_h) + cropped_image[y_slice, x_slice, :] = image[y0:y1, x0:x1, :] + + # crop detections + cropped_detections = detections.copy() + cropped_detections[:, 0:4:2] -= x0 + cropped_detections[:, 1:4:2] -= y0 + cropped_detections[:, 0:4:2] += cropped_ctx - left_w + cropped_detections[:, 1:4:2] += cropped_cty - top_h + cropped_extreme_pts = extreme_pts.copy() + cropped_extreme_pts[:, :, 0] -= x0 + cropped_extreme_pts[:, :, 1] -= y0 + cropped_extreme_pts[:, :, 0] += cropped_ctx - left_w + cropped_extreme_pts[:, :, 1] += cropped_cty - top_h + + + return cropped_image, cropped_detections, cropped_extreme_pts \ No newline at end of file diff --git a/test.py b/test.py new file mode 100755 index 0000000..402806e --- /dev/null +++ b/test.py @@ -0,0 +1,99 @@ +#!/usr/bin/env python +import os +import json +import torch +import pprint +import argparse +import importlib +import numpy as np + +import matplotlib +matplotlib.use("Agg") + +from config import system_configs +from nnet.py_factory import NetworkFactory +from db.datasets import datasets + +torch.backends.cudnn.benchmark = False + +def parse_args(): + parser = argparse.ArgumentParser(description="Test CornerNet") + parser.add_argument("cfg_file", help="config file", type=str) + parser.add_argument("--testiter", dest="testiter", + help="test at iteration i", + default=None, type=int) + parser.add_argument("--split", dest="split", + help="which split to use", + default="validation", type=str) + parser.add_argument("--suffix", dest="suffix", default=None, type=str) + parser.add_argument("--debug", action="store_true") + + args = parser.parse_args() + return args + +def make_dirs(directories): + for directory in directories: + if not os.path.exists(directory): + os.makedirs(directory) + +def test(db, split, testiter, debug=False, suffix=None): + result_dir = system_configs.result_dir + result_dir = os.path.join(result_dir, str(testiter), split) + + if suffix is not None: + result_dir = os.path.join(result_dir, suffix) + + make_dirs([result_dir]) + + test_iter = system_configs.max_iter if testiter is None else testiter + print("loading parameters at iteration: {}".format(test_iter)) + + print("building neural network...") + nnet = NetworkFactory(db) + print("loading parameters...") + nnet.load_params(test_iter) + + test_file = "test.{}".format(db.data) + testing = importlib.import_module(test_file).testing + + nnet.cuda() + nnet.eval_mode() + testing(db, nnet, result_dir, debug=debug) + +if __name__ == "__main__": + args = parse_args() + + if args.suffix is None: + cfg_file = os.path.join(system_configs.config_dir, args.cfg_file + ".json") + else: + cfg_file = os.path.join(system_configs.config_dir, args.cfg_file + "-{}.json".format(args.suffix)) + print("cfg_file: {}".format(cfg_file)) + + with open(cfg_file, "r") as f: + configs = json.load(f) + + configs["system"]["snapshot_name"] = args.cfg_file + system_configs.update_config(configs["system"]) + + train_split = system_configs.train_split + val_split = system_configs.val_split + test_split = system_configs.test_split + + split = { + "training": train_split, + "validation": val_split, + "testing": test_split + }[args.split] + + print("loading all datasets...") + dataset = system_configs.dataset + print("split: {}".format(split)) + testing_db = datasets[dataset](configs["db"], split) + + print("system config...") + pprint.pprint(system_configs.full) + + print("db config...") + pprint.pprint(testing_db.configs) + + test(testing_db, args.split, args.testiter, args.debug, args.suffix) diff --git a/test/__init__.py b/test/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/test/coco.py b/test/coco.py new file mode 100644 index 0000000..d712fcd --- /dev/null +++ b/test/coco.py @@ -0,0 +1,204 @@ +import os +import cv2 +import json +import numpy as np +import torch +import matplotlib.pyplot as plt + +from tqdm import tqdm +from config import system_configs +from utils import crop_image, normalize_ +from external.nms import soft_nms, soft_nms_merge + +def _rescale_dets(detections, ratios, borders, sizes): + xs, ys = detections[..., 0:4:2], detections[..., 1:4:2] + xs /= ratios[:, 1][:, None, None] + ys /= ratios[:, 0][:, None, None] + xs -= borders[:, 2][:, None, None] + ys -= borders[:, 0][:, None, None] + np.clip(xs, 0, sizes[:, 1][:, None, None], out=xs) + np.clip(ys, 0, sizes[:, 0][:, None, None], out=ys) + +def save_image(data, fn): + sizes = np.shape(data) + height = float(sizes[0]) + width = float(sizes[1]) + + fig = plt.figure() + fig.set_size_inches(width/height, 1, forward=False) + ax = plt.Axes(fig, [0., 0., 1., 1.]) + ax.set_axis_off() + fig.add_axes(ax) + + ax.imshow(data) + plt.savefig(fn, dpi = height) + plt.close() + +def kp_decode(nnet, images, K, ae_threshold=0.5, kernel=3): + detections = nnet.test([images], ae_threshold=ae_threshold, K=K, kernel=kernel) + detections = detections.data.cpu().numpy() + return detections + +def kp_detection(db, nnet, result_dir, debug=False, decode_func=kp_decode): + debug_dir = os.path.join(result_dir, "debug") + if not os.path.exists(debug_dir): + os.makedirs(debug_dir) + + if db.split != "trainval": + db_inds = db.db_inds[:100] if debug else db.db_inds + else: + db_inds = db.db_inds[:100] if debug else db.db_inds[:5000] + num_images = db_inds.size + + K = db.configs["top_k"] + ae_threshold = db.configs["ae_threshold"] + nms_kernel = db.configs["nms_kernel"] + + scales = db.configs["test_scales"] + weight_exp = db.configs["weight_exp"] + merge_bbox = db.configs["merge_bbox"] + categories = db.configs["categories"] + nms_threshold = db.configs["nms_threshold"] + max_per_image = db.configs["max_per_image"] + nms_algorithm = { + "nms": 0, + "linear_soft_nms": 1, + "exp_soft_nms": 2 + }[db.configs["nms_algorithm"]] + + top_bboxes = {} + for ind in tqdm(range(0, num_images), ncols=80, desc="locating kps"): + db_ind = db_inds[ind] + + image_id = db.image_ids(db_ind) + image_file = db.image_file(db_ind) + image = cv2.imread(image_file) + + height, width = image.shape[0:2] + + detections = [] + + for scale in scales: + new_height = int(height * scale) + new_width = int(width * scale) + new_center = np.array([new_height // 2, new_width // 2]) + + inp_height = new_height | 127 + inp_width = new_width | 127 + + images = np.zeros((1, 3, inp_height, inp_width), dtype=np.float32) + ratios = np.zeros((1, 2), dtype=np.float32) + borders = np.zeros((1, 4), dtype=np.float32) + sizes = np.zeros((1, 2), dtype=np.float32) + + out_height, out_width = (inp_height + 1) // 4, (inp_width + 1) // 4 + height_ratio = out_height / inp_height + width_ratio = out_width / inp_width + + resized_image = cv2.resize(image, (new_width, new_height)) + resized_image, border, offset = crop_image(resized_image, new_center, [inp_height, inp_width]) + + resized_image = resized_image / 255. + normalize_(resized_image, db.mean, db.std) + + images[0] = resized_image.transpose((2, 0, 1)) + borders[0] = border + sizes[0] = [int(height * scale), int(width * scale)] + ratios[0] = [height_ratio, width_ratio] + + images = np.concatenate((images, images[:, :, :, ::-1]), axis=0) + images = torch.from_numpy(images) + dets = decode_func(nnet, images, K, ae_threshold=ae_threshold, kernel=nms_kernel) + dets = dets.reshape(2, -1, 8) + dets[1, :, [0, 2]] = out_width - dets[1, :, [2, 0]] + dets = dets.reshape(1, -1, 8) + + _rescale_dets(dets, ratios, borders, sizes) + dets[:, :, 0:4] /= scale + detections.append(dets) + + detections = np.concatenate(detections, axis=1) + + classes = detections[..., -1] + classes = classes[0] + detections = detections[0] + + # reject detections with negative scores + keep_inds = (detections[:, 4] > -1) + detections = detections[keep_inds] + classes = classes[keep_inds] + + top_bboxes[image_id] = {} + for j in range(categories): + keep_inds = (classes == j) + top_bboxes[image_id][j + 1] = detections[keep_inds][:, 0:7].astype(np.float32) + if merge_bbox: + soft_nms_merge(top_bboxes[image_id][j + 1], Nt=nms_threshold, method=nms_algorithm, weight_exp=weight_exp) + else: + soft_nms(top_bboxes[image_id][j + 1], Nt=nms_threshold, method=nms_algorithm) + top_bboxes[image_id][j + 1] = top_bboxes[image_id][j + 1][:, 0:5] + + scores = np.hstack([ + top_bboxes[image_id][j][:, -1] + for j in range(1, categories + 1) + ]) + if len(scores) > max_per_image: + kth = len(scores) - max_per_image + thresh = np.partition(scores, kth)[kth] + for j in range(1, categories + 1): + keep_inds = (top_bboxes[image_id][j][:, -1] >= thresh) + top_bboxes[image_id][j] = top_bboxes[image_id][j][keep_inds] + + if debug: + image_file = db.image_file(db_ind) + image = cv2.imread(image_file) + + bboxes = {} + for j in range(1, categories + 1): + keep_inds = (top_bboxes[image_id][j][:, -1] > 0.5) + cat_name = db.class_name(j) + cat_size = cv2.getTextSize(cat_name, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)[0] + color = np.random.random((3, )) * 0.6 + 0.4 + color = color * 255 + color = color.astype(np.int32).tolist() + for bbox in top_bboxes[image_id][j][keep_inds]: + bbox = bbox[0:4].astype(np.int32) + if bbox[1] - cat_size[1] - 2 < 0: + cv2.rectangle(image, + (bbox[0], bbox[1] + 2), + (bbox[0] + cat_size[0], bbox[1] + cat_size[1] + 2), + color, -1 + ) + cv2.putText(image, cat_name, + (bbox[0], bbox[1] + cat_size[1] + 2), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), thickness=1 + ) + else: + cv2.rectangle(image, + (bbox[0], bbox[1] - cat_size[1] - 2), + (bbox[0] + cat_size[0], bbox[1] - 2), + color, -1 + ) + cv2.putText(image, cat_name, + (bbox[0], bbox[1] - 2), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), thickness=1 + ) + cv2.rectangle(image, + (bbox[0], bbox[1]), + (bbox[2], bbox[3]), + color, 2 + ) + debug_file = os.path.join(debug_dir, "{}.jpg".format(db_ind)) + + result_json = os.path.join(result_dir, "results.json") + detections = db.convert_to_coco(top_bboxes) + with open(result_json, "w") as f: + json.dump(detections, f) + + cls_ids = list(range(1, categories + 1)) + image_ids = [db.image_ids(ind) for ind in db_inds] + db.evaluate(result_json, cls_ids, image_ids) + return 0 + +def testing(db, nnet, result_dir, debug=False): + return globals()[system_configs.sampling_function](db, nnet, result_dir, debug=debug) diff --git a/test/coco_extreme.py b/test/coco_extreme.py new file mode 100644 index 0000000..5b2a113 --- /dev/null +++ b/test/coco_extreme.py @@ -0,0 +1,251 @@ +import os +import cv2 +import json +import numpy as np +import torch +import matplotlib.pyplot as plt + +from tqdm import tqdm +from config import system_configs +from utils import crop_image, normalize_ +from external.nms import soft_nms_with_points as soft_nms + +def _rescale_dets(detections, ratios, borders, sizes): + xs, ys = detections[..., 0:4:2], detections[..., 1:4:2] + xs /= ratios[:, 1][:, None, None] + ys /= ratios[:, 0][:, None, None] + xs -= borders[:, 2][:, None, None] + ys -= borders[:, 0][:, None, None] + np.clip(xs, 0, sizes[:, 1][:, None, None], out=xs) + np.clip(ys, 0, sizes[:, 0][:, None, None], out=ys) + +def _rescale_ex_pts(detections, ratios, borders, sizes): + xs, ys = detections[..., 5:13:2], detections[..., 6:13:2] + xs /= ratios[:, 1][:, None, None] + ys /= ratios[:, 0][:, None, None] + xs -= borders[:, 2][:, None, None] + ys -= borders[:, 0][:, None, None] + np.clip(xs, 0, sizes[:, 1][:, None, None], out=xs) + np.clip(ys, 0, sizes[:, 0][:, None, None], out=ys) + +def save_image(data, fn): + sizes = np.shape(data) + height = float(sizes[0]) + width = float(sizes[1]) + + fig = plt.figure() + fig.set_size_inches(width/height, 1, forward=False) + ax = plt.Axes(fig, [0., 0., 1., 1.]) + ax.set_axis_off() + fig.add_axes(ax) + + ax.imshow(data) + plt.savefig(fn, dpi = height) + plt.close() + +def _box_inside(box2, box1): + inside = (box2[0] >= box1[0] and box2[1] >= box1[1] and \ + box2[2] <= box1[2] and box2[3] <= box1[3]) + return inside + +def kp_decode(nnet, images, K, kernel=3, aggr_weight=0.1, + scores_thresh=0.1, center_thresh=0.1, debug=False): + detections = nnet.test( + [images], kernel=kernel, aggr_weight=aggr_weight, + scores_thresh=scores_thresh, center_thresh=center_thresh, debug=debug) + detections = detections.data.cpu().numpy() + return detections + +def kp_detection(db, nnet, result_dir, debug=False, decode_func=kp_decode): + debug_dir = os.path.join(result_dir, "debug") + if not os.path.exists(debug_dir): + os.makedirs(debug_dir) + + if db.split != "trainval": + db_inds = db.db_inds[:100] if debug else db.db_inds + else: + db_inds = db.db_inds[:100] if debug else db.db_inds[:5000] + num_images = db_inds.size + + K = db.configs["top_k"] + aggr_weight = db.configs["aggr_weight"] + scores_thresh = db.configs["scores_thresh"] + center_thresh = db.configs["center_thresh"] + suppres_ghost = db.configs["suppres_ghost"] + nms_kernel = db.configs["nms_kernel"] + + scales = db.configs["test_scales"] + categories = db.configs["categories"] + nms_threshold = db.configs["nms_threshold"] + max_per_image = db.configs["max_per_image"] + nms_algorithm = { + "nms": 0, + "linear_soft_nms": 1, + "exp_soft_nms": 2 + }[db.configs["nms_algorithm"]] + + top_bboxes = {} + for ind in tqdm(range(0, num_images), ncols=80, desc="locating kps"): + db_ind = db_inds[ind] + + image_id = db.image_ids(db_ind) + image_file = db.image_file(db_ind) + image = cv2.imread(image_file) + + height, width = image.shape[0:2] + + detections = [] + + for scale in scales: + new_height = int(height * scale) + new_width = int(width * scale) + new_center = np.array([new_height // 2, new_width // 2]) + + inp_height = new_height | 127 + inp_width = new_width | 127 + + images = np.zeros((1, 3, inp_height, inp_width), dtype=np.float32) + ratios = np.zeros((1, 2), dtype=np.float32) + borders = np.zeros((1, 4), dtype=np.float32) + sizes = np.zeros((1, 2), dtype=np.float32) + + out_height, out_width = (inp_height + 1) // 4, (inp_width + 1) // 4 + height_ratio = out_height / inp_height + width_ratio = out_width / inp_width + + resized_image = cv2.resize(image, (new_width, new_height)) + resized_image, border, offset = crop_image( + resized_image, new_center, [inp_height, inp_width]) + + resized_image = resized_image / 255. + normalize_(resized_image, db.mean, db.std) + + images[0] = resized_image.transpose((2, 0, 1)) + borders[0] = border + sizes[0] = [int(height * scale), int(width * scale)] + ratios[0] = [height_ratio, width_ratio] + + images = np.concatenate((images, images[:, :, :, ::-1]), axis=0) + images = torch.from_numpy(images) + dets = decode_func( + nnet, images, K, aggr_weight=aggr_weight, + scores_thresh=scores_thresh, center_thresh=center_thresh, + kernel=nms_kernel, debug=debug) + dets = dets.reshape(2, -1, 14) + dets[1, :, [0, 2]] = out_width - dets[1, :, [2, 0]] + dets[1, :, [5, 7, 9, 11]] = out_width - dets[1, :, [5, 7, 9, 11]] + dets[1, :, [7, 8, 11, 12]] = dets[1, :, [11, 12, 7, 8]].copy() + dets = dets.reshape(1, -1, 14) + + _rescale_dets(dets, ratios, borders, sizes) + _rescale_ex_pts(dets, ratios, borders, sizes) + dets[:, :, 0:4] /= scale + dets[:, :, 5:13] /= scale + detections.append(dets) + + detections = np.concatenate(detections, axis=1) + + classes = detections[..., -1] + classes = classes[0] + detections = detections[0] + + # reject detections with negative scores + keep_inds = (detections[:, 4] > 0) + detections = detections[keep_inds] + classes = classes[keep_inds] + + top_bboxes[image_id] = {} + for j in range(categories): + keep_inds = (classes == j) + top_bboxes[image_id][j + 1] = \ + detections[keep_inds].astype(np.float32) + soft_nms(top_bboxes[image_id][j + 1], + Nt=nms_threshold, method=nms_algorithm) + # top_bboxes[image_id][j + 1] = top_bboxes[image_id][j + 1][:, 0:5] + + scores = np.hstack([ + top_bboxes[image_id][j][:, 4] + for j in range(1, categories + 1) + ]) + if len(scores) > max_per_image: + kth = len(scores) - max_per_image + thresh = np.partition(scores, kth)[kth] + for j in range(1, categories + 1): + keep_inds = (top_bboxes[image_id][j][:, 4] >= thresh) + top_bboxes[image_id][j] = top_bboxes[image_id][j][keep_inds] + + if suppres_ghost: + for j in range(1, categories + 1): + n = len(top_bboxes[image_id][j]) + for k in range(n): + inside_score = 0 + if top_bboxes[image_id][j][k, 4] > 0.2: + for t in range(n): + if _box_inside(top_bboxes[image_id][j][t], + top_bboxes[image_id][j][k]): + inside_score += top_bboxes[image_id][j][t, 4] + if inside_score > top_bboxes[image_id][j][k, 4] * 3: + top_bboxes[image_id][j][k, 4] /= 2 + + if debug: + image_file = db.image_file(db_ind) + image = cv2.imread(image_file) + + bboxes = {} + for j in range(1, categories + 1): + keep_inds = (top_bboxes[image_id][j][:, 4] > 0.5) + cat_name = db.class_name(j) + cat_size = cv2.getTextSize( + cat_name + '0', cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)[0] + color = np.random.random((3, )) * 0.6 + 0.4 + color = color * 255 + color = color.astype(np.int32).tolist() + for bbox in top_bboxes[image_id][j][keep_inds]: + sc = bbox[4] + bbox = bbox[0:4].astype(np.int32) + txt = '{}{:.0f}'.format(cat_name, sc * 10) + if bbox[1] - cat_size[1] - 2 < 0: + cv2.rectangle(image, + (bbox[0], bbox[1] + 2), + (bbox[0] + cat_size[0], bbox[1] + cat_size[1] + 2), + color, -1 + ) + cv2.putText(image, txt, + (bbox[0], bbox[1] + cat_size[1] + 2), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), + thickness=1, lineType=cv2.LINE_AA + ) + else: + cv2.rectangle(image, + (bbox[0], bbox[1] - cat_size[1] - 2), + (bbox[0] + cat_size[0], bbox[1] - 2), + color, -1 + ) + cv2.putText(image, txt, + (bbox[0], bbox[1] - 2), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), + thickness=1, lineType=cv2.LINE_AA + ) + cv2.rectangle(image, + (bbox[0], bbox[1]), + (bbox[2], bbox[3]), + color, 2 + ) + debug_file = os.path.join(debug_dir, "{}.jpg".format(db_ind)) + cv2.imwrite(debug_file, image) + cv2.imshow('out', image) + cv2.waitKey() + + result_json = os.path.join(result_dir, "results.json") + detections = db.convert_to_coco(top_bboxes) + with open(result_json, "w") as f: + json.dump(detections, f) + + cls_ids = list(range(1, categories + 1)) + image_ids = [db.image_ids(ind) for ind in db_inds] + db.evaluate(result_json, cls_ids, image_ids) + return 0 + +def testing(db, nnet, result_dir, debug=False): + return globals()[system_configs.sampling_function]( + db, nnet, result_dir, debug=debug) diff --git a/tools/gen_coco_extreme_points.py b/tools/gen_coco_extreme_points.py new file mode 100644 index 0000000..0085782 --- /dev/null +++ b/tools/gen_coco_extreme_points.py @@ -0,0 +1,128 @@ +import pycocotools.coco as cocoapi +import sys +import cv2 +import numpy as np +import pickle +import json +SPLITS = ['val', 'train'] +ANN_PATH = '../data/coco/annotations/instances_{}2017.json' +OUT_PATH = '../data/coco/annotations/instances_extreme_{}2017.json' +IMG_DIR = '../data/coco/{}2017/' +DEBUG = False +from scipy.spatial import ConvexHull + +def _coco_box_to_bbox(box): + bbox = np.array([box[0], box[1], box[0] + box[2], box[1] + box[3]], + dtype=np.int32) + return bbox + +def _get_extreme_points(pts): + l, t = min(pts[:, 0]), min(pts[:, 1]) + r, b = max(pts[:, 0]), max(pts[:, 1]) + # 3 degrees + thresh = 0.02 + w = r - l + 1 + h = b - t + 1 + + pts = np.concatenate([pts[-1:], pts, pts[:1]], axis=0) + t_idx = np.argmin(pts[:, 1]) + t_idxs = [t_idx] + tmp = t_idx + 1 + while tmp < pts.shape[0] and pts[tmp, 1] - pts[t_idx, 1] <= thresh * h: + t_idxs.append(tmp) + tmp += 1 + tmp = t_idx - 1 + while tmp >= 0 and pts[tmp, 1] - pts[t_idx, 1] <= thresh * h: + t_idxs.append(tmp) + tmp -= 1 + tt = [(max(pts[t_idxs, 0]) + min(pts[t_idxs, 0])) // 2, t] + + b_idx = np.argmax(pts[:, 1]) + b_idxs = [b_idx] + tmp = b_idx + 1 + while tmp < pts.shape[0] and pts[b_idx, 1] - pts[tmp, 1] <= thresh * h: + b_idxs.append(tmp) + tmp += 1 + tmp = b_idx - 1 + while tmp >= 0 and pts[b_idx, 1] - pts[tmp, 1] <= thresh * h: + b_idxs.append(tmp) + tmp -= 1 + bb = [(max(pts[b_idxs, 0]) + min(pts[b_idxs, 0])) // 2, b] + + l_idx = np.argmin(pts[:, 0]) + l_idxs = [l_idx] + tmp = l_idx + 1 + while tmp < pts.shape[0] and pts[tmp, 0] - pts[l_idx, 0] <= thresh * w: + l_idxs.append(tmp) + tmp += 1 + tmp = l_idx - 1 + while tmp >= 0 and pts[tmp, 0] - pts[l_idx, 0] <= thresh * w: + l_idxs.append(tmp) + tmp -= 1 + ll = [l, (max(pts[l_idxs, 1]) + min(pts[l_idxs, 1])) // 2] + + r_idx = np.argmax(pts[:, 0]) + r_idxs = [r_idx] + tmp = r_idx + 1 + while tmp < pts.shape[0] and pts[r_idx, 0] - pts[tmp, 0] <= thresh * w: + r_idxs.append(tmp) + tmp += 1 + tmp = r_idx - 1 + while tmp >= 0 and pts[r_idx, 0] - pts[tmp, 0] <= thresh * w: + r_idxs.append(tmp) + tmp -= 1 + rr = [r, (max(pts[r_idxs, 1]) + min(pts[r_idxs, 1])) // 2] + + return np.array([tt, ll, bb, rr]) + +if __name__ == '__main__': + for split in SPLITS: + data = json.load(open(ANN_PATH.format(split), 'r')) + coco = cocoapi.COCO(ANN_PATH.format(split)) + img_ids = coco.getImgIds() + num_images = len(img_ids) + num_classes = 80 + tot_box = 0 + print('num_images', num_images) + anns_all = data['annotations'] + for i, ann in enumerate(anns_all): + tot_box += 1 + bbox = ann['bbox'] + seg = ann['segmentation'] + if type(seg) == list: + if len(seg) == 1: + pts = np.array(seg[0]).reshape(-1, 2) + else: + pts = [] + for v in seg: + pts += v + pts = np.array(pts).reshape(-1, 2) + else: + mask = coco.annToMask(ann) * 255 + tmp = np.where(mask > 0) + pts = np.asarray(tmp).transpose()[:, ::-1].astype(np.int32) + extreme_points = _get_extreme_points(pts).astype(np.int32) + anns_all[i]['extreme_points'] = extreme_points.copy().tolist() + if DEBUG: + img_id = ann['image_id'] + img_info = coco.loadImgs(ids=[img_id])[0] + img_path = IMG_DIR.format(split) + img_info['file_name'] + img = cv2.imread(img_path) + if type(seg) == list: + mask = np.zeros((img.shape[0], img.shape[1], 1), dtype=np.uint8) + cv2.fillPoly(mask, [pts.astype(np.int32).reshape(-1, 1, 2)], (255,0,0)) + else: + mask = mask.reshape(img.shape[0], img.shape[1], 1) + img = (0.4 * img + 0.6 * mask).astype(np.uint8) + bbox = _coco_box_to_bbox(ann['bbox']) + cl = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 0, 255)] + for j in range(extreme_points.shape[0]): + cv2.circle(img, (extreme_points[j, 0], extreme_points[j, 1]), + 5, cl[j], -1) + cv2.imshow('img', img) + cv2.waitKey() + print('tot_box', tot_box) + data['annotations'] = anns_all + json.dump(data, open(OUT_PATH.format(split), 'w')) + + diff --git a/tools/suppress_ghost.py b/tools/suppress_ghost.py new file mode 100644 index 0000000..bf57db3 --- /dev/null +++ b/tools/suppress_ghost.py @@ -0,0 +1,65 @@ +import pycocotools.coco as coco +from pycocotools.cocoeval import COCOeval +import sys +import cv2 +import numpy as np +import pickle +import json +ANN_PATH = '../data/coco/annotations/instances_val2017.json' +DEBUG = True + +def _coco_box_to_bbox(box): + bbox = np.array([box[0], box[1], box[0] + box[2], box[1] + box[3]], + dtype=np.int32) + return bbox + +def _overlap(box1, box2): + area1 = (box1[2] - box1[0] + 1) * (box1[3] - box1[1] + 1) + inter = max(min(box1[2], box2[2]) - max(box1[0], box2[0]) + 1, 0) * \ + max(min(box1[3], box2[3]) - max(box1[1], box2[1]) + 1, 0) + iou = 1.0 * inter / (area1 + 1e-5) + return iou + +def _box_inside(box2, box1): + inside = (box2[0] >= box1[0] and box2[1] >= box1[1] and \ + box2[2] <= box1[2] and box2[3] <= box1[3]) + return inside + +if __name__ == '__main__': + if len(sys.argv) > 2: + ANN_PATH = sys.argv[2] + coco = coco.COCO(ANN_PATH) + pred_path = sys.argv[1] + out_path = pred_path[:-5] + '_no_ghost.json' + dets = coco.loadRes(pred_path) + img_ids = coco.getImgIds() + num_images = len(img_ids) + thresh = 4 + out = [] + for i, img_id in enumerate(img_ids): + if i % 500 == 0: + print(i) + pred_ids = dets.getAnnIds(imgIds=[img_id]) + preds = dets.loadAnns(pred_ids) + num_preds = len(preds) + for j in range(num_preds): + overlap_score = 0 + if preds[j]['score'] > 0.2: + for k in range(num_preds): + if preds[j]['category_id'] == preds[k]['category_id'] and \ + _box_inside(_coco_box_to_bbox(preds[k]['bbox']), + _coco_box_to_bbox(preds[j]['bbox'])) > 0.8: + overlap_score += preds[k]['score'] + if overlap_score > thresh * preds[j]['score']: + # print('overlap_score', overlap_score, preds[j]['score']) + preds[j]['score'] = preds[j]['score'] / 2 + # preds[j]['score'] = preds[j]['score'] * np.exp(-(overlap_score / preds[j]['score'] - thresh)**2/2) + out.append(preds[j]) + json.dump(out, open(out_path, 'w')) + dets_refined = coco.loadRes(out_path) + coco_eval = COCOeval(coco, dets_refined, "bbox") + coco_eval.evaluate() + coco_eval.accumulate() + coco_eval.summarize() + + diff --git a/train.py b/train.py new file mode 100755 index 0000000..ed371c8 --- /dev/null +++ b/train.py @@ -0,0 +1,225 @@ +#!/usr/bin/env python +import os + +import json +import torch +import numpy as np +import queue +import pprint +import random +import argparse +import importlib +import threading +import traceback + +from tqdm import tqdm +from utils import stdout_to_tqdm +from config import system_configs +from nnet.py_factory import NetworkFactory +from torch.multiprocessing import Process, Queue, Pool +from db.datasets import datasets + +torch.backends.cudnn.enabled = True +torch.backends.cudnn.benchmark = True + +def parse_args(): + parser = argparse.ArgumentParser(description="Train CornerNet") + parser.add_argument("cfg_file", help="config file", type=str) + parser.add_argument("--iter", dest="start_iter", + help="train at iteration i", + default=0, type=int) + parser.add_argument("--threads", dest="threads", default=4, type=int) + parser.add_argument("--debug", action="store_true") + + args = parser.parse_args() + return args + +class AverageMeter(object): + """Computes and stores the average and current value""" + def __init__(self): + self.reset() + + def reset(self): + self.val = 0 + self.avg = 0 + self.sum = 0 + self.count = 0 + + def update(self, val, n=1): + self.val = val + self.sum += val * n + self.count += n + if self.count > 0: + self.avg = self.sum / self.count + +def prefetch_data(db, queue, sample_data, data_aug, debug=False): + ind = 0 + print("start prefetching data...") + np.random.seed(os.getpid()) + while True: + try: + data, ind = sample_data(db, ind, data_aug=data_aug, debug=debug) + queue.put(data) + except Exception as e: + traceback.print_exc() + raise e + +def pin_memory(data_queue, pinned_data_queue, sema): + while True: + data = data_queue.get() + + data["xs"] = [x.pin_memory() for x in data["xs"]] + data["ys"] = [y.pin_memory() for y in data["ys"]] + + pinned_data_queue.put(data) + + if sema.acquire(blocking=False): + return + +def init_parallel_jobs(dbs, queue, fn, data_aug, debug=False): + tasks = [Process(target=prefetch_data, + args=(db, queue, fn, data_aug, debug)) for db in dbs] + for task in tasks: + task.daemon = True + task.start() + return tasks + +def train(training_dbs, validation_db, start_iter=0, debug=False): + learning_rate = system_configs.learning_rate + max_iteration = system_configs.max_iter + pretrained_model = system_configs.pretrain + snapshot = system_configs.snapshot + # val_iter = system_configs.val_iter + display = system_configs.display + decay_rate = system_configs.decay_rate + stepsize = system_configs.stepsize + + # getting the size of each database + training_size = len(training_dbs[0].db_inds) + # validation_size = len(validation_db.db_inds) + + # queues storing data for training + training_queue = Queue(system_configs.prefetch_size) + # validation_queue = Queue(5) + + # queues storing pinned data for training + pinned_training_queue = queue.Queue(system_configs.prefetch_size) + # pinned_validation_queue = queue.Queue(5) + + # load data sampling function + data_file = "sample.{}".format(training_dbs[0].data) + sample_data = importlib.import_module(data_file).sample_data + + # allocating resources for parallel reading + training_tasks = init_parallel_jobs( + training_dbs, training_queue, sample_data, True, debug) + # if val_iter: + # validation_tasks = init_parallel_jobs([validation_db], validation_queue, sample_data, False) + + training_pin_semaphore = threading.Semaphore() + # validation_pin_semaphore = threading.Semaphore() + training_pin_semaphore.acquire() + # validation_pin_semaphore.acquire() + + training_pin_args = (training_queue, pinned_training_queue, training_pin_semaphore) + training_pin_thread = threading.Thread(target=pin_memory, args=training_pin_args) + training_pin_thread.daemon = True + training_pin_thread.start() + + # validation_pin_args = (validation_queue, pinned_validation_queue, validation_pin_semaphore) + # validation_pin_thread = threading.Thread(target=pin_memory, args=validation_pin_args) + # validation_pin_thread.daemon = True + # validation_pin_thread.start() + + print("building model...") + nnet = NetworkFactory(training_dbs[0]) + + if pretrained_model is not None: + if not os.path.exists(pretrained_model): + raise ValueError("pretrained model does not exist") + print("loading from pretrained model") + nnet.load_pretrained_params(pretrained_model) + + if start_iter: + learning_rate /= (decay_rate ** (start_iter // stepsize)) + + nnet.load_params(start_iter) + nnet.set_lr(learning_rate) + print("training starts from iteration {} with learning_rate {}".format(start_iter + 1, learning_rate)) + else: + nnet.set_lr(learning_rate) + + print("training start...") + nnet.cuda() + nnet.train_mode() + avg_loss = AverageMeter() + with stdout_to_tqdm() as save_stdout: + for iteration in tqdm(range(start_iter + 1, max_iteration + 1), file=save_stdout, ncols=80): + training = pinned_training_queue.get(block=True) + training_loss = nnet.train(**training) + avg_loss.update(training_loss.item()) + + if display and iteration % display == 0: + print("training loss at iteration {}: {:.6f} ({:.6f})".format( + iteration, training_loss.item(), avg_loss.avg)) + del training_loss + + # if val_iter and validation_db.db_inds.size and iteration % val_iter == 0: + # nnet.eval_mode() + # validation = pinned_validation_queue.get(block=True) + # validation_loss = nnet.validate(**validation) + # print("validation loss at iteration {}: {}".format(iteration, validation_loss.item())) + # nnet.train_mode() + + if iteration % snapshot == 0: + nnet.save_params(iteration) + + if iteration % 1000 == 0: + nnet.save_params(-1) + avg_loss = AverageMeter() + + if iteration % stepsize == 0: + learning_rate /= decay_rate + nnet.set_lr(learning_rate) + + # sending signal to kill the thread + training_pin_semaphore.release() + # validation_pin_semaphore.release() + + # terminating data fetching processes + for training_task in training_tasks: + training_task.terminate() + # for validation_task in validation_tasks: + # validation_task.terminate() + +if __name__ == "__main__": + args = parse_args() + + cfg_file = os.path.join(system_configs.config_dir, args.cfg_file + ".json") + with open(cfg_file, "r") as f: + configs = json.load(f) + + configs["system"]["snapshot_name"] = args.cfg_file + system_configs.update_config(configs["system"]) + + train_split = system_configs.train_split + val_split = system_configs.val_split + + print("loading all datasets...") + dataset = system_configs.dataset + # threads = max(torch.cuda.device_count() * 2, 4) + threads = args.threads + print("using {} threads".format(threads)) + training_dbs = [datasets[dataset](configs["db"], train_split) for _ in range(threads)] + # Remove validation to save GPU resources + # validation_db = datasets[dataset](configs["db"], val_split) + + print("system config...") + pprint.pprint(system_configs.full) + + print("db config...") + pprint.pprint(training_dbs[0].configs) + + print("len of db: {}".format(len(training_dbs[0].db_inds))) + # train(training_dbs, validation_db, args.start_iter) + train(training_dbs, None, args.start_iter, args.debug) diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 0000000..7badeae --- /dev/null +++ b/utils/__init__.py @@ -0,0 +1,4 @@ +from .tqdm import stdout_to_tqdm + +from .image import crop_image +from .image import color_jittering_, lighting_, normalize_ diff --git a/utils/color_map.py b/utils/color_map.py new file mode 100644 index 0000000..bc6869f --- /dev/null +++ b/utils/color_map.py @@ -0,0 +1,113 @@ +# Copyright (c) 2017-present, Facebook, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +############################################################################## + +"""An awesome colormap for really neat visualizations.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function +from __future__ import unicode_literals + +import numpy as np + + +def colormap(rgb=False): + color_list = np.array( + [ + 0.000, 0.447, 0.741, + 0.850, 0.325, 0.098, + 0.929, 0.694, 0.125, + 0.494, 0.184, 0.556, + 0.466, 0.674, 0.188, + 0.301, 0.745, 0.933, + 0.635, 0.078, 0.184, + 0.300, 0.300, 0.300, + 0.600, 0.600, 0.600, + 1.000, 0.000, 0.000, + 1.000, 0.500, 0.000, + 0.749, 0.749, 0.000, + 0.000, 1.000, 0.000, + 0.000, 0.000, 1.000, + 0.667, 0.000, 1.000, + 0.333, 0.333, 0.000, + 0.333, 0.667, 0.000, + 0.333, 1.000, 0.000, + 0.667, 0.333, 0.000, + 0.667, 0.667, 0.000, + 0.667, 1.000, 0.000, + 1.000, 0.333, 0.000, + 1.000, 0.667, 0.000, + 1.000, 1.000, 0.000, + 0.000, 0.333, 0.500, + 0.000, 0.667, 0.500, + 0.000, 1.000, 0.500, + 0.333, 0.000, 0.500, + 0.333, 0.333, 0.500, + 0.333, 0.667, 0.500, + 0.333, 1.000, 0.500, + 0.667, 0.000, 0.500, + 0.667, 0.333, 0.500, + 0.667, 0.667, 0.500, + 0.667, 1.000, 0.500, + 1.000, 0.000, 0.500, + 1.000, 0.333, 0.500, + 1.000, 0.667, 0.500, + 1.000, 1.000, 0.500, + 0.000, 0.333, 1.000, + 0.000, 0.667, 1.000, + 0.000, 1.000, 1.000, + 0.333, 0.000, 1.000, + 0.333, 0.333, 1.000, + 0.333, 0.667, 1.000, + 0.333, 1.000, 1.000, + 0.667, 0.000, 1.000, + 0.667, 0.333, 1.000, + 0.667, 0.667, 1.000, + 0.667, 1.000, 1.000, + 1.000, 0.000, 1.000, + 1.000, 0.333, 1.000, + 1.000, 0.667, 1.000, + 0.167, 0.000, 0.000, + 0.333, 0.000, 0.000, + 0.500, 0.000, 0.000, + 0.667, 0.000, 0.000, + 0.833, 0.000, 0.000, + 1.000, 0.000, 0.000, + 0.000, 0.167, 0.000, + 0.000, 0.333, 0.000, + 0.000, 0.500, 0.000, + 0.000, 0.667, 0.000, + 0.000, 0.833, 0.000, + 0.000, 1.000, 0.000, + 0.000, 0.000, 0.167, + 0.000, 0.000, 0.333, + 0.000, 0.000, 0.500, + 0.000, 0.000, 0.667, + 0.000, 0.000, 0.833, + 0.000, 0.000, 1.000, + 0.000, 0.000, 0.000, + 0.143, 0.143, 0.143, + 0.286, 0.286, 0.286, + 0.429, 0.429, 0.429, + 0.571, 0.571, 0.571, + 0.714, 0.714, 0.714, + 0.857, 0.857, 0.857, + 1.000, 1.000, 1.000 + ] + ).astype(np.float32) + color_list = color_list.reshape((-1, 3)) * 255 + if not rgb: + color_list = color_list[:, ::-1] + return color_list diff --git a/utils/debugger.py b/utils/debugger.py new file mode 100644 index 0000000..f2bd2a6 --- /dev/null +++ b/utils/debugger.py @@ -0,0 +1,205 @@ +import numpy as np +import cv2 +import matplotlib.pyplot as plt + +color_list = np.array( + [ + 0.000, 0.447, 0.741, + 0.850, 0.325, 0.098, + 0.929, 0.694, 0.125, + 0.494, 0.184, 0.556, + 0.466, 0.674, 0.188, + 0.301, 0.745, 0.933, + 0.635, 0.078, 0.184, + 0.300, 0.300, 0.300, + 0.600, 0.600, 0.600, + 1.000, 0.000, 0.000, + 1.000, 0.500, 0.000, + 0.749, 0.749, 0.000, + 0.000, 1.000, 0.000, + 0.000, 0.000, 1.000, + 0.667, 0.000, 1.000, + 0.333, 0.333, 0.000, + 0.333, 0.667, 0.000, + 0.333, 1.000, 0.000, + 0.667, 0.333, 0.000, + 0.667, 0.667, 0.000, + 0.667, 1.000, 0.000, + 1.000, 0.333, 0.000, + 1.000, 0.667, 0.000, + 1.000, 1.000, 0.000, + 0.000, 0.333, 0.500, + 0.000, 0.667, 0.500, + 0.000, 1.000, 0.500, + 0.333, 0.000, 0.500, + 0.333, 0.333, 0.500, + 0.333, 0.667, 0.500, + 0.333, 1.000, 0.500, + 0.667, 0.000, 0.500, + 0.667, 0.333, 0.500, + 0.667, 0.667, 0.500, + 0.667, 1.000, 0.500, + 1.000, 0.000, 0.500, + 1.000, 0.333, 0.500, + 1.000, 0.667, 0.500, + 1.000, 1.000, 0.500, + 0.000, 0.333, 1.000, + 0.000, 0.667, 1.000, + 0.000, 1.000, 1.000, + 0.333, 0.000, 1.000, + 0.333, 0.333, 1.000, + 0.333, 0.667, 1.000, + 0.333, 1.000, 1.000, + 0.667, 0.000, 1.000, + 0.667, 0.333, 1.000, + 0.667, 0.667, 1.000, + 0.667, 1.000, 1.000, + 1.000, 0.000, 1.000, + 1.000, 0.333, 1.000, + 1.000, 0.667, 1.000, + 0.167, 0.000, 0.000, + 0.333, 0.000, 0.000, + 0.500, 0.000, 0.000, + 0.667, 0.000, 0.000, + 0.833, 0.000, 0.000, + 1.000, 0.000, 0.000, + 0.000, 0.167, 0.000, + 0.000, 0.333, 0.000, + 0.000, 0.500, 0.000, + 0.000, 0.667, 0.000, + 0.000, 0.833, 0.000, + 0.000, 1.000, 0.000, + 0.000, 0.000, 0.167, + 0.000, 0.000, 0.333, + 0.000, 0.000, 0.500, + 0.000, 0.000, 0.667, + 0.000, 0.000, 0.833, + 0.000, 0.000, 1.000, + 0.000, 0.000, 0.000, + 0.143, 0.143, 0.143, + 0.286, 0.286, 0.286, + 0.429, 0.429, 0.429, + 0.571, 0.571, 0.571, + 0.714, 0.714, 0.714, + 0.857, 0.857, 0.857, + 1.000, 1.000, 1.000, + 0.50, 0.5, 0 + ] + ).astype(np.float32) +color_list = color_list.reshape((-1, 3)) * 255 + +def show_2d(img, points, c, edges): + num_joints = points.shape[0] + points = ((points.reshape(num_joints, -1))).astype(np.int32) + for j in range(num_joints): + cv2.circle(img, (points[j, 0], points[j, 1]), 3, c, -1) + for e in edges: + if points[e].min() > 0: + cv2.line(img, (points[e[0], 0], points[e[0], 1]), + (points[e[1], 0], points[e[1], 1]), c, 2) + return img + +class Debugger(object): + def __init__(self, ipynb = False, num_classes=80): + self.ipynb = ipynb + if not self.ipynb: + self.plt = plt + self.fig = self.plt.figure() + self.imgs = {} + # colors = [((np.random.random((3, )) * 0.6 + 0.4)*255).astype(np.uint8) \ + # for _ in range(num_classes)] + colors = [(color_list[_]).astype(np.uint8) \ + for _ in range(num_classes)] + self.colors = np.array(colors, dtype=np.uint8).reshape(len(colors), 1, 1, 3) + + def add_img(self, img, imgId = 'default', revert_color=False): + if revert_color: + img = 255 - img + self.imgs[imgId] = img.copy() + + def add_mask(self, mask, bg, imgId = 'default', trans = 0.8): + self.imgs[imgId] = (mask.reshape(mask.shape[0], mask.shape[1], 1) * 255 * trans + \ + bg * (1 - trans)).astype(np.uint8) + + def add_point_2d(self, point, c, edges, imgId = 'default'): + self.imgs[imgId] = show_2d(self.imgs[imgId], point, c, edges) + + def show_img(self, pause = False, imgId = 'default'): + cv2.imshow('{}'.format(imgId), self.imgs[imgId]) + if pause: + cv2.waitKey() + + def add_blend_img(self, back, fore, imgId='blend', trans=0.5): + # fore = 255 - fore + if fore.shape[0] != back.shape[0] or fore.shape[0] != back.shape[1]: + fore = cv2.resize(fore, (back.shape[1], back.shape[0])) + if len(fore.shape) == 2: + fore = fore.reshape(fore.shape[0], fore.shape[1], 1) + self.imgs[imgId] = (back * (1. - trans) + fore * trans) + self.imgs[imgId][self.imgs[imgId] > 255] = 255 + self.imgs[imgId] = self.imgs[imgId].astype(np.uint8) + + def gen_colormap(self, img, s=4): + num_classes = len(self.colors) + img[img < 0] = 0 + h, w = img.shape[1], img.shape[2] + color_map = np.zeros((h*s, w*s, 3), dtype=np.uint8) + for i in range(num_classes): + resized = cv2.resize(img[i], (w*s, h*s)).reshape(h*s, w*s, 1) + cl = self.colors[i] + color_map = np.maximum(color_map, (resized * cl).astype(np.uint8)) + return color_map + + def add_rect(self, rect1, rect2, c, conf=1, imgId = 'default'): + cv2.rectangle(self.imgs[imgId], (rect1[0], rect1[1]), (rect2[0], rect2[1]), c, 2) + if conf < 1: + cv2.circle(self.imgs[imgId], (rect1[0], rect1[1]), int(10 * conf), c, 1) + cv2.circle(self.imgs[imgId], (rect2[0], rect2[1]), int(10 * conf), c, 1) + cv2.circle(self.imgs[imgId], (rect1[0], rect2[1]), int(10 * conf), c, 1) + cv2.circle(self.imgs[imgId], (rect2[0], rect1[1]), int(10 * conf), c, 1) + + def add_points(self, points, img_id = 'default'): + num_classes = len(points) + assert num_classes == len(self.colors) + for i in range(num_classes): + for j in range(len(points[i])): + c = self.colors[i, 0, 0] + cv2.circle(self.imgs[img_id], (points[i][j][0] * 4, points[i][j][1] * 4), + 5, (255, 255, 255), -1) + cv2.circle(self.imgs[img_id], (points[i][j][0] * 4, points[i][j][1] * 4), + 3, (int(c[0]), int(c[1]), int(c[2])), -1) + + def show_all_imgs(self, pause=False): + if not self.ipynb: + for i, v in self.imgs.items(): + cv2.imshow('{}'.format(i), v) + if pause: + cv2.waitKey() + else: + self.ax = None + nImgs = len(self.imgs) + fig=plt.figure(figsize=(nImgs * 10,10)) + nCols = nImgs + nRows = nImgs // nCols + for i, (k, v) in enumerate(self.imgs.items()): + fig.add_subplot(1, nImgs, i + 1) + if len(v.shape) == 3: + plt.imshow(cv2.cvtColor(v, cv2.COLOR_BGR2RGB)) + else: + plt.imshow(v) + plt.show() + + def save_img(self, imgId='default', path='./cache/debug/'): + cv2.imwrite(path + '{}.png'.format(imgId), self.imgs[imgId]) + + def save_all_imgs(self, path='./cache/debug/', prefix='', genID=False): + if genID: + try: + idx = int(np.loadtxt(path + '/id.txt')) + except: + idx = 0 + prefix=idx + np.savetxt(path + '/id.txt', np.ones(1) * (idx + 1), fmt='%d') + for i, v in self.imgs.items(): + cv2.imwrite(path + '/{}{}.png'.format(prefix, i), v) + diff --git a/utils/image.py b/utils/image.py new file mode 100644 index 0000000..ff2a43e --- /dev/null +++ b/utils/image.py @@ -0,0 +1,71 @@ +import cv2 +import numpy as np +import random + +def grayscale(image): + return cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) + +def normalize_(image, mean, std): + image -= mean + image /= std + +def lighting_(data_rng, image, alphastd, eigval, eigvec): + alpha = data_rng.normal(scale=alphastd, size=(3, )) + image += np.dot(eigvec, eigval * alpha) + +def blend_(alpha, image1, image2): + image1 *= alpha + image2 *= (1 - alpha) + image1 += image2 + +def saturation_(data_rng, image, gs, gs_mean, var): + alpha = 1. + data_rng.uniform(low=-var, high=var) + blend_(alpha, image, gs[:, :, None]) + +def brightness_(data_rng, image, gs, gs_mean, var): + alpha = 1. + data_rng.uniform(low=-var, high=var) + image *= alpha + +def contrast_(data_rng, image, gs, gs_mean, var): + alpha = 1. + data_rng.uniform(low=-var, high=var) + blend_(alpha, image, gs_mean) + +def color_jittering_(data_rng, image): + functions = [brightness_, contrast_, saturation_] + random.shuffle(functions) + + gs = grayscale(image) + gs_mean = gs.mean() + for f in functions: + f(data_rng, image, gs, gs_mean, 0.4) + +def crop_image(image, center, size): + cty, ctx = center + height, width = size + im_height, im_width = image.shape[0:2] + cropped_image = np.zeros((height, width, 3), dtype=image.dtype) + + x0, x1 = max(0, ctx - width // 2), min(ctx + width // 2, im_width) + y0, y1 = max(0, cty - height // 2), min(cty + height // 2, im_height) + + left, right = ctx - x0, x1 - ctx + top, bottom = cty - y0, y1 - cty + + cropped_cty, cropped_ctx = height // 2, width // 2 + y_slice = slice(cropped_cty - top, cropped_cty + bottom) + x_slice = slice(cropped_ctx - left, cropped_ctx + right) + cropped_image[y_slice, x_slice, :] = image[y0:y1, x0:x1, :] + + border = np.array([ + cropped_cty - top, + cropped_cty + bottom, + cropped_ctx - left, + cropped_ctx + right + ], dtype=np.float32) + + offset = np.array([ + cty - height // 2, + ctx - width // 2 + ]) + + return cropped_image, border, offset diff --git a/utils/tqdm.py b/utils/tqdm.py new file mode 100755 index 0000000..334dfc1 --- /dev/null +++ b/utils/tqdm.py @@ -0,0 +1,25 @@ +import sys +import numpy as np +import contextlib + +from tqdm import tqdm + +class TqdmFile(object): + dummy_file = None + def __init__(self, dummy_file): + self.dummy_file = dummy_file + + def write(self, x): + if len(x.rstrip()) > 0: + tqdm.write(x, file=self.dummy_file) + +@contextlib.contextmanager +def stdout_to_tqdm(): + save_stdout = sys.stdout + try: + sys.stdout = TqdmFile(sys.stdout) + yield save_stdout + except Exception as exc: + raise exc + finally: + sys.stdout = save_stdout diff --git a/utils/visualize.py b/utils/visualize.py new file mode 100644 index 0000000..dfc8f43 --- /dev/null +++ b/utils/visualize.py @@ -0,0 +1,127 @@ +import cv2 +import numpy as np +import matplotlib.pyplot as plt +from matplotlib.patches import Polygon +import pycocotools.mask as mask_util + +_GRAY = (218, 227, 218) +_GREEN = (18, 127, 15) +_WHITE = (255, 255, 255) + +def vis_mask(img, mask, col, alpha=0.4, show_border=True, border_thick=2): + """Visualizes a single binary mask.""" + + img = img.astype(np.float32) + idx = np.nonzero(mask) + + img[idx[0], idx[1], :] *= 1.0 - alpha + img[idx[0], idx[1], :] += alpha * col + + if show_border: + _, contours, _ = cv2.findContours( + mask.copy(), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE) + cv2.drawContours(img, contours, -1, _WHITE, border_thick, cv2.LINE_AA) + + return img.astype(np.uint8) + + +def vis_octagon(img, extreme_points, col, border_thick=2): + """Visualizes a single binary mask.""" + + img = img.astype(np.uint8) + # COL = (col).astype(np.uint8).tolist() + # print('col', COL) + # octagon = get_octagon(extreme_points) + # octagon = np.array(octagon).reshape(8, 1, 2).astype(np.int32) + # cv2.polylines(img, [octagon], + # True, COL, border_thick) + mask = extreme_point_to_octagon_mask( + extreme_points, img.shape[0], img.shape[1]) + + img = vis_mask(img, mask, col) + + return img.astype(np.uint8) + +def vis_ex(img, extreme_points, col, border_thick=2): + """Visualizes a single binary mask.""" + + img = img.astype(np.uint8) + COL = (col).astype(np.uint8).tolist() + # print('col', COL) + ex = np.array(extreme_points).reshape(4, 2).astype(np.int32) + + L = 10 + T = 0.7 + cv2.arrowedLine(img, (ex[0][0], ex[0][1] + L), (ex[0][0], ex[0][1]), COL, border_thick, tipLength=T) + cv2.arrowedLine(img, (ex[1][0] + L, ex[1][1]), (ex[1][0], ex[1][1]), COL, border_thick, tipLength=T) + cv2.arrowedLine(img, (ex[2][0], ex[2][1] - L), (ex[2][0], ex[2][1]), COL, border_thick, tipLength=T) + cv2.arrowedLine(img, (ex[3][0] - L, ex[3][1]), (ex[3][0], ex[3][1]), COL, border_thick, tipLength=T) + + ''' + R = 6 + cv2.circle(img, (ex[0][0], ex[0][1]), R, COL, -1) + cv2.circle(img, (ex[1][0], ex[1][1]), R, COL, -1) + cv2.circle(img, (ex[2][0], ex[2][1]), R, COL, -1) + cv2.circle(img, (ex[3][0], ex[3][1]), R, COL, -1) + + cv2.circle(img, (ex[0][0], ex[0][1]), R, _WHITE, 2) + cv2.circle(img, (ex[1][0], ex[1][1]), R, _WHITE, 2) + cv2.circle(img, (ex[2][0], ex[2][1]), R, _WHITE, 2) + cv2.circle(img, (ex[3][0], ex[3][1]), R, _WHITE, 2) + ''' + return img.astype(np.uint8) + + +def vis_class(img, pos, class_str, font_scale=0.35): + """Visualizes the class.""" + img = img.astype(np.uint8) + x0, y0 = int(pos[0]), int(pos[1]) + # Compute text size. + txt = class_str + font = cv2.FONT_HERSHEY_SIMPLEX + ((txt_w, txt_h), _) = cv2.getTextSize(txt, font, font_scale, 1) + # Place text background. + if y0 - int(1.3 * txt_h) < 0: + y0 = y0 + int(1.6 * txt_h) + back_tl = x0, y0 - int(1.3 * txt_h) + back_br = x0 + txt_w, y0 + cv2.rectangle(img, back_tl, back_br, _GREEN, -1) + # cv2.rectangle(img, back_tl, back_br, _GRAY, -1) + # Show text. + txt_tl = x0, y0 - int(0.3 * txt_h) + cv2.putText(img, txt, txt_tl, font, font_scale, _GRAY, lineType=cv2.LINE_AA) + # cv2.putText(img, txt, txt_tl, font, font_scale, (46, 52, 54), lineType=cv2.LINE_AA) + return img + + +def vis_bbox(img, bbox, thick=2): + """Visualizes a bounding box.""" + img = img.astype(np.uint8) + (x0, y0, w, h) = bbox + x1, y1 = int(x0 + w), int(y0 + h) + x0, y0 = int(x0), int(y0) + cv2.rectangle(img, (x0, y0), (x1, y1), _GREEN, thickness=thick) + return img + +def get_octagon(ex): + ex = np.array(ex).reshape(4, 2) + w, h = ex[3][0] - ex[1][0], ex[2][1] - ex[0][1] + t, l, b, r = ex[0][1], ex[1][0], ex[2][1], ex[3][0] + x = 8. + octagon = [[min(ex[0][0] + w / x, r), ex[0][1], \ + max(ex[0][0] - w / x, l), ex[0][1], \ + ex[1][0], max(ex[1][1] - h / x, t), \ + ex[1][0], min(ex[1][1] + h / x, b), \ + max(ex[2][0] - w / x, l), ex[2][1], \ + min(ex[2][0] + w / x, r), ex[2][1], \ + ex[3][0], min(ex[3][1] + h / x, b), \ + ex[3][0], max(ex[3][1] - h / x, t) + ]] + return octagon + +def extreme_point_to_octagon_mask(extreme_points, h, w): + octagon = get_octagon(extreme_points) + rles = mask_util.frPyObjects(octagon, h, w) + rle = mask_util.merge(rles) + mask = mask_util.decode(rle) + return mask \ No newline at end of file