Refactor for pip installation

.gitignore

@@ -4,3 +4,5 @@ ckpt/
 vis/
 log/
 pretrained/
+
+.idea/

README.md

@@ -228,6 +228,19 @@ python3 eval_multipro.py --gpus GPUS --cfg config/ade20k-resnet50dilated-ppm_dee
 python3 eval_multipro.py --gpus GPUS --cfg config/ade20k-resnet101-upernet.yaml
 ```
 
+## Integration with other projects
+This library can be installed via `pip` to easily integrate with another codebase
+```bash
+pip install git+https://github.com/CSAILVision/semantic-segmentation-pytorch.git@master
+```
+
+Now this library can easily be consumed programmatically. For example
+```python
+from mit_semseg.config import cfg
+from mit_semseg.dataset import TestDataset
+from mit_semseg.models import ModelBuilder, SegmentationModule
+```
+
 ## Reference
 
 If you find the code or pre-trained models useful, please cite the following papers:

eval.py

@@ -9,12 +9,12 @@
 import torch.nn as nn
 from scipy.io import loadmat
 # Our libs
-from config import cfg
-from dataset import ValDataset
-from models import ModelBuilder, SegmentationModule
-from utils import AverageMeter, colorEncode, accuracy, intersectionAndUnion, setup_logger
-from lib.nn import user_scattered_collate, async_copy_to
-from lib.utils import as_numpy
+from mit_semseg.config import cfg
+from mit_semseg.dataset import ValDataset
+from mit_semseg.models import ModelBuilder, SegmentationModule
+from mit_semseg.utils import AverageMeter, colorEncode, accuracy, intersectionAndUnion, setup_logger
+from mit_semseg.lib.nn import user_scattered_collate, async_copy_to
+from mit_semseg.lib.utils import as_numpy
 from PIL import Image
 from tqdm import tqdm
 

eval_multipro.py

@@ -10,12 +10,12 @@
 import torch.nn as nn
 from scipy.io import loadmat
 # Our libs
-from config import cfg
-from dataset import ValDataset
-from models import ModelBuilder, SegmentationModule
-from utils import AverageMeter, colorEncode, accuracy, intersectionAndUnion, parse_devices, setup_logger
-from lib.nn import user_scattered_collate, async_copy_to
-from lib.utils import as_numpy
+from mit_semseg.config import cfg
+from mit_semseg.dataset import ValDataset
+from mit_semseg.models import ModelBuilder, SegmentationModule
+from mit_semseg.utils import AverageMeter, colorEncode, accuracy, intersectionAndUnion, parse_devices, setup_logger
+from mit_semseg.lib.nn import user_scattered_collate, async_copy_to
+from mit_semseg.lib.utils import as_numpy
 from PIL import Image
 from tqdm import tqdm
 

mit_semseg/__init__.py

@@ -0,0 +1,5 @@
+"""
+MIT CSAIL Semantic Segmentation
+"""
+
+__version__ = '1.0.0'

models/hrnet.py → mit_semseg/models/hrnet.py

@@ -8,7 +8,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from .utils import load_url
-from lib.nn import SynchronizedBatchNorm2d
+from mit_semseg.lib.nn import SynchronizedBatchNorm2d
 
 BatchNorm2d = SynchronizedBatchNorm2d
 BN_MOMENTUM = 0.1

models/mobilenet.py → mit_semseg/models/mobilenet.py

@@ -6,7 +6,7 @@
 import torch.nn as nn
 import math
 from .utils import load_url
-from lib.nn import SynchronizedBatchNorm2d
+from mit_semseg.lib.nn import SynchronizedBatchNorm2d
 
 BatchNorm2d = SynchronizedBatchNorm2d
 

models/models.py → mit_semseg/models/models.py

@@ -1,8 +1,7 @@
 import torch
 import torch.nn as nn
-import torchvision
 from . import resnet, resnext, mobilenet, hrnet
-from lib.nn import SynchronizedBatchNorm2d
+from mit_semseg.lib.nn import SynchronizedBatchNorm2d
 BatchNorm2d = SynchronizedBatchNorm2d
 
 

models/resnet.py → mit_semseg/models/resnet.py

@@ -1,7 +1,7 @@
 import torch.nn as nn
 import math
 from .utils import load_url
-from lib.nn import SynchronizedBatchNorm2d
+from mit_semseg.lib.nn import SynchronizedBatchNorm2d
 BatchNorm2d = SynchronizedBatchNorm2d
 
 

models/resnext.py → mit_semseg/models/resnext.py

@@ -1,7 +1,7 @@
 import torch.nn as nn
 import math
 from .utils import load_url
-from lib.nn import SynchronizedBatchNorm2d
+from mit_semseg.lib.nn import SynchronizedBatchNorm2d
 BatchNorm2d = SynchronizedBatchNorm2d
 
 

setup.py

@@ -0,0 +1,29 @@
+import setuptools
+
+with open('README.md', 'r') as fh:
+    long_description = fh.read()
+
+setuptools.setup(
+    name='mit_semseg',
+    version='1.0.0',
+    author='MIT CSAIL',
+    description='Pytorch implementation for Semantic Segmentation/Scene Parsing on MIT ADE20K dataset',
+    long_description=long_description,
+    long_description_content_type='text/markdown',
+    url='https://github.com/CSAILVision/semantic-segmentation-pytorch',
+    packages=setuptools.find_packages(),
+    classifiers=(
+        'Programming Language :: Python :: 3',
+        'License :: OSI Approved :: BSD License',
+        'Operating System :: OS Independent',
+    ),
+    install_requires=[
+        'numpy',
+        'torch>=0.4.1',
+        'torchvision',
+        'opencv-python',
+        'yacs',
+        'scipy',
+        'tqdm'
+    ]
+)

test.py

@@ -9,14 +9,14 @@
 from scipy.io import loadmat
 import csv
 # Our libs
-from dataset import TestDataset
-from models import ModelBuilder, SegmentationModule
-from utils import colorEncode, find_recursive, setup_logger
-from lib.nn import user_scattered_collate, async_copy_to
-from lib.utils import as_numpy
+from mit_semseg.dataset import TestDataset
+from mit_semseg.models import ModelBuilder, SegmentationModule
+from mit_semseg.utils import colorEncode, find_recursive, setup_logger
+from mit_semseg.lib.nn import user_scattered_collate, async_copy_to
+from mit_semseg.lib.utils import as_numpy
 from PIL import Image
 from tqdm import tqdm
-from config import cfg
+from mit_semseg.config import cfg
 
 colors = loadmat('data/color150.mat')['colors']
 names = {}

train.py

@@ -9,11 +9,11 @@
 import torch
 import torch.nn as nn
 # Our libs
-from config import cfg
-from dataset import TrainDataset
-from models import ModelBuilder, SegmentationModule
-from utils import AverageMeter, parse_devices, setup_logger
-from lib.nn import UserScatteredDataParallel, user_scattered_collate, patch_replication_callback
+from mit_semseg.config import cfg
+from mit_semseg.dataset import TrainDataset
+from mit_semseg.models import ModelBuilder, SegmentationModule
+from mit_semseg.utils import AverageMeter, parse_devices, setup_logger
+from mit_semseg.lib.nn import UserScatteredDataParallel, user_scattered_collate, patch_replication_callback
 
 
 # train one epoch