Awesome
Semantic Segmentation with Prototypes
This repository is a fork of https://github.com/cfchen-duke/ProtoPNet, adapted for ProtoSeg.
We also use pytorch implemention of DeepLab from: https://github.com/kazuto1011/deeplab-pytorch
Setting environment
Install conda environment:
conda install env -f env.yml
Before any script run:
source env.sh
Setting up data
Downloading data
Please use the following links to download Cityscapes or PASCAL VOC 2012:
- Cityscapes: https://www.cityscapes-dataset.com/downloads/ (gtFine_trainvaltest.zip and leftImgbit_trainvaltest.zip are needed)
- Pascal VOC 2012: Folllow instruction here: https://github.com/kazuto1011/deeplab-pytorch/blob/master/data/datasets/voc12/README.md. Please download the augmentated variant ("SegmentationAug" and "SegmentationClassAug").
Preprocessing data
We convert images to .npy files for faster loading during training. To do this, run the following commands:
# replace <N> with the number of parallel jobs
python -m segmentation.preprocess_cityscapes <N> # for cityscapes
python -m segmentation.preprocess_pascal <N> # for cityscapes
python -m segmentation.img_to_numpy
Training model
warmup + joint training + prototype push
# train on Cityscapes, no similarity loss, ImageNet pretraining
python -m segmentation.train cityscapes_no_kld_imnet <your_training_run_name>
# train on Cityscapes, with similarity loss, COCO pretraining
python -m segmentation.train cityscapes_kld_coco <your_training_run_name>
# train on Cityscapes, with similarity loss, ImageNet pretraining
python -m segmentation.train cityscapes_kld_imnet <your_training_run_name>
# train on PASCAL VOC 2012, no similarity loss, ImageNet pretraining
python -m segmentation.train pascal_no_kld_imnet <your_training_run_name>
# train on PASCAL VOC 2012, with similarity loss, COCO pretraining
python -m segmentation.train pascal_kld_coco <your_training_run_name>
# train on PASCAL VOC 2012, with similarity loss, ImageNet pretraining
python -m segmentation.train pascal_kld_imnet <your_training_run_name>
pruning and finetuning after pruning
python -m segmentation.run_pruning pascal_kld_imnet <your_training_run_name>
python -m segmentation.train pascal_kld_imnet <your_training_run_name> --pruned
Evaluation
# Evaluation on valid set for Cityscapes ('cityscapes_kld_imnet' should be replaced with your training config)
# The evaluation saves mIOU results in model directory. It also generates a few additional plots.
# <training_staged> should be one of: 'warmup', 'nopush', 'push', 'pruned'
python -m segmentation.eval_valid <your_training_run_name> <training_stage>
# For evaluating on Pascal dataset, add '-p' flag
python -m segmentation.eval_valid <your_training_run_name> <training_stage> -p
# Generating predictions on cityscapes or pascal test set:
python -m segmentation.eval_test <your_cityscapes_training_run_name> <training_stage>
python -m segmentation.eval_test <your_pascal_training_run_name> <training_stage> -p