longlongaaago / exe-gan

Facial image inpainting is a task of filling visually realistic and semantically meaningful contents for missing or masked pixels in a face image. This paper presents EXE-GAN, a novel diverse and interactive facial inpainting framework, which can not only preserve the high-quality visual effect of the whole image but also complete the face image with exemplar-like facial attributes.

Do Inpainting Yourself: Generative Facial Inpainting Guided by Exemplars (EXE-GAN)

Official PyTorch implementation of EXE-GAN. [Homepage] [paper] [demo_youtube] [demo_bilibili]

We present EXE-GAN, a novel exemplar-guided facial inpainting framework using generative adversarial networks. Our approach can not only preserve the quality of the input facial image but also complete the image with exemplar-like facial attributes.

Notice

Our paper was first released on Sun, 13 Feb 2022. We are thankful for the community's recognition and attention to our project. We also recognized that there have been some great papers published after ours, and we encourage you to check out their projects as well:

• Paint by Example, codes (released at Wed, 23 Nov 2022, CVPR 2023)
• Reference-Guided Face Inpainting, codes (released at Mon, 13 Mar 2023, TCSVT 2023)
• PATMAT, codes (released at Wed, 12 Apr 2023, ICCV 2023)

Requirements

cd EXE-GAN project
pip install -r requirements.txt

• Note that other versions of PyTorch (e.g., higher than 1.7) also work well, but you have to install the corresponding CUDA version.

What we have released

• Training and testing codes
• Pre-trained models

Training

• Prepare your dataset (download FFHQ, and CelebA-HQ)
• The folder structure of training and testing data is shown below:

root/
    test/
        xxx.png
        ...
        xxz.png
    train/
        xxx.png
        ...
        xxz.png

• Prepare pre-trained checkpoints: Arcface.pth and psp_ffhq_encode.pt (put models in ./pre-train)

• Training

python train.py --path /root/train --test_path /root/test --size 256 --embedding_weight 0.1 --id_loss_weight 0.1 --percept_loss_weight 0.5 --arcface_path ./pre-train/Arcface.pth --psp_checkpoint_path ./pre-train/psp_ffhq_encode.pt

Testing

Notice

• For editing images from the web, photos should be aligned by face landmarks and cropped to 256x256 by align_face.

• Irregular masks (optional, if you would like to test on irregular masks, download Testing Set masks)

• (use our FFHQ_60k pre-trained model EXE_GAN_model.pt or trained *pt file by yourself.)

python test.py --path /root/test --size 256 --psp_checkpoint_path ./pre-train/psp_ffhq_encode.pt --ckpt ./checkpoint/EXE_GAN_model.pt --mask_root ./dataset/mask/testing_mask_dataset --mask_file_root ./dataset/mask --mask_type test_6.txt

- mask_root Irregular masks root
- mask_file_root file name list file folder
- mask_type could be ["center", "test_2.txt", "test_3.txt", "test_4.txt", "test_5.txt", "test_6.txt", "all"]

• If you don't have irregular masks, just using center masks is also fine.

python test.py --path /root/test --size 256 --psp_checkpoint_path ./pre-train/psp_ffhq_encode.pt --ckpt ./checkpoint/EXE_GAN_model.pt --mask_type center

Exemplar-guided facial image recovery

Notice

• For editing images from the web, photos should be aligned by face landmarks and cropped to 256x256 by align_face.

(use our FFHQ_60k pre-trained model EXE_GAN_model.pt or trained *pt file by yourself.)

python guided_recovery.py --psp_checkpoint_path ./pre-train/psp_ffhq_encode.pt --ckpt ./checkpoint/EXE_GAN_model.pt --masked_dir ./imgs/exe_guided_recovery/mask --gt_dir ./imgs/exe_guided_recovery/target --exemplar_dir ./imgs/exe_guided_recovery/exemplar --sample_times 10 --eval_dir ./recover_out

- masked_dir: mask input folder
- gt_dir: the input gt_dir, used for  editing 
- exemplar_dir: exemplar_dir, the exemplar dir, for guiding the editing
- eval_dir: output dir

Ground-truth	Mask	Exemplar	Inpainted

• Inherent diversity, set --sample_times 10 higher to get more diverse results.

diversity 1	diversity 2	diversity 3	diversity 4

Exemplar guided style mixing

Notice

• For editing images from the web, photos should be aligned by face landmarks and cropped to 256x256 by align_face.

(use our FFHQ_60k pre-trained model EXE_GAN_model.pt or trained *pt file by yourself.)

python exemplar_style_mixing.py --psp_checkpoint_path ./pre-train/psp_ffhq_encode.pt --ckpt ./checkpoint/EXE_GAN_model.pt --masked_dir ./imgs/exe_guided_recovery/mask --gt_dir ./imgs/exe_guided_recovery/target --exemplar_dir ./imgs/exe_guided_recovery/exemplar --sample_times 2 --eval_dir mixing_out

- masked_dir: mask input folder
- gt_dir: the input gt_dir, used for  editing 
- exemplar_dir: exemplar_dir, the exemplar dir, for guiding the editing
- eval_dir: output dir

• Inputs are shown below:

