/pytorch-semseg

Semantic Segmentation Architectures Implemented in PyTorch

Primary LanguagePythonMIT LicenseMIT

pytorch-semseg

license pypi DOI

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

DataLoaders implemented

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