Editing Out-of-domain GAN Inversion via Differential Activations

This is the official implementation of the paper "Editing Out-of-domain GAN Inversion via Differential Activations"

Prerequisite

• Linux
• NVIDIA GPU + CUDA CuDNN
• Python 3.7
• Pytorch >= 1.7, torchvision >= 0.8.2
• mmcv-full is required for some modules. The installation can be done as follows:

  pip install mmcv-full
  

• other packages (ttach, opencv-python):

  pip install ttach opencv-python
  

• more detailed environment configuration can be found in environment.yaml, which is directly exported by anaconda.

Getting Started

Preparation

• Clone the repository and enter the main folder.

  git clone git@github.com:HaoruiSong622/Editing-Out-of-Domain.git
  cd Editing-Out-of-Domain
  

• Download the pretrained checkpoints.

  Path	Description
  diff_cam_weight.pt	The weight for the DiffCAM in our model.
  deghosting.pt	The weight for the deghosting network.
  pSp_ffhq_encode.pt	Here we chose pSp encoder to do StyleGAN Inversion. Please Download the pretrained pSp checkpoint.

Run the Model

python image_process.py --device 0
--diffcam_ckpt_path path_to_diffcam_weight
--deghosting_ckpt_path path_to_deghosting_weight
--direction_path ./directions/Beard.npy
--image_dir ./sample_images
--output_dir path_to_output_dir
--psp_ckptpath path_to_psp_encoder_ffhq_weight

Training

In order to train our model, you need to train the Diff-CAM module and deghosting network one by one.

Training Diff-CAM Module

The first step is to train the Diff-CAM module. Run the following command to train the module.

python trainerDA.py --trainset_path path_to_training_dataset
--testset_path path_to_testing_dataset
--device 0
--DA_batch_size your_batch_size
--num_workers your_dataloader_num_workers
--direction_path ./directions
--exp_dir path_to_experiment_directory
--psp_ckptpath path_to_psp_encoder_ffhq_weight

Generating Ghosting Images

Before training the deghosting network, you need to generate the ghosting images dataset. Run the following command to generate ghosting images.

python gen_dataset.py --diffcam_ckpt_path path_to_diffcam_path
--direction_dir ./directions
--src_image_dir path_to_source_dataset
--dst_image_dir path_to_output_ghosting_dataset
--psp_ckptpath path_to_psp_encoder_ffhq_weight

Training Deghosting Network

After generating the ghosting image dataset, you can finally train the deghosting network. Run the following command to train the network.

python train_deghosting.py --trainset_lq_path path_to_trainset_lq
--testset_lq_path path_to_testset_lq
--trainset_tg_path path_to_trainset_tg
--testset_tg_path path_to_testset_tg
--exp_dir path_to_experiment_directory