This project is an open source semantic segmentation toolbox based on PyTorch. It is based on the codes of our Tianchi competition in 2021 (https://tianchi.aliyun.com/competition/entrance/531860/introduction).
In the competition, our team won the third place (please see Tianchi_README.md).
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.
- ResNet (CVPR'2016)
- SeNet (CVPR'2018)
- IBN-Net (CVPR'2018)
- EfficientNet (CVPR'2020)
- UNet
- DLink-Net
- Res-UNet
- Efficient-UNet
- Deeplab v3+
- HR-Net
- MixUp /CutMix /CopyPaste
- SWA
- LovaszSoftmax Loss /LargeMarginSoftmax Loss
- FP16
- Multi-scale
- large image inference (cut and merge)
- post process (crf/superpixels)
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
python inference.py --config_file ${CONFIG_FILE}
CONFIG_FILE
: File of inference config about model
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 modelIMAGE_FILE_PATH
: File of large input image to predictTEMP_CUT_PATH
: Temp folder of small cutting samplesTEMP_SAVE_PATH
: Temp folder of predict results of cutting samplesSAVE_SEG_FILE
: Predict result of the large image