Awesome
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