KbeautyHair / KbeautyBaseline

Backbone model: StarGAN v2

Paper
StarGAN v2: Diverse Image Synthesis for Multiple Domains link
Yunjey Choi*, Youngjung Uh*, Jaejun Yoo*, Jung-Woo Ha
In CVPR 2020. (* indicates equal contribution)

Official implementation in Pytorch
The official Pytorch implementation of StarGAN v2 can be found at clovaai/stargan-v2

TensorFlow implementation
The TensorFlow implementation of StarGAN v2 by our team member junho can be found at clovaai/stargan-v2-tensorflow.

Software installation

Clone this repository:

git clone https://github.com/KbeautyHair/KbeautyBaseline.git
cd KbeautyBaseline/

Install the dependencies:

conda create -n KbeautyBaseline python=3.6.7
conda activate KbeautyBaseline

• CUDA version 10.0

conda install -y pytorch=1.6.0 torchvision=0.7.0 cudatoolkit=10.0 -c pytorch

• CUDA version 11.0

conda install -y pytorch=1.6.0 torchvision=0.7.0 cudatoolkit=11.0 -c pytorch

conda install x264=='1!152.20180717' ffmpeg=4.0.2 -c conda-forge
pip install opencv-python==4.1.2.30 ffmpeg-python==0.2.0 scikit-image==0.16.2
pip install pillow==7.0.0 scipy==1.2.1 tqdm==4.43.0 munch==2.5.0

Datasets and pre-trained networks

We provide links to download K-hairstyle dataset we used to train the baseline model and the corresponding pre-trained networks. The datasets and checkpoints of the pre-trained networks are required to be downloaded and stored in the data and expr/checkpoints directories, respectively.

To download the K-hairstyle dataset, please visit this dataset link and for the pre-trained networks, visit this checkpoint link, please.

Translating images

After downloading the pre-trained networks, you can synthesize output images reflecting diverse hairstyles of reference images. The following commands will save generated images to the expr/results directory. To generate images, run the following command:

python main.py --mode sample --img_size 512 --num_domains 2 --resume_iter 60000 --w_hpf 0 \
               --checkpoint_dir expr/checkpoints/k-hairstyle --result_dir expr/results/k-hairstyle \
               --trg_domain [TARGET DOMAIN (e.g., 0)] --src_dir sample_images/src --ref_dir sample_images/ref               

Evaluation metrics

To evaluate the baseline model using Fréchet Inception Distance (FID), run the following commands:

python main.py --mode eval --img_size 512 --num_domains 2 --w_hpf 0 \
               --resume_iter 60000 --num_sample 1000 --val_batch_size 50 \        
               --train_img_dir data/mqset \
               --val_img_dir data/mqset \
               --checkpoint_dir expr/checkpoints/k-hairstyle \
               --eval_dir expr/eval/k-hairstyle --dataset_dir imagelists

Training networks

To train the baseline model from scratch, run the following commands. Generated images and network checkpoints will be stored in the expr/samples and expr/checkpoints directories, respectively. Training takes about three days on a single Tesla V100 GPU. Please see here for training arguments and a description of them.

python main.py --mode train --img_size 512 --num_domains 2 --w_hpf 0 \
               --lambda_reg 1 --lambda_sty 2 --lambda_ds 1 --lambda_cyc 2 \
               --batch_size 5 --val_batch_size 30 \
               --train_img_dir data/mqset --val_img_dir data/mqset --dataset_dir imagelists \
               --checkpoint_dir expr/checkpoints/k-hairstyle --eval_dir expr/eval/k-hairstyle --sample_dir expr/samples/k-hairstyle

Reference

If you want to get more details of the original model Stargan v2, this repository will be useful. The source code, pre-trained models of StarGAN v2 are available under Creative Commons BY-NC 4.0 license by NAVER Corporation. You can use, copy, tranform and build upon the material for non-commercial purposes as long as you give appropriate credit by citing our paper, and indicate if changes were made.

Contact

For business inquiries, please contact [e-mail] or [e-mail].
For technical and other inquires, please contact [e-mail] or [e-mail].