<div align="center">Semantic Segmentation</div>

<div align="center"> <p>Easy to use and customizable SOTA Semantic Segmentation models with abundant datasets in PyTorch</p>

<a href="https://colab.research.google.com/github/sithu31296/semantic-segmentation/blob/main/notebooks/tutorial.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a>

</div>

<div align="center">Major Rework! Stay tuned...</div>

<p>A lot has been changed since 2022, nowadays there are even open-world segmentation models (Segment Anything). However, traditional segmentation models are still in demand for high accuracy and custom use cases. This repo will be updated according to new PyTorch version, updated models, and documentations for how to use with a custom dataset, etc.</p> <p>Expected release date -> May 2024</p>

Planned features:

• Rework on whole training pipeline
• Baseline pre-trained model
• New updated ideas
• Easy integration with SOTA backbone models (with tutorials)
• Tutorial for custom dataset
• Distributed training

Current features to be discarded:

• Amount of datasets provided will be reduced. But instead, representative ones will be remained with a tutorial for custom dataset.
• Amount of models provided will be reducted. Instead, valuable tricks and modules will be remained, and can be easily integrated with any models.
• Augmentations will be replaced with official torchvisionv2 transforms.
• Conversion and inference with other frameworks

</div>

<div align="center">Features</div>

• Applicable to following tasks:
  • Scene Parsing
  • Human Parsing
  • Face Parsing
  • Medical Image Segmentation (Coming Soon)
• 20+ Datasets
• 15+ SOTA Backbones
• 10+ SOTA Semantic Segmentation Models
• PyTorch, ONNX, TFLite, OpenVINO Export & Inference

</div>

<div align="center">Model Zoo</div>

Supported Backbones:

• ResNet (CVPR 2016)
• ResNetD (ArXiv 2018)
• MobileNetV2 (CVPR 2018)
• MobileNetV3 (ICCV 2019)
• MiT (NeurIPS 2021)
• ResT (NeurIPS 2021)
• MicroNet (ICCV 2021)
• ResNet+ (ArXiv 2021)
• PVTv2 (CVMJ 2022)
• PoolFormer (CVPR 2022)
• ConvNeXt (CVPR 2022)
• UniFormer (ArXiv 2022)
• VAN (ArXiv 2022)
• DaViT (ArXiv 2022)

Supported Heads/Methods:

• FCN (CVPR 2015)
• UPerNet (ECCV 2018)
• BiSeNetv1 (ECCV 2018)
• FPN (CVPR 2019)
• SFNet (ECCV 2020)
• SegFormer (NeurIPS 2021)
• FaPN (ICCV 2021)
• CondNet (IEEE SPL 2021)
• Light-Ham (ICLR 2021)
• Lawin (ArXiv 2022)
• TopFormer (CVPR 2022)

Supported Standalone Models:

• BiSeNetv2 (IJCV 2021)
• DDRNet (ArXiv 2021)

Supported Modules:

• PPM (CVPR 2017)
• PSA (ArXiv 2021)

Refer to MODELS for benchmarks and available pre-trained models.

And check BACKBONES for supported backbones.

Notes: Most of the methods do not have pre-trained models. It's very difficult to combine different models with pre-trained weights in one repository and limited resource to re-train myself.

<div align="center">Supported Datasets</div>

Scene Parsing:

• ADE20K
• CityScapes
• COCO-Stuff
• CamVid
• PASCAL-Context
• Mapillary Vistas
• Sun RGB-D

Human Parsing:

• MHPv2
• MHPv1
• LIP
• CCIHP
• CIHP
• ATR

Face Parsing:

• HELEN
• LaPa
• iBugMask
• CelebAMaskHQ
• FaceSynthetics

Others:

• SUIM

Refer to DATASETS for more details and dataset preparation.

<details> <summary><strong>Available Augmentations</strong> (click to expand)</summary>

Check the notebook here to test the augmentation effects.

Pixel-level Transforms:

• ColorJitter (Brightness, Contrast, Saturation, Hue)
• Gamma, Sharpness, AutoContrast, Equalize, Posterize
• GaussianBlur, Grayscale

Spatial-level Transforms:

• Affine, RandomRotation
• HorizontalFlip, VerticalFlip
• CenterCrop, RandomCrop
• Pad, ResizePad, Resize
• RandomResizedCrop

