Latent Subspace Optimization

The official Pytorch implementation of our paper Where is My Spot? Few-shot Image Generation via Latent Subspace Optimization, CVPR 2023.

framework

Where is My Spot? Few-shot Image Generation via Latent Subspace Optimization

Chenxi Zheng, Bangzhen Liu, Xuemiao Xu, Huaidong Zhang, and Shengfeng He

Paper

Environment

Python 3.7
Pytorch 1.7.1

Prepare the environment using either conda or pip.

cd envs
conda env create -f environment.yaml
cd ..

cd envs
conda create -n LSO python=3.7
conda activate LSO
pip install -r requirements.txt
cd ..

If installation of pytorch fails or a custom Pytorch version is needed, please install pytorch following the official guidance manually.

Getting started

Before training, prepare the pretrained ckpts optimized with seen categories.

Dataset: Follow LoFGAN and download the datasets.
Pretrained StyleGAN2: download the pretrained checkpoints.
Inverted latent codes: download inverted latent codes.
Weights for feature extractor: download the weights.

Note that the latent codes only contain the subset for image generation $\mathbb{S}_{gen}^{c}$ in Sec. 4.2. For example, the shape of tensor in flowers_unseen17_0-10_step1300.npy is $[17*10, 12, 512]$, which is corresponding to the slice $[85:102, 0:10, :, :, :]$ of images in flower_c8189_s128_data_rgb.npy which has a tensor shape of $[102, 40, 128, 128, 3]$.

Unzip the files and set up the paths in configs/default_configs.py.

DATA_PATH = <PATH_TO_DATASET_NPY>
CKPT_PATH = <PATH_TO_STYLEGAN2_CKPT>
WS_PATH = <PATH_TO_WS>
IDCKPT_PATH = <PATH_TO_IDWEIGHTS>

Training

Multi-task for quantitative evaluation

python train_unseen.py \
    --outdir <output_dir> \
    --k_shot <k> \
    --dataset_name <dataset_name>

The code does NOT support multi-GPU processing in the current version but allows separate runs for different categories. For instance, execute a run for categories $[0, 8)$ on GPU:0 and execute another for categories $[8, 17)$ on GPU:1. The indices can be specified with --classes=<start_idx>-<end_idx>.
Results are saved in directory <output_dir>/<run_dir>/few-shot_samples/magnitude=<magnitude-intensity>. For each category, we randomly generate 128 images for quantitative evaluation.
Refer to the main definition in train_unseen.py for command examples.

Single-task for detailed visualization

We also provide single-task optimization for visualization and detailed optimization evaluation.

python train_unseen.py \
    --outdir <output_dir> \
    --k_shot <k> \
    --single_task <cidx> <idx_1,...,idx_k> \
    --dataset_name <dataset_name>

Specify the argument --single_task with the category number and the image indices.
To save the checkpoints, add argument --save_ckpt as true.

Evaluation

(optional) If the images are generated in the separated runs, use merger.py to combine all the images.

python merger.py \
    --path <output_dir> \
    --idx <runidx_1>,...,<runidx_n>

Refer to the main definition in merger.py for command examples.

Quantitative evaluation of the generated images.

python main_metric_calculate.py \
    --real_dir <real_directory> \
    --fake_dir <fake_directory> \
    --dataset_name <dataset_name>

Refer to the main definition in main_metric_calculate.py for command examples.

The calculation of LPIPS is significantly accelerated by first extracting the features of each image. Please refer to metrics/lpips_fs/lpips_fs.py.

Citation

If you use this code for your research, please cite our paper.

@inproceedings{zheng2023my,
title={Where Is My Spot? Few-Shot Image Generation via Latent Subspace Optimization},
author={Zheng, Chenxi and Liu, Bangzhen and Zhang, Huaidong and Xu, Xuemiao and He, Shengfeng},
booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
pages={3272--3281},
year={2023}
}

Acknowledgement

This project builds upon and has been inspired by the following repositories:

Training and evaluation code are based on StyleGAN2-ADA and LoFGAN.
The code for calculating FID is based on pytorch-fid.

We would like to thank the entire open-source community for fostering an environment of collaboration and knowledge sharing.

License

This repository is under MIT license.

chansey0529 / LSO

readme