SemanticGAN
This is the official code for:
Semantic Segmentation with Generative Models: Semi-Supervised Learning and Strong Out-of-Domain Generalization
Daiqing Li, Junlin Yang, Karsten Kreis, Antonio Torralba, Sanja Fidler
Requirements
- Python 3.6 or 3.7 are supported.
- Pytorch 1.4.0 + is recommended.
- This code is tested with CUDA 10.2 toolkit and CuDNN 7.5.
- Please check the python package requirement from
requirements.txt
, and install using
pip install -r requirements.txt
Training
To reproduce paper Semantic Segmentation with Generative Models: Semi-Supervised Learning and Strong Out-of-Domain Generalization:
- Run Step1: Semantic GAN training
- Run Step2: Encoder training
- Run Inference & Optimization.
0. Prepare for FID calculation
In order to calculate FID score, you need to prepare inception features for your dataset,
python prepare_inception.py \
--size [resolution of the image] \
--batch [batch size] \
--output [path to save the inception file, in .pkl] \
--dataset_name celeba-mask \
[positional argument 1, path to the image folder]] \
1. GAN Training
For training GAN with both image and its label,
python train_seg_gan.py \
--img_dataset [path-to-img-folder] \
--seg_dataset [path-to-seg-folder] \
--inception [path-to-inception file] \
--seg_name celeba-mask \
--checkpoint_dir [path-to-ckpt-dir] \
To use multi-gpus training in the cloud,
python -m torch.distributed.launch \
--nproc_per_node=N_GPU \
--master_port=PORTtrain_gan.py \
train_gan.py \
--img_dataset [path-to-img-folder] \
--inception [path-to-inception file] \
--dataset_name celeba-mask \
--checkpoint_dir [path-to-ckpt-dir] \
2. Encoder Triaining
python train_enc.py \
--img_dataset [path-to-img-folder] \
--seg_dataset [path-to-seg-folder] \
--ckpt [path-to-pretrained GAN model] \
--seg_name celeba-mask \
--enc_backboend [fpn|res] \
--checkpoint_dir [path-to-ckpt-dir] \
Inference
For Face Parts Segmentation Task
python inference.py \
--ckpt [path-to-ckpt] \
--img_dir [path-to-test-folder] \
--outdir [path-to-output-folder] \
--dataset_name celeba-mask \
--w_plus \
--image_mode RGB \
--seg_dim 8 \
--step 200 [optimization steps] \
Visualization of different optimization steps