Ground-truth	Mask	Exemplar 1	Exemplar 2

• Style mixing results

Editing masks by yourself

We also uploaded the mask editing tool. You can try this tool to generate your masks for editing.

python mask_gui.py

Bibtex

• If you find our code useful, please cite our paper:

  @misc{lu2022inpainting,
      title={Do Inpainting Yourself: Generative Facial Inpainting Guided by Exemplars}, 
      author={Wanglong Lu and Hanli Zhao and Xianta Jiang and Xiaogang Jin and Yongliang Yang and Min Wang and Jiankai Lyu and Kaijie Shi},
      year={2022},
      eprint={2202.06358},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
  }
  

Acknowledgements

Model details and custom CUDA kernel codes are from official repositories: https://github.com/NVlabs/stylegan2

Codes for Learned Perceptual Image Patch Similarity, LPIPS came from https://github.com/richzhang/PerceptualSimilarity

To match FID scores more closely to tensorflow official implementations, I have used FID Inception V3 implementations in https://github.com/mseitzer/pytorch-fid

exe-gan's Issues

Question regarding algorithm 1

Hello! thank you for your great work!

I checked train.py and got a question.

        rand_num = random.randint(1,rand_end)
        # for inference
        if rand_num == rand_end:
            infer_img = real_img
        else:
            infer_img = torch.flip(real_img, dims=[0])

It looks like given integer (rand_end), pick a integer N ~ Uniform({1,...,rand_end}), if N==rand_end, flip the real_img.

I believe it is different from the description in the algorithm 1 of your paper.

exemplar_style_mixing.py generate only the images in the provided examples

Hi, I am excited about your model and tried to use it with my images for style mixing. I have a ground truth image and an example image which are face aligned and 256x256. also have a mask with 256x256. when I ran the following command in colab after all the installation steps, I got the same results mentioned in the ReadME. not for my images.

command I ran is as follows:

%cd /content/EXE-GAN
!python exemplar_style_mixing.py
--psp_checkpoint_path /content/psp_ffhq_encode.pt
--ckpt /content/EXE_GAN_model.pt
--masked_dir /content/EXE-GAN/imgs/exe_guided_recovery/mask
--gt_dir /content/EXE-GAN/imgs/exe_guided_recovery/target
--exemplar_dir /content/EXE-GAN/imgs/exe_guided_recovery/exemplar
--sample_times 2
--eval_dir /content/EXE-GAN/imgs/exe_guided_recovery/mixing_out

Any help on this would be highly appriciated

Error loading PSP checkpoint:

Hi, I ran into this issue when loading psp checkpoint to perform image guided recovery:

⚡ main ~/EXE-GAN python3 guided_recovery.py --psp_checkpoint_path ./pre-train/psp_ffhq_encode.pt --ckpt ./checkpoint/EXE_GAN_model.pt --masked_dir ./imgs/exe_guided_recovery/mask --gt_dir ./imgs/exe_guided_recovery/target --exemplar_dir ./imgs/exe_guided_recovery/exemplar --sample_times 10 --eval_dir .imgs/exe_guided_recovery/recover_out_gpu_run from .fused_act import FusedLeakyReLU, fused_leaky_relu from .upfirdn2d import upfirdn2d model name: exe_gan !!!!!!!!!!!!!!! start_latent :4 n_psp_latent :10 Loading ResNet ArcFace Traceback (most recent call last): File "/teamspace/studios/this_studio/EXE-GAN/guided_recovery.py", line 199, in <module> generator = get_model(args.arch, model_path=args.ckpt, psp_path=args.psp_checkpoint_path) File "/teamspace/studios/this_studio/EXE-GAN/test.py", line 299, in get_model generator = EXE_GAN(exe_ckpt_path=model_path, psp_ckpt_path=psp_path) File "/teamspace/studios/this_studio/EXE-GAN/test.py", line 46, in __init__ self.psp_embedding = Psp_Embedding(psp_ckpt_path, psp_start_latent, num_psp_latent).to(device) File "/teamspace/studios/this_studio/EXE-GAN/Loss/psp_embedding.py", line 38, in __init__ self.psp_encoder = self.get_psp_encoder(psp_encoder_path=psp_encoder_path).eval().cuda() File "/teamspace/studios/this_studio/EXE-GAN/Loss/psp_embedding.py", line 111, in get_psp_encoder ckpt = torch.load(psp_encoder_path, map_location='cpu') File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/serialization.py", line 815, in load return _legacy_load(opened_file, map_location, pickle_module, **pickle_load_args) File "/home/zeus/miniconda3/envs/cloudspace/lib/python3.10/site-packages/torch/serialization.py", line 1051, in _legac y_load typed_storage._untyped_storage._set_from_file( RuntimeError: unexpected EOF, expected 820379 more bytes. The file might be corrupted.

I used Python 3.10 and followed most of the dependencies. My PyTorch versions are:
torch==2.0.0 torchaudio==2.0.1 torchvision==0.15.1
and my virtual environment is conda.

Do you know what I might have done wrong? Can't figure out after a bit of search. Thank you!

Do you have a pretrained model?

Do you have a pretrained model?

有源码吗？

什么时候释放源码？