pytorch-semseg

Semantic Segmentation Algorithms Implemented in PyTorch

This repository aims at mirroring popular semantic segmentation architectures in PyTorch.

Networks implemented

• PSPNet - With support for loading pretrained models w/o caffe dependency
• ICNet - With optional batchnorm and pretrained models
• FRRN - Model A and B
• FCN - All 1 (FCN32s), 2 (FCN16s) and 3 (FCN8s) stream variants
• U-Net - With optional deconvolution and batchnorm
• Link-Net - With multiple resnet backends
• Segnet - With Unpooling using Maxpool indices

Upcoming

• E-Net
• RefineNet

DataLoaders implemented

• CamVid
• Pascal VOC
• ADE20K
• MIT Scene Parsing Benchmark
• Cityscapes
• NYUDv2
• Sun-RGBD

Requirements

• pytorch >=0.3.0
• torchvision ==0.2.0
• visdom >=1.0.1 (for loss and results visualization)
• scipy
• tqdm

One-line installation

pip install -r requirements.txt

Data

• Download data for desired dataset(s) from list of URLs here.
• Extract the zip / tar and modify the path appropriately in config.json

Usage

Launch visdom by running (in a separate terminal window)

python -m visdom.server

To train the model :

python train.py [-h] [--arch [ARCH]] [--dataset [DATASET]]
                [--img_rows [IMG_ROWS]] [--img_cols [IMG_COLS]]
                [--n_epoch [N_EPOCH]] [--batch_size [BATCH_SIZE]]
                [--l_rate [L_RATE]] [--feature_scale [FEATURE_SCALE]]
                [--visdom [VISDOM]]

  --arch           Architecture to use ['fcn8s, unet, segnet etc']
  --dataset        Dataset to use ['pascal, camvid, ade20k etc']
  --img_rows       Height of the input image
  --img_cols       Width of the input image
  --n_epoch        # of the epochs
  --batch_size     Batch Size
  --l_rate         Learning Rate
  --feature_scale  Divider for # of features to use
  --visdom         Show visualization(s) on visdom | False by default

To validate the model :

python validate.py [-h] [--model_path [MODEL_PATH]] [--dataset [DATASET]]
                   [--img_rows [IMG_ROWS]] [--img_cols [IMG_COLS]]
                   [--batch_size [BATCH_SIZE]] [--split [SPLIT]]

  --model_path   Path to the saved model
  --dataset      Dataset to use ['pascal, camvid, ade20k etc']
  --img_rows     Height of the input image
  --img_cols     Width of the input image
  --batch_size   Batch Size
  --split        Split of dataset to validate on

To test the model w.r.t. a dataset on custom images(s):

python test.py [-h] [--model_path [MODEL_PATH]] [--dataset [DATASET]]
               [--dcrf [DCRF]] [--img_path [IMG_PATH]] [--out_path [OUT_PATH]]
 
  --model_path          Path to the saved model
  --dataset             Dataset to use ['pascal, camvid, ade20k etc']
  --dcrf                Enable DenseCRF based post-processing
  --img_path            Path of the input image
  --out_path            Path of the output segmap

If you find this code useful in your research, please consider citing:

@article{mshahsemseg,
    Author = {Meet P Shah},
    Title = {Semantic Segmentation Architectures Implemented in PyTorch.},
    Journal = {https://github.com/meetshah1995/pytorch-semseg},
    Year = {2017}
}