image_seg

This project is an open source semantic segmentation toolbox based on PyTorch. It is based on the codes of our Tianchi competition in 2021 .
In the competition, our team won the third place (please see Tianchi_README.md).

Overview

The master branch works with PyTorch 1.6+.The project now supports popular and contemporary semantic segmentation frameworks, e.g. UNet, DeepLabV3+, HR-Net etc.

Requirements

Support

Backbone

  • ResNet (CVPR'2016)
  • SeNet (CVPR'2018)
  • IBN-Net (CVPR'2018)
  • EfficientNet (CVPR'2020)

Methods

  • UNet
  • DLink-Net
  • Res-UNet
  • Efficient-UNet
  • Deeplab v3+
  • HR-Net

Tricks

  • MixUp /CutMix /CopyPaste
  • SWA
  • LovaszSoftmax Loss /LargeMarginSoftmax Loss
  • FP16
  • Multi-scale

Tools

  • large image inference (cut and merge)
  • post process (crf/superpixels)

Quick Start

Train a model

python train.py --config_file ${CONFIG_FILE}
  • CONFIG_FILE: File of training config about model

Examples:
We trained our model in Tianchi competition according to the following script:
Stage 1 (160e)

python train.py --config_file configs/tc_seg/tc_seg_res_unet_r34_ibn_a_160e.yml

Stage 2 (swa 24e)

python train.py --config_file configs/tc_seg/tc_seg_res_unet_r34_ibn_a_swa.yml

Inference with pretrained models

python inference.py --config_file ${CONFIG_FILE}
  • CONFIG_FILE: File of inference config about model

Predict large image with pretrained models

python predict_demo.py --config_file ${CONFIG_FILE} --rs_img_file ${IMAGE_FILE_PATH} --temp_img_save_path ${TEMP_CUT_PATH} -temp_seg_map_save_path ${TEMP_SAVE_PATH} --save_seg_map_file ${SAVE_SEG_FILE}
  • CONFIG_FILE: File of inference config about model
  • IMAGE_FILE_PATH: File of large input image to predict
  • TEMP_CUT_PATH: Temp folder of small cutting samples
  • TEMP_SAVE_PATH: Temp folder of predict results of cutting samples
  • SAVE_SEG_FILE: Predict result of the large image

GitHub

https://github.com/whut2962575697/image_seg