</details>

<div align="center">Usage</div>

<details open> <summary><strong>Installation</strong></summary>

• python >= 3.6
• torch >= 1.8.1
• torchvision >= 0.9.1

Then, clone the repo and install the project with:

$ git clone https://github.com/sithu31296/semantic-segmentation
$ cd semantic-segmentation
$ pip install -e .

</details> <br> <details> <summary><strong>Configuration</strong> (click to expand)</summary>

Create a configuration file in configs. Sample configuration for ADE20K dataset can be found here. Then edit the fields you think if it is needed. This configuration file is needed for all of training, evaluation and prediction scripts.

</details> <br> <details> <summary><strong>Training</strong> (click to expand)</summary>

To train with a single GPU:

$ python tools/train.py --cfg configs/CONFIG_FILE.yaml

To train with multiple gpus, set DDP field in config file to true and run as follows:

$ python -m torch.distributed.launch --nproc_per_node=2 --use_env tools/train.py --cfg configs/<CONFIG_FILE_NAME>.yaml

</details> <br> <details> <summary><strong>Evaluation</strong> (click to expand)</summary>

Make sure to set MODEL_PATH of the configuration file to your trained model directory.

$ python tools/val.py --cfg configs/<CONFIG_FILE_NAME>.yaml

To evaluate with multi-scale and flip, change ENABLE field in MSF to true and run the same command as above.

</details> <br> <details open> <summary><strong>Inference</strong></summary>

To make an inference, edit the parameters of the config file from below.

• Change MODEL >> NAME and BACKBONE to your desired pretrained model.
• Change DATASET >> NAME to the dataset name depending on the pretrained model.
• Set TEST >> MODEL_PATH to pretrained weights of the testing model.
• Change TEST >> FILE to the file or image folder path you want to test.
• Testing results will be saved in SAVE_DIR.

## example using ade20k pretrained models
$ python tools/infer.py --cfg configs/ade20k.yaml

Example test results (SegFormer-B2):

</details> <br> <details> <summary><strong>Convert to other Frameworks</strong> (ONNX, CoreML, OpenVINO, TFLite)</summary>

To convert to ONNX and CoreML, run:

$ python tools/export.py --cfg configs/<CONFIG_FILE_NAME>.yaml

To convert to OpenVINO and TFLite, see torch_optimize.

</details> <br> <details> <summary><strong>Inference</strong> (ONNX, OpenVINO, TFLite)</summary>

## ONNX Inference
$ python scripts/onnx_infer.py --model <ONNX_MODEL_PATH> --img-path <TEST_IMAGE_PATH>

## OpenVINO Inference
$ python scripts/openvino_infer.py --model <OpenVINO_MODEL_PATH> --img-path <TEST_IMAGE_PATH>

## TFLite Inference
$ python scripts/tflite_infer.py --model <TFLite_MODEL_PATH> --img-path <TEST_IMAGE_PATH>

</details> <br> <details> <summary><strong>References</strong> (click to expand)</summary>

• https://github.com/CoinCheung/BiSeNet
• https://github.com/open-mmlab/mmsegmentation
• https://github.com/rwightman/pytorch-image-models

</details> <br> <details> <summary><strong>Citations</strong> (click to expand)</summary>

@article{xie2021segformer,
  title={SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers},
  author={Xie, Enze and Wang, Wenhai and Yu, Zhiding and Anandkumar, Anima and Alvarez, Jose M and Luo, Ping},
  journal={arXiv preprint arXiv:2105.15203},
  year={2021}
}

@misc{xiao2018unified,
  title={Unified Perceptual Parsing for Scene Understanding}, 
  author={Tete Xiao and Yingcheng Liu and Bolei Zhou and Yuning Jiang and Jian Sun},
  year={2018},
  eprint={1807.10221},
  archivePrefix={arXiv},
  primaryClass={cs.CV}
}

@article{hong2021deep,
  title={Deep Dual-resolution Networks for Real-time and Accurate Semantic Segmentation of Road Scenes},
  author={Hong, Yuanduo and Pan, Huihui and Sun, Weichao and Jia, Yisong},
  journal={arXiv preprint arXiv:2101.06085},
  year={2021}
}

@misc{zhang2021rest,
  title={ResT: An Efficient Transformer for Visual Recognition}, 
  author