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:

  1. Run Step1: Semantic GAN training
  2. Run Step2: Encoder training
  3. 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

face-parts-seg

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

face-parts-opt-steps

GitHub

https://github.com/nv-tlabs/semanticGAN_code