├── .github └── workflows │ ├── pypi_release.yml │ └── pythonpublish.yml ├── .gitignore ├── LICENSE ├── README.md ├── gluoncvth ├── __init__.py ├── models │ ├── __init__.py │ ├── base.py │ ├── deeplab.py │ ├── fcn.py │ ├── model_store.py │ ├── model_zoo.py │ ├── pspnet.py │ ├── resnet.py │ └── wideresnet.py └── utils │ ├── __init__.py │ ├── files.py │ ├── pallete.py │ └── preset.py ├── image └── demo_deeplab_ade.png └── setup.py /.github/workflows/pypi_release.yml: -------------------------------------------------------------------------------- 1 | # This workflows will upload a Python Package using Twine when a release is created 2 | # For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries 3 | 4 | name: Pypi Release 5 | 6 | on: 7 | release: 8 | types: [created] 9 | 10 | jobs: 11 | deploy: 12 | 13 | runs-on: ubuntu-18.04 14 | 15 | steps: 16 | - uses: actions/checkout@master 17 | - name: Set up Python 18 | uses: actions/setup-python@v1 19 | with: 20 | python-version: '3.7' 21 | - name: Install dependencies 22 | run: | 23 | python -m pip install --upgrade pip 24 | pip install setuptools wheel twine pypandoc 25 | - name: Build and publish 26 | env: 27 | TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }} 28 | TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }} 29 | RELEASE: 1 30 | run: | 31 | python setup.py sdist bdist_wheel 32 | twine upload dist/* --verbose 33 | -------------------------------------------------------------------------------- /.github/workflows/pythonpublish.yml: -------------------------------------------------------------------------------- 1 | name: Upload Python Package 2 | 3 | on: 4 | push: 5 | branches: [ master ] 6 | 7 | jobs: 8 | deploy: 9 | runs-on: ubuntu-latest 10 | steps: 11 | - uses: actions/checkout@v1 12 | - name: Set up Python 13 | uses: actions/setup-python@v1 14 | with: 15 | python-version: '3.x' 16 | - name: Install dependencies 17 | run: | 18 | python -m pip install --upgrade pip 19 | pip install setuptools wheel twine 20 | - name: Build and publish 21 | env: 22 | TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }} 23 | TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }} 24 | run: | 25 | python setup.py sdist bdist_wheel 26 | twine upload dist/* 27 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | *.egg-info/ 24 | .installed.cfg 25 | *.egg 26 | MANIFEST 27 | 28 | # PyInstaller 29 | # Usually these files are written by a python script from a template 30 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 31 | *.manifest 32 | *.spec 33 | 34 | # Installer logs 35 | pip-log.txt 36 | pip-delete-this-directory.txt 37 | 38 | # Unit test / coverage reports 39 | htmlcov/ 40 | .tox/ 41 | .coverage 42 | .coverage.* 43 | .cache 44 | nosetests.xml 45 | coverage.xml 46 | *.cover 47 | .hypothesis/ 48 | .pytest_cache/ 49 | 50 | # Translations 51 | *.mo 52 | *.pot 53 | 54 | # Django stuff: 55 | *.log 56 | local_settings.py 57 | db.sqlite3 58 | 59 | # Flask stuff: 60 | instance/ 61 | .webassets-cache 62 | 63 | # Scrapy stuff: 64 | .scrapy 65 | 66 | # Sphinx documentation 67 | docs/_build/ 68 | 69 | # PyBuilder 70 | target/ 71 | 72 | # Jupyter Notebook 73 | .ipynb_checkpoints 74 | 75 | # pyenv 76 | .python-version 77 | 78 | # celery beat schedule file 79 | celerybeat-schedule 80 | 81 | # SageMath parsed files 82 | *.sage.py 83 | 84 | # Environments 85 | .env 86 | .venv 87 | env/ 88 | venv/ 89 | ENV/ 90 | env.bak/ 91 | venv.bak/ 92 | 93 | # Spyder project settings 94 | .spyderproject 95 | .spyproject 96 | 97 | # Rope project settings 98 | .ropeproject 99 | 100 | # mkdocs documentation 101 | /site 102 | 103 | # mypy 104 | .mypy_cache/ 105 | 106 | # garbadge 107 | *.DS_Store 108 | *.swp 109 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 Hang Zhang 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | [![PyPI](https://img.shields.io/pypi/v/gluoncv-torch.svg)](https://pypi.python.org/pypi/gluoncv-torch) 2 | [![PyPI Pre-release](https://img.shields.io/badge/pypi--prerelease-v0.0.6-ff69b4.svg)](https://pypi.org/project/gluoncv-torch/#history) 3 | [![Upload Python Package](https://github.com/zhanghang1989/gluoncv-torch/workflows/Upload%20Python%20Package/badge.svg)](https://github.com/zhanghang1989/gluoncv-torch/actions) 4 | [![Downloads](http://pepy.tech/badge/gluoncv-torch)](http://pepy.tech/project/gluoncv-torch) 5 | # PyTorch-Encoding 6 | # GluonCV-Torch 7 | 8 | Load [GluonCV](https://gluon-cv.mxnet.io/) Models in PyTorch. 9 | Simply `import gluoncvth` to getting better pretrained model than `torchvision`: 10 | 11 | ```python 12 | import gluoncvth as gcv 13 | model = gcv.models.resnet50(pretrained=True) 14 | ``` 15 | 16 | **Installation**: 17 | 18 | ```bash 19 | pip install gluoncv-torch 20 | ``` 21 | 22 | 23 | ## Available Models 24 | 25 | ### ImageNet 26 | 27 | ImageNet models single-crop error rates, comparing to the `torchvision` models: 28 | 29 | | | torchvision | | gluoncvth | | 30 | |---------------------------------|-----------------|---------------|---------------|-------------| 31 | | Model | Top-1 error | Top-5 error | Top-1 error | Top-5 error | 32 | | [ResNet18](#resnet) | 30.24 | 10.92 | 29.06 | 10.17 | 33 | | [ResNet34](#resnet) | 26.70 | 8.58 | 25.35 | 7.92 | 34 | | [ResNet50](#resnet) | 23.85 | 7.13 | 22.33 | 6.18 | 35 | | [ResNet101](#resnet) | 22.63 | 6.44 | 20.80 | 5.39 | 36 | | [ResNet152](#resnet) | 21.69 | 5.94 | 20.56 | 5.39 | 37 | | Inception v3 | 22.55 | 6.44 | 21.33 | 5.61 | 38 | 39 | More models available at [GluonCV Image Classification ModelZoo](https://gluon-cv.mxnet.io/model_zoo/classification.html#imagenet) 40 | 41 | ### Semantic Segmentation 42 | 43 | Results on Pascal VOC dataset: 44 | 45 | | Model | Base Network | mIoU | 46 | |-------------------------|---------------|------------| 47 | | [FCN](#fcn) | ResNet101 | 83.6 | 48 | | [PSPNet](#pspnet) | ResNet101 | 85.1 | 49 | | [DeepLabV3](#deeplabv3) | ResNet101 | 86.2 | 50 | 51 | Results on ADE20K dataset: 52 | 53 | | Model | Base Network | PixAcc | mIoU | 54 | |-------------------------|---------------|-----------|------------| 55 | | [FCN](#fcn) | ResNet101 | 80.6 | 41.6 | 56 | | [PSPNet](#pspnet) | ResNet101 | 80.8 | 42.9 | 57 | | [DeepLabV3](#deeplabv3) | ResNet101 | 81.1 | 44.1 | 58 | 59 | **Quick Demo** 60 | 61 | ```python 62 | import torch 63 | import gluoncvth 64 | 65 | # Get the model 66 | model = gluoncvth.models.get_deeplab_resnet101_ade(pretrained=True) 67 | model.eval() 68 | 69 | # Prepare the image 70 | url = 'https://github.com/zhanghang1989/image-data/blob/master/encoding/' + \ 71 | 'segmentation/ade20k/ADE_val_00001142.jpg?raw=true' 72 | filename = 'example.jpg' 73 | img = gluoncvth.utils.load_image( 74 | gluoncvth.utils.download(url, filename)).unsqueeze(0) 75 | 76 | # Make prediction 77 | output = model.evaluate(img) 78 | predict = torch.max(output, 1)[1].cpu().numpy() + 1 79 | 80 | # Get color pallete for visualization 81 | mask = gluoncvth.utils.get_mask_pallete(predict, 'ade20k') 82 | mask.save('output.png') 83 | ``` 84 | 85 | ![](./image/demo_deeplab_ade.png) 86 | 87 | 88 | More models available at [GluonCV Semantic Segmentation ModelZoo](https://gluon-cv.mxnet.io/model_zoo/segmentation.html) 89 | 90 | ## API Reference 91 | 92 | ### ResNet 93 | 94 | - `gluoncvth.models.resnet18(pretrained=True)` 95 | - `gluoncvth.models.resnet34(pretrained=True)` 96 | - `gluoncvth.models.resnet50(pretrained=True)` 97 | - `gluoncvth.models.resnet101(pretrained=True)` 98 | - `gluoncvth.models.resnet152(pretrained=True)` 99 | 100 | ### FCN 101 | 102 | - `gluoncvth.models.get_fcn_resnet101_voc(pretrained=True)` 103 | - `gluoncvth.models.get_fcn_resnet101_ade(pretrained=True)` 104 | 105 | ### PSPNet 106 | 107 | - `gluoncvth.models.get_psp_resnet101_voc(pretrained=True)` 108 | - `gluoncvth.models.get_psp_resnet101_ade(pretrained=True)` 109 | 110 | ### DeepLabV3 111 | 112 | - `gluoncvth.models.get_deeplab_resnet101_voc(pretrained=True)` 113 | - `gluoncvth.models.get_deeplab_resnet101_ade(pretrained=True)` 114 | 115 | ### 116 | 117 | ## Why [GluonCV](https://gluon-cv.mxnet.io/)? 118 | 119 | **1. State-of-the-art Implementations** 120 | 121 | **2. Pretrained Models and Tutorials** 122 | 123 | **3. Community Support** 124 | 125 | We expect this PyTorch inference API for GluonCV models will be beneficial to the entire computer vision comunity. 126 | -------------------------------------------------------------------------------- /gluoncvth/__init__.py: -------------------------------------------------------------------------------- 1 | from . import models 2 | from . import utils 3 | -------------------------------------------------------------------------------- /gluoncvth/models/__init__.py: -------------------------------------------------------------------------------- 1 | from . import model_zoo 2 | from .resnet import * 3 | from .fcn import * 4 | from .pspnet import * 5 | from .deeplab import * 6 | -------------------------------------------------------------------------------- /gluoncvth/models/base.py: -------------------------------------------------------------------------------- 1 | 2 | import torch.nn as nn 3 | from . import resnet 4 | 5 | up_kwargs = {'mode': 'bilinear', 'align_corners': True} 6 | 7 | __all__ = ['BaseNet', 'acronyms'] 8 | 9 | acronyms = { 10 | 'coco': 'coco', 11 | 'pascal_voc': 'voc', 12 | 'pascal_aug': 'voc', 13 | 'ade20k': 'ade', 14 | 'citys': 'citys', 15 | 'minc': 'minc', 16 | } 17 | 18 | nclass = { 19 | 'coco': 21, 20 | 'pascal_voc': 21, 21 | 'pascal_aug': 21, 22 | 'ade20k': 150, 23 | 'pcontext': 59, 24 | 'citys': 19, 25 | } 26 | 27 | class BaseNet(nn.Module): 28 | def __init__(self, nclass, backbone, aux, se_loss, dilated=True, norm_layer=None, 29 | base_size=520, crop_size=480, mean=[.485, .456, .406], 30 | std=[.229, .224, .225], root='~/.gluoncvth/models'): 31 | super(BaseNet, self).__init__() 32 | self.nclass = nclass 33 | self.aux = aux 34 | self.se_loss = se_loss 35 | self.mean = mean 36 | self.std = std 37 | self.base_size = base_size 38 | self.crop_size = crop_size 39 | # copying modules from pretrained models 40 | if backbone == 'resnet50': 41 | self.pretrained = resnet.resnet50(pretrained=False, dilated=dilated, deep_base=True, 42 | norm_layer=norm_layer, root=root) 43 | elif backbone == 'resnet101': 44 | self.pretrained = resnet.resnet101(pretrained=False, dilated=dilated, deep_base=True, 45 | norm_layer=norm_layer, root=root) 46 | elif backbone == 'resnet152': 47 | self.pretrained = resnet.resnet152(pretrained=False, dilated=dilated, deep_base=True, 48 | norm_layer=norm_layer, root=root) 49 | else: 50 | raise RuntimeError('unknown backbone: {}'.format(backbone)) 51 | # bilinear upsample options 52 | self._up_kwargs = up_kwargs 53 | 54 | def base_forward(self, x): 55 | x = self.pretrained.conv1(x) 56 | x = self.pretrained.bn1(x) 57 | x = self.pretrained.relu(x) 58 | x = self.pretrained.maxpool(x) 59 | c1 = self.pretrained.layer1(x) 60 | c2 = self.pretrained.layer2(c1) 61 | c3 = self.pretrained.layer3(c2) 62 | c4 = self.pretrained.layer4(c3) 63 | return c1, c2, c3, c4 64 | 65 | def evaluate(self, x): 66 | pred = self.forward(x) 67 | if isinstance(pred, (tuple, list)): 68 | pred = pred[0] 69 | return pred 70 | -------------------------------------------------------------------------------- /gluoncvth/models/deeplab.py: -------------------------------------------------------------------------------- 1 | from __future__ import division 2 | import os 3 | import numpy as np 4 | import torch 5 | import torch.nn as nn 6 | from torch.nn.functional import interpolate 7 | 8 | from .base import BaseNet 9 | from .fcn import FCNHead 10 | 11 | class DeepLabV3(BaseNet): 12 | def __init__(self, nclass, backbone, aux=True, se_loss=False, norm_layer=nn.BatchNorm2d, **kwargs): 13 | super(DeepLabV3, self).__init__(nclass, backbone, aux, se_loss, norm_layer=norm_layer, **kwargs) 14 | self.head = DeepLabV3Head(2048, nclass, norm_layer, self._up_kwargs) 15 | if aux: 16 | self.auxlayer = FCNHead(1024, nclass, norm_layer) 17 | 18 | def forward(self, x): 19 | _, _, h, w = x.size() 20 | _, _, c3, c4 = self.base_forward(x) 21 | 22 | outputs = [] 23 | x = self.head(c4) 24 | x = interpolate(x, (h,w), **self._up_kwargs) 25 | outputs.append(x) 26 | if self.aux: 27 | auxout = self.auxlayer(c3) 28 | auxout = interpolate(auxout, (h,w), **self._up_kwargs) 29 | outputs.append(auxout) 30 | return tuple(outputs) 31 | 32 | 33 | class DeepLabV3Head(nn.Module): 34 | def __init__(self, in_channels, out_channels, norm_layer, up_kwargs, atrous_rates=[12, 24, 36], **kwargs): 35 | super(DeepLabV3Head, self).__init__() 36 | inter_channels = in_channels // 8 37 | self.aspp = ASPP_Module(in_channels, atrous_rates, norm_layer, up_kwargs, **kwargs) 38 | self.block = nn.Sequential( 39 | nn.Conv2d(inter_channels, inter_channels, 3, padding=1, bias=False), 40 | norm_layer(inter_channels), 41 | nn.ReLU(True), 42 | nn.Dropout2d(0.1, False), 43 | nn.Conv2d(inter_channels, out_channels, 1)) 44 | 45 | def forward(self, x): 46 | x = self.aspp(x) 47 | x = self.block(x) 48 | return x 49 | 50 | 51 | def ASPPConv(in_channels, out_channels, atrous_rate, norm_layer): 52 | block = nn.Sequential( 53 | nn.Conv2d(in_channels, out_channels, 3, padding=atrous_rate, 54 | dilation=atrous_rate, bias=False), 55 | norm_layer(out_channels), 56 | nn.ReLU(True)) 57 | return block 58 | 59 | class AsppPooling(nn.Module): 60 | def __init__(self, in_channels, out_channels, norm_layer, up_kwargs): 61 | super(AsppPooling, self).__init__() 62 | self._up_kwargs = up_kwargs 63 | self.gap = nn.Sequential(nn.AdaptiveAvgPool2d(1), 64 | nn.Conv2d(in_channels, out_channels, 1, bias=False), 65 | norm_layer(out_channels), 66 | nn.ReLU(True)) 67 | 68 | def forward(self, x): 69 | _, _, h, w = x.size() 70 | pool = self.gap(x) 71 | return interpolate(pool, (h,w), **self._up_kwargs) 72 | 73 | class ASPP_Module(nn.Module): 74 | def __init__(self, in_channels, atrous_rates, norm_layer, up_kwargs): 75 | super(ASPP_Module, self).__init__() 76 | out_channels = in_channels // 8 77 | rate1, rate2, rate3 = tuple(atrous_rates) 78 | self.b0 = nn.Sequential( 79 | nn.Conv2d(in_channels, out_channels, 1, bias=False), 80 | norm_layer(out_channels), 81 | nn.ReLU(True)) 82 | self.b1 = ASPPConv(in_channels, out_channels, rate1, norm_layer) 83 | self.b2 = ASPPConv(in_channels, out_channels, rate2, norm_layer) 84 | self.b3 = ASPPConv(in_channels, out_channels, rate3, norm_layer) 85 | self.b4 = AsppPooling(in_channels, out_channels, norm_layer, up_kwargs) 86 | 87 | self.project = nn.Sequential( 88 | nn.Conv2d(5*out_channels, out_channels, 1, bias=False), 89 | norm_layer(out_channels), 90 | nn.ReLU(True), 91 | nn.Dropout2d(0.5, False)) 92 | 93 | def forward(self, x): 94 | feat0 = self.b0(x) 95 | feat1 = self.b1(x) 96 | feat2 = self.b2(x) 97 | feat3 = self.b3(x) 98 | feat4 = self.b4(x) 99 | y = torch.cat((feat0, feat1, feat2, feat3, feat4), 1) 100 | return self.project(y) 101 | 102 | def get_deeplab(dataset='pascal_voc', backbone='resnet50', pretrained=False, 103 | root='~/.gluoncvth/models', **kwargs): 104 | # infer number of classes 105 | from .base import nclass, acronyms 106 | model = DeepLabV3(nclass[dataset.lower()], backbone=backbone, root=root, **kwargs) 107 | if pretrained: 108 | from .model_store import get_model_file 109 | model.load_state_dict(torch.load( 110 | get_model_file('deeplab_%s_%s'%(backbone, acronyms[dataset]), root=root))) 111 | return model 112 | 113 | def get_deeplab_resnet101_voc(pretrained=False, root='~/.gluoncvth/models', **kwargs): 114 | r"""DeepLabV3 model from the paper `"Context Encoding for Semantic Segmentation" 115 | `_ 116 | 117 | Parameters 118 | ---------- 119 | pretrained : bool, default False 120 | Whether to load the pretrained weights for model. 121 | root : str, default '~/.gluoncvth/models' 122 | Location for keeping the model parameters. 123 | 124 | 125 | Examples 126 | -------- 127 | >>> model = get_deeplab_resnet50_ade(pretrained=True) 128 | >>> print(model) 129 | """ 130 | return get_deeplab('pascal_voc', 'resnet101', pretrained, root=root, **kwargs) 131 | 132 | 133 | def get_deeplab_resnet101_ade(pretrained=False, root='~/.gluoncvth/models', **kwargs): 134 | r"""DeepLabV3 model from the paper `"Context Encoding for Semantic Segmentation" 135 | `_ 136 | 137 | Parameters 138 | ---------- 139 | pretrained : bool, default False 140 | Whether to load the pretrained weights for model. 141 | root : str, default '~/.gluoncvth/models' 142 | Location for keeping the model parameters. 143 | 144 | 145 | Examples 146 | -------- 147 | >>> model = get_deeplab_resnet50_ade(pretrained=True) 148 | >>> print(model) 149 | """ 150 | return get_deeplab('ade20k', 'resnet101', pretrained, root=root, **kwargs) 151 | 152 | -------------------------------------------------------------------------------- /gluoncvth/models/fcn.py: -------------------------------------------------------------------------------- 1 | from __future__ import division 2 | import os 3 | import numpy as np 4 | import torch 5 | import torch.nn as nn 6 | from torch.nn.functional import interpolate 7 | 8 | from .base import BaseNet 9 | 10 | __all__ = ['FCN', 'get_fcn', 'get_fcn_resnet101_voc', 'get_fcn_resnet101_ade'] 11 | 12 | class FCN(BaseNet): 13 | r"""Fully Convolutional Networks for Semantic Segmentation 14 | 15 | Parameters 16 | ---------- 17 | nclass : int 18 | Number of categories for the training dataset. 19 | backbone : string 20 | Pre-trained dilated backbone network type (default:'resnet50'; 'resnet50', 21 | 'resnet101' or 'resnet152'). 22 | norm_layer : object 23 | Normalization layer used in backbone network (default: :class:`mxnet.gluon.nn.BatchNorm`; 24 | 25 | 26 | Reference: 27 | 28 | Long, Jonathan, Evan Shelhamer, and Trevor Darrell. "Fully convolutional networks 29 | for semantic segmentation." *CVPR*, 2015 30 | 31 | Examples 32 | -------- 33 | >>> model = FCN(nclass=21, backbone='resnet50') 34 | >>> print(model) 35 | """ 36 | def __init__(self, nclass, backbone, aux=True, se_loss=False, norm_layer=nn.BatchNorm2d, **kwargs): 37 | super(FCN, self).__init__(nclass, backbone, aux, se_loss, norm_layer=norm_layer, **kwargs) 38 | self.head = FCNHead(2048, nclass, norm_layer) 39 | if aux: 40 | self.auxlayer = FCNHead(1024, nclass, norm_layer) 41 | 42 | def forward(self, x): 43 | imsize = x.size()[2:] 44 | _, _, c3, c4 = self.base_forward(x) 45 | 46 | x = self.head(c4) 47 | x = interpolate(x, imsize, **self._up_kwargs) 48 | outputs = [x] 49 | if self.aux: 50 | auxout = self.auxlayer(c3) 51 | auxout = interpolate(auxout, imsize, **self._up_kwargs) 52 | outputs.append(auxout) 53 | return tuple(outputs) 54 | 55 | 56 | class FCNHead(nn.Module): 57 | def __init__(self, in_channels, out_channels, norm_layer): 58 | super(FCNHead, self).__init__() 59 | inter_channels = in_channels // 4 60 | self.conv5 = nn.Sequential(nn.Conv2d(in_channels, inter_channels, 3, padding=1, bias=False), 61 | norm_layer(inter_channels), 62 | nn.ReLU(), 63 | nn.Dropout2d(0.1, False), 64 | nn.Conv2d(inter_channels, out_channels, 1)) 65 | 66 | def forward(self, x): 67 | return self.conv5(x) 68 | 69 | 70 | def get_fcn(dataset='pascal_voc', backbone='resnet50', pretrained=True, 71 | root='~/.gluoncvth/models', **kwargs): 72 | r"""FCN model from the paper `"Fully Convolutional Network for semantic segmentation" 73 | `_ 74 | Parameters 75 | ---------- 76 | dataset : str, default pascal_voc 77 | The dataset that model pretrained on. (pascal_voc, ade20k) 78 | pretrained : bool, default False 79 | Whether to load the pretrained weights for model. 80 | root : str, default '~/.gluoncvth/models' 81 | Location for keeping the model parameters. 82 | Examples 83 | -------- 84 | >>> model = get_fcn(dataset='pascal_voc', backbone='resnet50', pretrained=True) 85 | >>> print(model) 86 | """ 87 | # infer number of classes 88 | from .base import nclass, acronyms 89 | model = FCN(nclass[dataset.lower()], backbone=backbone, **kwargs) 90 | if pretrained: 91 | from .model_store import get_model_file 92 | model.load_state_dict(torch.load( 93 | get_model_file('fcn_%s_%s'%(backbone, acronyms[dataset]), root=root))) 94 | return model 95 | 96 | def get_fcn_resnet101_voc(pretrained=True, root='~/.gluoncvth/models', **kwargs): 97 | r"""EncNet-PSP model from the paper `"Context Encoding for Semantic Segmentation" 98 | `_ 99 | 100 | Parameters 101 | ---------- 102 | pretrained : bool, default False 103 | Whether to load the pretrained weights for model. 104 | root : str, default '~/.gluoncvth/models' 105 | Location for keeping the model parameters. 106 | 107 | 108 | Examples 109 | -------- 110 | >>> model = get_fcn_resnet50_voc(pretrained=True) 111 | >>> print(model) 112 | """ 113 | return get_fcn('pascal_voc', 'resnet101', pretrained, root=root, **kwargs) 114 | 115 | def get_fcn_resnet101_ade(pretrained=True, root='~/.gluoncvth/models', **kwargs): 116 | r"""EncNet-PSP model from the paper `"Context Encoding for Semantic Segmentation" 117 | `_ 118 | 119 | Parameters 120 | ---------- 121 | pretrained : bool, default False 122 | Whether to load the pretrained weights for model. 123 | root : str, default '~/.gluoncvth/models' 124 | Location for keeping the model parameters. 125 | 126 | 127 | Examples 128 | -------- 129 | >>> model = get_fcn_resnet50_ade(pretrained=True) 130 | >>> print(model) 131 | """ 132 | return get_fcn('ade20k', 'resnet101', pretrained, root=root, **kwargs) 133 | -------------------------------------------------------------------------------- /gluoncvth/models/model_store.py: -------------------------------------------------------------------------------- 1 | """Model store which provides pretrained models.""" 2 | from __future__ import print_function 3 | __all__ = ['get_model_file', 'purge'] 4 | import os 5 | import zipfile 6 | 7 | from ..utils import download, check_sha1 8 | 9 | _model_sha1 = {name: checksum for checksum, name in [ 10 | ('7591854d34e97010f019e9f98f9aed9c4a463d58', 'resnet18'), 11 | ('64557eb2096a56ff2db4fdd637e00ede804c1fec', 'resnet34'), 12 | ('0ef8ed2db4162747fecb34ee542b944d141b3ef1', 'resnet50'), 13 | ('1834038c51dd60b2819b4391acfec0bb4525f986', 'resnet101'), 14 | ('990926f3c93c67aea2342d1a5b88ba63dfee32f4', 'resnet152'), 15 | ('357fb3777da3ebdf13ab06bee51fa6f83837967c', 'fcn_resnet101_voc'), 16 | ('8bb3bccd02da0e5431a616d3abe7e8c383e8f587', 'fcn_resnet101_ade'), 17 | ('6d90aaae73a3adcb20f186895b27bf45368601ab', 'psp_resnet101_voc'), 18 | ('fe990f00dda51d58718c43cf4705e0a61ca15ef0', 'psp_resnet101_ade'), 19 | ('5c25b7db003fb805df6574139bf04e1b85f0f37d', 'deeplab_resnet101_voc'), 20 | ('c0d88de54f3abbc358038c248f0863bef96fb0d4', 'deeplab_resnet101_ade'), 21 | ]} 22 | 23 | gluoncvth_repo_url = 'https://s3.us-west-1.wasabisys.com/resnest' 24 | _url_format = '{repo_url}gluoncvth/{file_name}.zip' 25 | 26 | def short_hash(name): 27 | if name not in _model_sha1: 28 | raise ValueError('Pretrained model for {name} is not available.'.format(name=name)) 29 | return _model_sha1[name][:8] 30 | 31 | def get_model_file(name, root=os.path.join('~', '.gluoncvth', 'models')): 32 | r"""Return location for the pretrained on local file system. 33 | 34 | This function will download from online model zoo when model cannot be found or has mismatch. 35 | The root directory will be created if it doesn't exist. 36 | 37 | Parameters 38 | ---------- 39 | name : str 40 | Name of the model. 41 | root : str, default '~/.gluoncvth/models' 42 | Location for keeping the model parameters. 43 | 44 | Returns 45 | ------- 46 | file_path 47 | Path to the requested pretrained model file. 48 | """ 49 | file_name = '{name}-{short_hash}'.format(name=name, short_hash=short_hash(name)) 50 | root = os.path.expanduser(root) 51 | file_path = os.path.join(root, file_name+'.pth') 52 | sha1_hash = _model_sha1[name] 53 | if os.path.exists(file_path): 54 | if check_sha1(file_path, sha1_hash): 55 | return file_path 56 | else: 57 | print('Mismatch in the content of model file {} detected.' + 58 | ' Downloading again.'.format(file_path)) 59 | else: 60 | print('Model file {} is not found. Downloading.'.format(file_path)) 61 | 62 | if not os.path.exists(root): 63 | os.makedirs(root) 64 | 65 | zip_file_path = os.path.join(root, file_name+'.zip') 66 | repo_url = os.environ.get('ENCODING_REPO', gluoncvth_repo_url) 67 | if repo_url[-1] != '/': 68 | repo_url = repo_url + '/' 69 | download(_url_format.format(repo_url=repo_url, file_name=file_name), 70 | path=zip_file_path, 71 | overwrite=True) 72 | with zipfile.ZipFile(zip_file_path) as zf: 73 | zf.extractall(root) 74 | os.remove(zip_file_path) 75 | 76 | if check_sha1(file_path, sha1_hash): 77 | return file_path 78 | else: 79 | raise ValueError('Downloaded file has different hash. Please try again.') 80 | 81 | def purge(root=os.path.join('~', '.gluoncvth', 'models')): 82 | r"""Purge all pretrained model files in local file store. 83 | 84 | Parameters 85 | ---------- 86 | root : str, default '~/.gluoncvth/models' 87 | Location for keeping the model parameters. 88 | """ 89 | root = os.path.expanduser(root) 90 | files = os.listdir(root) 91 | for f in files: 92 | if f.endswith(".pth"): 93 | os.remove(os.path.join(root, f)) 94 | 95 | def pretrained_model_list(): 96 | return list(_model_sha1.keys()) 97 | -------------------------------------------------------------------------------- /gluoncvth/models/model_zoo.py: -------------------------------------------------------------------------------- 1 | # pylint: disable=wildcard-import, unused-wildcard-import 2 | 3 | from .resnet import * 4 | 5 | __all__ = ['get_model'] 6 | 7 | 8 | def get_model(name, **kwargs): 9 | """Returns a pre-defined model by name 10 | 11 | Parameters 12 | ---------- 13 | name : str 14 | Name of the model. 15 | pretrained : bool 16 | Whether to load the pretrained weights for model. 17 | root : str, default '~/.gluoncvth/models' 18 | Location for keeping the model parameters. 19 | 20 | Returns 21 | ------- 22 | Module: 23 | The model. 24 | """ 25 | models = { 26 | 'resnet18': resnet18, 27 | 'resnet34': resnet34, 28 | 'resnet50': resnet50, 29 | 'resnet101': resnet101, 30 | 'resnet152': resnet152, 31 | } 32 | name = name.lower() 33 | if name not in models: 34 | raise ValueError('%s\n\t%s' % (str(e), '\n\t'.join(sorted(models.keys())))) 35 | net = models[name](**kwargs) 36 | return net 37 | -------------------------------------------------------------------------------- /gluoncvth/models/pspnet.py: -------------------------------------------------------------------------------- 1 | from __future__ import division 2 | import os 3 | import numpy as np 4 | import torch 5 | import torch.nn as nn 6 | import torch.nn.functional as F 7 | 8 | from .base import BaseNet 9 | from .fcn import FCNHead 10 | 11 | class PSP(BaseNet): 12 | def __init__(self, nclass, backbone, aux=True, se_loss=False, norm_layer=nn.BatchNorm2d, **kwargs): 13 | super(PSP, self).__init__(nclass, backbone, aux, se_loss, norm_layer=norm_layer, **kwargs) 14 | self.head = PSPHead(2048, nclass, norm_layer, self._up_kwargs) 15 | if aux: 16 | self.auxlayer = FCNHead(1024, nclass, norm_layer) 17 | 18 | def forward(self, x): 19 | _, _, h, w = x.size() 20 | _, _, c3, c4 = self.base_forward(x) 21 | 22 | outputs = [] 23 | x = self.head(c4) 24 | x = F.interpolate(x, (h,w), **self._up_kwargs) 25 | outputs.append(x) 26 | if self.aux: 27 | auxout = self.auxlayer(c3) 28 | auxout = F.interpolate(auxout, (h,w), **self._up_kwargs) 29 | outputs.append(auxout) 30 | return tuple(outputs) 31 | 32 | class PyramidPooling(nn.Module): 33 | """ 34 | Reference: 35 | Zhao, Hengshuang, et al. *"Pyramid scene parsing network."* 36 | """ 37 | def __init__(self, in_channels, norm_layer, up_kwargs): 38 | super(PyramidPooling, self).__init__() 39 | self.pool1 = nn.AdaptiveAvgPool2d(1) 40 | self.pool2 = nn.AdaptiveAvgPool2d(2) 41 | self.pool3 = nn.AdaptiveAvgPool2d(3) 42 | self.pool4 = nn.AdaptiveAvgPool2d(6) 43 | 44 | out_channels = int(in_channels/4) 45 | self.conv1 = nn.Sequential(nn.Conv2d(in_channels, out_channels, 1, bias=False), 46 | norm_layer(out_channels), 47 | nn.ReLU(True)) 48 | self.conv2 = nn.Sequential(nn.Conv2d(in_channels, out_channels, 1, bias=False), 49 | norm_layer(out_channels), 50 | nn.ReLU(True)) 51 | self.conv3 = nn.Sequential(nn.Conv2d(in_channels, out_channels, 1, bias=False), 52 | norm_layer(out_channels), 53 | nn.ReLU(True)) 54 | self.conv4 = nn.Sequential(nn.Conv2d(in_channels, out_channels, 1, bias=False), 55 | norm_layer(out_channels), 56 | nn.ReLU(True)) 57 | # bilinear interpolate options 58 | self._up_kwargs = up_kwargs 59 | 60 | def forward(self, x): 61 | _, _, h, w = x.size() 62 | feat1 = F.interpolate(self.conv1(self.pool1(x)), (h, w), **self._up_kwargs) 63 | feat2 = F.interpolate(self.conv2(self.pool2(x)), (h, w), **self._up_kwargs) 64 | feat3 = F.interpolate(self.conv3(self.pool3(x)), (h, w), **self._up_kwargs) 65 | feat4 = F.interpolate(self.conv4(self.pool4(x)), (h, w), **self._up_kwargs) 66 | return torch.cat((x, feat1, feat2, feat3, feat4), 1) 67 | 68 | class PSPHead(nn.Module): 69 | def __init__(self, in_channels, out_channels, norm_layer, up_kwargs): 70 | super(PSPHead, self).__init__() 71 | inter_channels = in_channels // 4 72 | self.conv5 = nn.Sequential(PyramidPooling(in_channels, norm_layer, up_kwargs), 73 | nn.Conv2d(in_channels * 2, inter_channels, 3, padding=1, bias=False), 74 | norm_layer(inter_channels), 75 | nn.ReLU(True), 76 | nn.Dropout2d(0.1, False), 77 | nn.Conv2d(inter_channels, out_channels, 1)) 78 | 79 | def forward(self, x): 80 | return self.conv5(x) 81 | 82 | def get_psp(dataset='pascal_voc', backbone='resnet50', pretrained=False, 83 | root='~/.gluoncvth/models', **kwargs): 84 | r"""PSP model from the paper `"Context Encoding for Semantic Segmentation" 85 | `_ 86 | Parameters 87 | ---------- 88 | dataset : str, default pascal_voc 89 | The dataset that model pretrained on. (pascal_voc, ade20k) 90 | pretrained : bool, default False 91 | Whether to load the pretrained weights for model. 92 | root : str, default '~/.gluoncvth/models' 93 | Location for keeping the model parameters. 94 | Examples 95 | -------- 96 | >>> model = get_psp(dataset='pascal_voc', backbone='resnet50', pretrained=True) 97 | >>> print(model) 98 | """ 99 | # infer number of classes 100 | from .base import nclass, acronyms 101 | model = PSP(nclass[dataset.lower()], backbone=backbone, root=root, **kwargs) 102 | if pretrained: 103 | from .model_store import get_model_file 104 | model.load_state_dict(torch.load( 105 | get_model_file('psp_%s_%s'%(backbone, acronyms[dataset]), root=root))) 106 | return model 107 | 108 | def get_psp_resnet101_voc(pretrained=False, root='~/.gluoncvth/models', **kwargs): 109 | r"""PSP model from the paper `"Context Encoding for Semantic Segmentation" 110 | `_ 111 | 112 | Parameters 113 | ---------- 114 | pretrained : bool, default False 115 | Whether to load the pretrained weights for model. 116 | root : str, default '~/.gluoncvth/models' 117 | Location for keeping the model parameters. 118 | 119 | 120 | Examples 121 | -------- 122 | >>> model = get_psp_resnet50_ade(pretrained=True) 123 | >>> print(model) 124 | """ 125 | return get_psp('pascal_voc', 'resnet101', pretrained, root=root, **kwargs) 126 | 127 | def get_psp_resnet101_ade(pretrained=False, root='~/.gluoncvth/models', **kwargs): 128 | r"""PSP model from the paper `"Context Encoding for Semantic Segmentation" 129 | `_ 130 | 131 | Parameters 132 | ---------- 133 | pretrained : bool, default False 134 | Whether to load the pretrained weights for model. 135 | root : str, default '~/.gluoncvth/models' 136 | Location for keeping the model parameters. 137 | 138 | 139 | Examples 140 | -------- 141 | >>> model = get_psp_resnet50_ade(pretrained=True) 142 | >>> print(model) 143 | """ 144 | return get_psp('ade20k', 'resnet101', pretrained, root=root, **kwargs) 145 | 146 | -------------------------------------------------------------------------------- /gluoncvth/models/resnet.py: -------------------------------------------------------------------------------- 1 | """Dilated ResNet""" 2 | import math 3 | import torch 4 | import torch.utils.model_zoo as model_zoo 5 | import torch.nn as nn 6 | from ..models.model_store import get_model_file 7 | 8 | __all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 9 | 'resnet152'] 10 | 11 | def conv3x3(in_planes, out_planes, stride=1): 12 | "3x3 convolution with padding" 13 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, 14 | padding=1, bias=False) 15 | 16 | 17 | class BasicBlock(nn.Module): 18 | """ResNet BasicBlock 19 | """ 20 | expansion = 1 21 | def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None, previous_dilation=1, 22 | norm_layer=None): 23 | super(BasicBlock, self).__init__() 24 | self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=3, stride=stride, 25 | padding=dilation, dilation=dilation, bias=False) 26 | self.bn1 = norm_layer(planes) 27 | self.relu = nn.ReLU(inplace=True) 28 | self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=1, 29 | padding=previous_dilation, dilation=previous_dilation, bias=False) 30 | self.bn2 = norm_layer(planes) 31 | self.downsample = downsample 32 | self.stride = stride 33 | 34 | def forward(self, x): 35 | residual = x 36 | 37 | out = self.conv1(x) 38 | out = self.bn1(out) 39 | out = self.relu(out) 40 | 41 | out = self.conv2(out) 42 | out = self.bn2(out) 43 | 44 | if self.downsample is not None: 45 | residual = self.downsample(x) 46 | 47 | out += residual 48 | out = self.relu(out) 49 | 50 | return out 51 | 52 | 53 | class Bottleneck(nn.Module): 54 | """ResNet Bottleneck 55 | """ 56 | # pylint: disable=unused-argument 57 | expansion = 4 58 | def __init__(self, inplanes, planes, stride=1, dilation=1, 59 | downsample=None, previous_dilation=1, norm_layer=None): 60 | super(Bottleneck, self).__init__() 61 | self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) 62 | self.bn1 = norm_layer(planes) 63 | self.conv2 = nn.Conv2d( 64 | planes, planes, kernel_size=3, stride=stride, 65 | padding=dilation, dilation=dilation, bias=False) 66 | self.bn2 = norm_layer(planes) 67 | self.conv3 = nn.Conv2d( 68 | planes, planes * 4, kernel_size=1, bias=False) 69 | self.bn3 = norm_layer(planes * 4) 70 | self.relu = nn.ReLU(inplace=True) 71 | self.downsample = downsample 72 | self.dilation = dilation 73 | self.stride = stride 74 | 75 | def _sum_each(self, x, y): 76 | assert(len(x) == len(y)) 77 | z = [] 78 | for i in range(len(x)): 79 | z.append(x[i]+y[i]) 80 | return z 81 | 82 | def forward(self, x): 83 | residual = x 84 | 85 | out = self.conv1(x) 86 | out = self.bn1(out) 87 | out = self.relu(out) 88 | 89 | out = self.conv2(out) 90 | out = self.bn2(out) 91 | out = self.relu(out) 92 | 93 | out = self.conv3(out) 94 | out = self.bn3(out) 95 | 96 | if self.downsample is not None: 97 | residual = self.downsample(x) 98 | 99 | out += residual 100 | out = self.relu(out) 101 | 102 | return out 103 | 104 | 105 | class ResNet(nn.Module): 106 | """Dilated Pre-trained ResNet Model, which preduces the stride of 8 featuremaps at conv5. 107 | 108 | Parameters 109 | ---------- 110 | block : Block 111 | Class for the residual block. Options are BasicBlockV1, BottleneckV1. 112 | layers : list of int 113 | Numbers of layers in each block 114 | classes : int, default 1000 115 | Number of classification classes. 116 | dilated : bool, default False 117 | Applying dilation strategy to pretrained ResNet yielding a stride-8 model, 118 | typically used in Semantic Segmentation. 119 | norm_layer : object 120 | Normalization layer used in backbone network (default: :class:`mxnet.gluon.nn.BatchNorm`; 121 | for Synchronized Cross-GPU BachNormalization). 122 | 123 | Reference: 124 | 125 | - He, Kaiming, et al. "Deep residual learning for image recognition." Proceedings of the IEEE conference on computer vision and pattern recognition. 2016. 126 | 127 | - Yu, Fisher, and Vladlen Koltun. "Multi-scale context aggregation by dilated convolutions." 128 | """ 129 | # pylint: disable=unused-variable 130 | def __init__(self, block, layers, num_classes=1000, dilated=True, 131 | deep_base=False, norm_layer=nn.BatchNorm2d): 132 | self.inplanes = 128 if deep_base else 64 133 | super(ResNet, self).__init__() 134 | if deep_base: 135 | self.conv1 = nn.Sequential( 136 | nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, bias=False), 137 | norm_layer(64), 138 | nn.ReLU(inplace=True), 139 | nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False), 140 | norm_layer(64), 141 | nn.ReLU(inplace=True), 142 | nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1, bias=False), 143 | ) 144 | else: 145 | self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, 146 | bias=False) 147 | self.bn1 = norm_layer(self.inplanes) 148 | self.relu = nn.ReLU(inplace=True) 149 | self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) 150 | self.layer1 = self._make_layer(block, 64, layers[0], norm_layer=norm_layer) 151 | self.layer2 = self._make_layer(block, 128, layers[1], stride=2, norm_layer=norm_layer) 152 | if dilated: 153 | self.layer3 = self._make_layer(block, 256, layers[2], stride=1, 154 | dilation=2, norm_layer=norm_layer) 155 | self.layer4 = self._make_layer(block, 512, layers[3], stride=1, 156 | dilation=4, norm_layer=norm_layer) 157 | else: 158 | self.layer3 = self._make_layer(block, 256, layers[2], stride=2, 159 | norm_layer=norm_layer) 160 | self.layer4 = self._make_layer(block, 512, layers[3], stride=2, 161 | norm_layer=norm_layer) 162 | self.avgpool = nn.AvgPool2d(7, stride=1) 163 | self.fc = nn.Linear(512 * block.expansion, num_classes) 164 | 165 | for m in self.modules(): 166 | if isinstance(m, nn.Conv2d): 167 | n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels 168 | m.weight.data.normal_(0, math.sqrt(2. / n)) 169 | elif isinstance(m, norm_layer): 170 | m.weight.data.fill_(1) 171 | m.bias.data.zero_() 172 | 173 | def _make_layer(self, block, planes, blocks, stride=1, dilation=1, norm_layer=None): 174 | downsample = None 175 | if stride != 1 or self.inplanes != planes * block.expansion: 176 | downsample = nn.Sequential( 177 | nn.Conv2d(self.inplanes, planes * block.expansion, 178 | kernel_size=1, stride=stride, bias=False), 179 | norm_layer(planes * block.expansion), 180 | ) 181 | 182 | layers = [] 183 | if dilation == 1 or dilation == 2: 184 | layers.append(block(self.inplanes, planes, stride, dilation=1, 185 | downsample=downsample, previous_dilation=dilation, norm_layer=norm_layer)) 186 | elif dilation == 4: 187 | layers.append(block(self.inplanes, planes, stride, dilation=2, 188 | downsample=downsample, previous_dilation=dilation, norm_layer=norm_layer)) 189 | else: 190 | raise RuntimeError("=> unknown dilation size: {}".format(dilation)) 191 | 192 | self.inplanes = planes * block.expansion 193 | for i in range(1, blocks): 194 | layers.append(block(self.inplanes, planes, dilation=dilation, previous_dilation=dilation, 195 | norm_layer=norm_layer)) 196 | 197 | return nn.Sequential(*layers) 198 | 199 | def forward(self, x): 200 | x = self.conv1(x) 201 | x = self.bn1(x) 202 | x = self.relu(x) 203 | x = self.maxpool(x) 204 | 205 | x = self.layer1(x) 206 | x = self.layer2(x) 207 | x = self.layer3(x) 208 | x = self.layer4(x) 209 | 210 | x = self.avgpool(x) 211 | x = x.view(x.size(0), -1) 212 | x = self.fc(x) 213 | 214 | return x 215 | 216 | def resnet18(pretrained=True, root='~/.gluoncvth/models', **kwargs): 217 | """Constructs a ResNet-18 model. 218 | 219 | Args: 220 | pretrained (bool): If True, returns a model pre-trained on ImageNet 221 | """ 222 | model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs) 223 | if pretrained: 224 | model.load_state_dict(torch.load( 225 | get_model_file('resnet18', root=root)), strict=False) 226 | return model 227 | 228 | 229 | def resnet34(pretrained=True, root='~/.gluoncvth/models', **kwargs): 230 | """Constructs a ResNet-34 model. 231 | 232 | Args: 233 | pretrained (bool): If True, returns a model pre-trained on ImageNet 234 | """ 235 | model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs) 236 | if pretrained: 237 | model.load_state_dict(torch.load( 238 | get_model_file('resnet34', root=root)), strict=False) 239 | return model 240 | 241 | 242 | def resnet50(pretrained=True, root='~/.gluoncvth/models', **kwargs): 243 | """Constructs a ResNet-50 model. 244 | 245 | Args: 246 | pretrained (bool): If True, returns a model pre-trained on ImageNet 247 | """ 248 | model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs) 249 | if pretrained: 250 | model.load_state_dict(torch.load( 251 | get_model_file('resnet50', root=root)), strict=False) 252 | return model 253 | 254 | 255 | def resnet101(pretrained=True, root='~/.gluoncvth/models', **kwargs): 256 | """Constructs a ResNet-101 model. 257 | 258 | Args: 259 | pretrained (bool): If True, returns a model pre-trained on ImageNet 260 | """ 261 | model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs) 262 | if pretrained: 263 | model.load_state_dict(torch.load( 264 | get_model_file('resnet101', root=root)), strict=False) 265 | return model 266 | 267 | 268 | def resnet152(pretrained=True, root='~/.gluoncvth/models', **kwargs): 269 | """Constructs a ResNet-152 model. 270 | 271 | Args: 272 | pretrained (bool): If True, returns a model pre-trained on ImageNet 273 | """ 274 | model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs) 275 | if pretrained: 276 | model.load_state_dict(torch.load( 277 | get_model_file('resnet152', root=root)), strict=False) 278 | return model 279 | -------------------------------------------------------------------------------- /gluoncvth/models/wideresnet.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from collections import OrderedDict 3 | from functools import partial 4 | 5 | import torch 6 | import torch.nn as nn 7 | import torch.nn.functional as F 8 | 9 | from ..models.model_store import get_model_file 10 | 11 | __all__ = ['WideResNet', 'wideresnet38'] 12 | 13 | class ABN(nn.BatchNorm2d): 14 | """Hacky for inference only 15 | """ 16 | def __init__(self, num_features, eps=1e-5, momentum=0.1, sync=True, activation="none", 17 | slope=0.01): 18 | super(ABN, self).__init__(num_features, eps=eps, momentum=momentum, affine=True) 19 | self.slope = slope 20 | 21 | def forward(self, x): 22 | y = F.batch_norm( 23 | x, self.running_mean, self.running_var, self.weight, self.bias, 24 | self.training or not self.track_running_stats, 25 | self.momentum, self.eps) 26 | return F.leaky_relu_(y, self.slope) 27 | 28 | class GlobalAvgPool2d(nn.Module): 29 | def __init__(self): 30 | """Global average pooling over the input's spatial dimensions""" 31 | super(GlobalAvgPool2d, self).__init__() 32 | 33 | def forward(self, inputs): 34 | return F.adaptive_avg_pool2d(inputs, 1).view(inputs.size(0), -1) 35 | 36 | class IdentityResidualBlock(nn.Module): 37 | def __init__(self, 38 | in_channels, 39 | channels, 40 | stride=1, 41 | dilation=1, 42 | groups=1, 43 | norm_act=ABN, 44 | dropout=None): 45 | """Configurable identity-mapping residual block 46 | Parameters 47 | ---------- 48 | in_channels : int 49 | Number of input channels. 50 | channels : list of int 51 | Number of channels in the internal feature maps. Can either have two or three elements: if three construct 52 | a residual block with two `3 x 3` convolutions, otherwise construct a bottleneck block with `1 x 1`, then 53 | `3 x 3` then `1 x 1` convolutions. 54 | stride : int 55 | Stride of the first `3 x 3` convolution 56 | dilation : int 57 | Dilation to apply to the `3 x 3` convolutions. 58 | groups : int 59 | Number of convolution groups. This is used to create ResNeXt-style blocks and is only compatible with 60 | bottleneck blocks. 61 | norm_act : callable 62 | Function to create normalization / activation Module. 63 | dropout: callable 64 | Function to create Dropout Module. 65 | """ 66 | super(IdentityResidualBlock, self).__init__() 67 | 68 | # Check parameters for inconsistencies 69 | if len(channels) != 2 and len(channels) != 3: 70 | raise ValueError("channels must contain either two or three values") 71 | if len(channels) == 2 and groups != 1: 72 | raise ValueError("groups > 1 are only valid if len(channels) == 3") 73 | 74 | is_bottleneck = len(channels) == 3 75 | need_proj_conv = stride != 1 or in_channels != channels[-1] 76 | 77 | self.bn1 = norm_act(in_channels) 78 | if not is_bottleneck: 79 | layers = [ 80 | ("conv1", nn.Conv2d(in_channels, channels[0], 3, stride=stride, padding=dilation, bias=False, 81 | dilation=dilation)), 82 | ("bn2", norm_act(channels[0])), 83 | ("conv2", nn.Conv2d(channels[0], channels[1], 3, stride=1, padding=dilation, bias=False, 84 | dilation=dilation)) 85 | ] 86 | if dropout is not None: 87 | layers = layers[0:2] + [("dropout", dropout())] + layers[2:] 88 | else: 89 | layers = [ 90 | ("conv1", nn.Conv2d(in_channels, channels[0], 1, stride=1, padding=0, bias=False)), 91 | ("bn2", norm_act(channels[0])), 92 | ("conv2", nn.Conv2d(channels[0], channels[1], 3, stride=stride, padding=dilation, bias=False, 93 | groups=groups, dilation=dilation)), 94 | ("bn3", norm_act(channels[1])), 95 | ("conv3", nn.Conv2d(channels[1], channels[2], 1, stride=1, padding=0, bias=False)) 96 | ] 97 | if dropout is not None: 98 | layers = layers[0:4] + [("dropout", dropout())] + layers[4:] 99 | self.convs = nn.Sequential(OrderedDict(layers)) 100 | 101 | if need_proj_conv: 102 | self.proj_conv = nn.Conv2d(in_channels, channels[-1], 1, stride=stride, padding=0, bias=False) 103 | 104 | def forward(self, x): 105 | if hasattr(self, "proj_conv"): 106 | bn1 = self.bn1(x) 107 | shortcut = self.proj_conv(bn1) 108 | else: 109 | shortcut = x.clone() 110 | bn1 = self.bn1(x) 111 | 112 | out = self.convs(bn1) 113 | out.add_(shortcut) 114 | 115 | return out 116 | 117 | class WideResNet(nn.Module): 118 | def __init__(self, structure, norm_act=ABN, classes=1000, dilation=False): 119 | """Wider ResNet with pre-activation (identity mapping) blocks 120 | This variant uses down-sampling by max-pooling in the first two blocks and by strided convolution in the others. 121 | Parameters 122 | ---------- 123 | structure : list of int 124 | Number of residual blocks in each of the six modules of the network. 125 | norm_act : callable 126 | Function to create normalization / activation Module. 127 | classes : int 128 | If not `0` also include global average pooling and a fully-connected layer with `classes` outputs at the end 129 | of the network. 130 | dilation : bool 131 | If `True` apply dilation to the last three modules and change the down-sampling factor from 32 to 8. 132 | """ 133 | super(WideResNet, self).__init__() 134 | self.structure = structure 135 | self.dilation = dilation 136 | 137 | if len(structure) != 6: 138 | raise ValueError("Expected a structure with six values") 139 | 140 | # Initial layers 141 | self.mod1 = nn.Sequential(OrderedDict([ 142 | ("conv1", nn.Conv2d(3, 64, 3, stride=1, padding=1, bias=False)) 143 | ])) 144 | 145 | # Groups of residual blocks 146 | in_channels = 64 147 | channels = [(128, 128), (256, 256), (512, 512), (512, 1024), (512, 1024, 2048), (1024, 2048, 4096)] 148 | for mod_id, num in enumerate(structure): 149 | # Create blocks for module 150 | blocks = [] 151 | for block_id in range(num): 152 | if not dilation: 153 | dil = 1 154 | stride = 2 if block_id == 0 and 2 <= mod_id <= 4 else 1 155 | else: 156 | if mod_id == 3: 157 | dil = 2 158 | elif mod_id > 3: 159 | dil = 4 160 | else: 161 | dil = 1 162 | stride = 2 if block_id == 0 and mod_id == 2 else 1 163 | 164 | if mod_id == 4: 165 | drop = partial(nn.Dropout2d, p=0.3) 166 | elif mod_id == 5: 167 | drop = partial(nn.Dropout2d, p=0.5) 168 | else: 169 | drop = None 170 | 171 | blocks.append(( 172 | "block%d" % (block_id + 1), 173 | IdentityResidualBlock(in_channels, channels[mod_id], norm_act=norm_act, stride=stride, dilation=dil, 174 | dropout=drop) 175 | )) 176 | 177 | # Update channels and p_keep 178 | in_channels = channels[mod_id][-1] 179 | 180 | # Create module 181 | if mod_id < 2: 182 | self.add_module("pool%d" % (mod_id + 2), nn.MaxPool2d(3, stride=2, padding=1)) 183 | self.add_module("mod%d" % (mod_id + 2), nn.Sequential(OrderedDict(blocks))) 184 | 185 | # Pooling and predictor 186 | self.bn_out = norm_act(in_channels) 187 | if classes != 0: 188 | self.classifier = nn.Sequential(OrderedDict([ 189 | ("avg_pool", GlobalAvgPool2d()), 190 | ("fc", nn.Linear(in_channels, classes)) 191 | ])) 192 | 193 | def forward(self, img): 194 | out = self.mod1(img) 195 | out = self.pool2(out) 196 | out = self.mod2(out) 197 | 198 | out = self.pool3(out) 199 | out = self.mod3(out) 200 | out = self.mod4(out) 201 | out = self.mod5(out) 202 | out = self.mod6(out) 203 | out = self.mod7(out) 204 | out = self.bn_out(out) 205 | 206 | if hasattr(self, "classifier"): 207 | return self.classifier(out) 208 | else: 209 | return out 210 | return out 211 | 212 | def wideresnet38(pretrained=False, root='~/.gluoncvth/models', **kwargs): 213 | """Constructs a ResNet-50 model. 214 | 215 | Args: 216 | pretrained (bool): If True, returns a model pre-trained on ImageNet 217 | """ 218 | model = WideResNet([3, 3, 6, 3, 1, 1], **kwargs) 219 | if pretrained: 220 | model.load_state_dict(torch.load( 221 | get_model_file('wideresnet38', root=root)), strict=False) 222 | return model 223 | -------------------------------------------------------------------------------- /gluoncvth/utils/__init__.py: -------------------------------------------------------------------------------- 1 | from .files import * 2 | from .preset import * 3 | from .pallete import * 4 | -------------------------------------------------------------------------------- /gluoncvth/utils/files.py: -------------------------------------------------------------------------------- 1 | import os 2 | import requests 3 | import errno 4 | import shutil 5 | import hashlib 6 | from tqdm import tqdm 7 | import torch 8 | 9 | __all__ = ['save_checkpoint', 'download', 'mkdir', 'check_sha1'] 10 | 11 | def save_checkpoint(state, args, is_best, filename='checkpoint.pth.tar'): 12 | """Saves checkpoint to disk""" 13 | directory = "runs/%s/%s/%s/"%(args.dataset, args.model, args.checkname) 14 | if not os.path.exists(directory): 15 | os.makedirs(directory) 16 | filename = directory + filename 17 | torch.save(state, filename) 18 | if is_best: 19 | shutil.copyfile(filename, directory + 'model_best.pth.tar') 20 | 21 | 22 | def download(url, path=None, overwrite=False, sha1_hash=None): 23 | """Download an given URL 24 | Parameters 25 | ---------- 26 | url : str 27 | URL to download 28 | path : str, optional 29 | Destination path to store downloaded file. By default stores to the 30 | current directory with same name as in url. 31 | overwrite : bool, optional 32 | Whether to overwrite destination file if already exists. 33 | sha1_hash : str, optional 34 | Expected sha1 hash in hexadecimal digits. Will ignore existing file when hash is specified 35 | but doesn't match. 36 | Returns 37 | ------- 38 | str 39 | The file path of the downloaded file. 40 | """ 41 | if path is None: 42 | fname = url.split('/')[-1] 43 | else: 44 | path = os.path.expanduser(path) 45 | if os.path.isdir(path): 46 | fname = os.path.join(path, url.split('/')[-1]) 47 | else: 48 | fname = path 49 | 50 | if overwrite or not os.path.exists(fname) or (sha1_hash and not check_sha1(fname, sha1_hash)): 51 | dirname = os.path.dirname(os.path.abspath(os.path.expanduser(fname))) 52 | if not os.path.exists(dirname): 53 | os.makedirs(dirname) 54 | 55 | print('Downloading %s from %s...'%(fname, url)) 56 | r = requests.get(url, stream=True) 57 | if r.status_code != 200: 58 | raise RuntimeError("Failed downloading url %s"%url) 59 | total_length = r.headers.get('content-length') 60 | with open(fname, 'wb') as f: 61 | if total_length is None: # no content length header 62 | for chunk in r.iter_content(chunk_size=1024): 63 | if chunk: # filter out keep-alive new chunks 64 | f.write(chunk) 65 | else: 66 | total_length = int(total_length) 67 | for chunk in tqdm(r.iter_content(chunk_size=1024), 68 | total=int(total_length / 1024. + 0.5), 69 | unit='KB', unit_scale=False, dynamic_ncols=True): 70 | f.write(chunk) 71 | 72 | if sha1_hash and not check_sha1(fname, sha1_hash): 73 | raise UserWarning('File {} is downloaded but the content hash does not match. ' \ 74 | 'The repo may be outdated or download may be incomplete. ' \ 75 | 'If the "repo_url" is overridden, consider switching to ' \ 76 | 'the default repo.'.format(fname)) 77 | 78 | return fname 79 | 80 | 81 | def check_sha1(filename, sha1_hash): 82 | """Check whether the sha1 hash of the file content matches the expected hash. 83 | Parameters 84 | ---------- 85 | filename : str 86 | Path to the file. 87 | sha1_hash : str 88 | Expected sha1 hash in hexadecimal digits. 89 | Returns 90 | ------- 91 | bool 92 | Whether the file content matches the expected hash. 93 | """ 94 | sha1 = hashlib.sha1() 95 | with open(filename, 'rb') as f: 96 | while True: 97 | data = f.read(1048576) 98 | if not data: 99 | break 100 | sha1.update(data) 101 | 102 | return sha1.hexdigest() == sha1_hash 103 | 104 | 105 | def mkdir(path): 106 | """make dir exists okay""" 107 | try: 108 | os.makedirs(path) 109 | except OSError as exc: # Python >2.5 110 | if exc.errno == errno.EEXIST and os.path.isdir(path): 111 | pass 112 | else: 113 | raise 114 | -------------------------------------------------------------------------------- /gluoncvth/utils/pallete.py: -------------------------------------------------------------------------------- 1 | from PIL import Image 2 | 3 | def get_mask_pallete(npimg, dataset='detail'): 4 | """Get image color pallete for visualizing masks""" 5 | # recovery boundary 6 | if dataset == 'pascal_voc': 7 | npimg[npimg==21] = 255 8 | # put colormap 9 | out_img = Image.fromarray(npimg.squeeze().astype('uint8')) 10 | if dataset == 'ade20k': 11 | out_img.putpalette(adepallete) 12 | elif dataset == 'cityscapes': 13 | out_img.putpalette(citypallete) 14 | elif dataset in ('detail', 'pascal_voc', 'pascal_aug'): 15 | out_img.putpalette(vocpallete) 16 | return out_img 17 | 18 | def _get_voc_pallete(num_cls): 19 | n = num_cls 20 | pallete = [0]*(n*3) 21 | for j in range(0,n): 22 | lab = j 23 | pallete[j*3+0] = 0 24 | pallete[j*3+1] = 0 25 | pallete[j*3+2] = 0 26 | i = 0 27 | while (lab > 0): 28 | pallete[j*3+0] |= (((lab >> 0) & 1) << (7-i)) 29 | pallete[j*3+1] |= (((lab >> 1) & 1) << (7-i)) 30 | pallete[j*3+2] |= (((lab >> 2) & 1) << (7-i)) 31 | i = i + 1 32 | lab >>= 3 33 | return pallete 34 | 35 | vocpallete = _get_voc_pallete(256) 36 | 37 | adepallete = [0,0,0,120,120,120,180,120,120,6,230,230,80,50,50,4,200,3,120,120,80,140,140,140,204,5,255,230,230,230,4,250,7,224,5,255,235,255,7,150,5,61,120,120,70,8,255,51,255,6,82,143,255,140,204,255,4,255,51,7,204,70,3,0,102,200,61,230,250,255,6,51,11,102,255,255,7,71,255,9,224,9,7,230,220,220,220,255,9,92,112,9,255,8,255,214,7,255,224,255,184,6,10,255,71,255,41,10,7,255,255,224,255,8,102,8,255,255,61,6,255,194,7,255,122,8,0,255,20,255,8,41,255,5,153,6,51,255,235,12,255,160,150,20,0,163,255,140,140,140,250,10,15,20,255,0,31,255,0,255,31,0,255,224,0,153,255,0,0,0,255,255,71,0,0,235,255,0,173,255,31,0,255,11,200,200,255,82,0,0,255,245,0,61,255,0,255,112,0,255,133,255,0,0,255,163,0,255,102,0,194,255,0,0,143,255,51,255,0,0,82,255,0,255,41,0,255,173,10,0,255,173,255,0,0,255,153,255,92,0,255,0,255,255,0,245,255,0,102,255,173,0,255,0,20,255,184,184,0,31,255,0,255,61,0,71,255,255,0,204,0,255,194,0,255,82,0,10,255,0,112,255,51,0,255,0,194,255,0,122,255,0,255,163,255,153,0,0,255,10,255,112,0,143,255,0,82,0,255,163,255,0,255,235,0,8,184,170,133,0,255,0,255,92,184,0,255,255,0,31,0,184,255,0,214,255,255,0,112,92,255,0,0,224,255,112,224,255,70,184,160,163,0,255,153,0,255,71,255,0,255,0,163,255,204,0,255,0,143,0,255,235,133,255,0,255,0,235,245,0,255,255,0,122,255,245,0,10,190,212,214,255,0,0,204,255,20,0,255,255,255,0,0,153,255,0,41,255,0,255,204,41,0,255,41,255,0,173,0,255,0,245,255,71,0,255,122,0,255,0,255,184,0,92,255,184,255,0,0,133,255,255,214,0,25,194,194,102,255,0,92,0,255] 38 | 39 | citypallete = [ 40 | 128,64,128,244,35,232,70,70,70,102,102,156,190,153,153,153,153,153,250,170,30,220,220,0,107,142,35,152,251,152,70,130,180,220,20,60,255,0,0,0,0,142,0,0,70,0,60,100,0,80,100,0,0,230,119,11,32,128,192,0,0,64,128,128,64,128,0,192,128,128,192,128,64,64,0,192,64,0,64,192,0,192,192,0,64,64,128,192,64,128,64,192,128,192,192,128,0,0,64,128,0,64,0,128,64,128,128,64,0,0,192,128,0,192,0,128,192,128,128,192,64,0,64,192,0,64,64,128,64,192,128,64,64,0,192,192,0,192,64,128,192,192,128,192,0,64,64,128,64,64,0,192,64,128,192,64,0,64,192,128,64,192,0,192,192,128,192,192,64,64,64,192,64,64,64,192,64,192,192,64,64,64,192,192,64,192,64,192,192,192,192,192,32,0,0,160,0,0,32,128,0,160,128,0,32,0,128,160,0,128,32,128,128,160,128,128,96,0,0,224,0,0,96,128,0,224,128,0,96,0,128,224,0,128,96,128,128,224,128,128,32,64,0,160,64,0,32,192,0,160,192,0,32,64,128,160,64,128,32,192,128,160,192,128,96,64,0,224,64,0,96,192,0,224,192,0,96,64,128,224,64,128,96,192,128,224,192,128,32,0,64,160,0,64,32,128,64,160,128,64,32,0,192,160,0,192,32,128,192,160,128,192,96,0,64,224,0,64,96,128,64,224,128,64,96,0,192,224,0,192,96,128,192,224,128,192,32,64,64,160,64,64,32,192,64,160,192,64,32,64,192,160,64,192,32,192,192,160,192,192,96,64,64,224,64,64,96,192,64,224,192,64,96,64,192,224,64,192,96,192,192,224,192,192,0,32,0,128,32,0,0,160,0,128,160,0,0,32,128,128,32,128,0,160,128,128,160,128,64,32,0,192,32,0,64,160,0,192,160,0,64,32,128,192,32,128,64,160,128,192,160,128,0,96,0,128,96,0,0,224,0,128,224,0,0,96,128,128,96,128,0,224,128,128,224,128,64,96,0,192,96,0,64,224,0,192,224,0,64,96,128,192,96,128,64,224,128,192,224,128,0,32,64,128,32,64,0,160,64,128,160,64,0,32,192,128,32,192,0,160,192,128,160,192,64,32,64,192,32,64,64,160,64,192,160,64,64,32,192,192,32,192,64,160,192,192,160,192,0,96,64,128,96,64,0,224,64,128,224,64,0,96,192,128,96,192,0,224,192,128,224,192,64,96,64,192,96,64,64,224,64,192,224,64,64,96,192,192,96,192,64,224,192,192,224,192,32,32,0,160,32,0,32,160,0,160,160,0,32,32,128,160,32,128,32,160,128,160,160,128,96,32,0,224,32,0,96,160,0,224,160,0,96,32,128,224,32,128,96,160,128,224,160,128,32,96,0,160,96,0,32,224,0,160,224,0,32,96,128,160,96,128,32,224,128,160,224,128,96,96,0,224,96,0,96,224,0,224,224,0,96,96,128,224,96,128,96,224,128,224,224,128,32,32,64,160,32,64,32,160,64,160,160,64,32,32,192,160,32,192,32,160,192,160,160,192,96,32,64,224,32,64,96,160,64,224,160,64,96,32,192,224,32,192,96,160,192,224,160,192,32,96,64,160,96,64,32,224,64,160,224,64,32,96,192,160,96,192,32,224,192,160,224,192,96,96,64,224,96,64,96,224,64,224,224,64,96,96,192,224,96,192,96,224,192,0,0,0] 41 | 42 | -------------------------------------------------------------------------------- /gluoncvth/utils/preset.py: -------------------------------------------------------------------------------- 1 | """Preset Transforms for Demos""" 2 | from PIL import Image 3 | import numpy as np 4 | import torch 5 | import torchvision.transforms as transform 6 | 7 | __all__ = ['load_image'] 8 | 9 | input_transform = transform.Compose([ 10 | transform.ToTensor(), 11 | transform.Normalize([.485, .456, .406], [.229, .224, .225])]) 12 | 13 | def load_image(filename, size=None, scale=None, keep_asp=True): 14 | """Load the image for demos""" 15 | img = Image.open(filename).convert('RGB') 16 | if size is not None: 17 | if keep_asp: 18 | size2 = int(size * 1.0 / img.size[0] * img.size[1]) 19 | img = img.resize((size, size2), Image.ANTIALIAS) 20 | else: 21 | img = img.resize((size, size), Image.ANTIALIAS) 22 | elif scale is not None: 23 | img = img.resize((int(img.size[0] / scale), int(img.size[1] / scale)), Image.ANTIALIAS) 24 | 25 | img = input_transform(img) 26 | return img 27 | -------------------------------------------------------------------------------- /image/demo_deeplab_ade.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/StacyYang/gluoncv-torch/8cab8da1ad546a39ee4821fbe2b01340958633a6/image/demo_deeplab_ade.png -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | import os 3 | import io 4 | import re 5 | import subprocess 6 | from setuptools import setup, find_packages 7 | import setuptools.command.develop 8 | import setuptools.command.install 9 | 10 | version = '0.0.6' 11 | 12 | try: 13 | if not os.getenv('RELEASE'): 14 | from datetime import date 15 | today = date.today() 16 | day = today.strftime("b%Y%m%d") 17 | version += day 18 | except Exception: 19 | pass 20 | 21 | def create_version_file(): 22 | global version, cwd 23 | print('-- Building version ' + version) 24 | version_path = os.path.join(cwd, 'gluoncvth', 'version.py') 25 | with open(version_path, 'w') as f: 26 | f.write('"""This is gluoncvth version file."""\n') 27 | f.write("__version__ = '{}'\n".format(version)) 28 | 29 | # run test scrip after installation 30 | class install(setuptools.command.install.install): 31 | def run(self): 32 | create_version_file() 33 | setuptools.command.install.install.run(self) 34 | 35 | class develop(setuptools.command.develop.develop): 36 | def run(self): 37 | create_version_file() 38 | setuptools.command.develop.develop.run(self) 39 | 40 | long_description = open('README.md').read() 41 | 42 | requirements = [ 43 | 'numpy', 44 | 'torch', 45 | 'tqdm', 46 | 'requests', 47 | 'Pillow', 48 | ] 49 | 50 | setup( 51 | # Metadata 52 | name='gluoncv-torch', 53 | version=version, 54 | author='Gluon CV Toolkit Contributors', 55 | url='https://github.com/dmlc/gluon-cv', 56 | description='MXNet Gluon CV Toolkit', 57 | long_description=long_description, 58 | long_description_content_type='text/markdown', 59 | license='Apache-2.0', 60 | # Package info 61 | packages=find_packages(exclude=('docs', 'tests', 'scripts')), 62 | zip_safe=True, 63 | include_package_data=True, 64 | install_requires=requirements, 65 | ) 66 | --------------------------------------------------------------------------------