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.4.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] [--config [CONFIG]] 

--config                Configuration file to use

To validate the model :

usage: validate.py [-h] [--config [CONFIG]] [--model_path [MODEL_PATH]]
                       [--eval_flip] [--measure_time]

  --config              Config file to be used
  --model_path          Path to the saved model
  --eval_flip           Enable evaluation with flipped image | True by default
  --measure_time        Enable evaluation with time (fps) measurement | True
                        by default

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}
}