KumapowerLIU / PD-GAN

The official pytorch code of PD-GAN: Probabilistic Diverse GAN for Image Inpainting (CVPR 2021)

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

License CC BY-NC-SA 4.0 Python 3.6

PD-GAN

PDGAN Show

Paper | BibTex

PD-GAN: Probabilistic Diverse GAN for Image Inpainting Hongyu Liu, Ziyu Wan, Wei Huang, Yibing Song, Xintong Han and JIng Liao.
In CVPR 2021

Installation

Clone this repo.

git clone https://github.com/KumapowerLIU/PD-GAN.git

Prerequisites

  • Python3
  • Pytorch >=1.0
  • Tensorboard
  • Torchvision
  • pillow
  • loguru

Dataset Preparation

We use Places2, CelebA and Paris Street-View datasets. To train a model on the full dataset, download datasets from official websites.

Our model is trained on the irregular mask dataset provided by Liu et al. You can download the publically available Irregular Mask Dataset from their website.

Code Structure

PDGAN Show

  • train.py: the entry point for training.
  • train.py: the entry point for test.
  • models/PConv.py: defines the loss, model, optimizetion, foward, backward and others of Partial Conv.
  • models/PDGAN.py: defines the loss, model, optimizetion, foward, backward and others of PDGAN.
  • models/network/pcconv.py: defines the network of Partial Conv.
  • models/network/pdgan.py: defines the network of PDGAN.
  • config/: creates option lists using argparse package. More individuals are dynamically added in other files as well.
  • data/: process the dataset before passing to the network.

Training New Models

# To train on the you dataset for ParitalConv, for example.
python train.py --gt_root [the path of structure images]   --mask_root [the path of mask images] --model PConv
# To train on the you dataset for PDGAN, for example.
python train.py --gt_root [the path of structure images]   --mask_root [the path of mask images] --model PDGAN --pre_ckpt_PConv_EN_dir [the path of pre-trained PartialConv encoder] --pre_ckpt_PConv_DE_dir [the path of pre-trained PartialConv decoder]

There are many options you can specify. Please use python train.py --help or see the options To log training, use --./logs for Tensorboard. The logs are stored at logs/[name].

Pre-trained weights and test model

There are three folders to present a pre-trained model for three datasets, respectively. Currently, we only upload the pre-trained model of Partical Conv. in each folder. We implement the Partial Conv by ourselves.

TODO

  • upload the test code
  • upload the pre-trained model of PDGAN.

Citation

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

@inproceedings{liu2021pd,
  title={Pd-gan: Probabilistic diverse gan for image inpainting},
  author={Liu, Hongyu and Wan, Ziyu and Huang, Wei and Song, Yibing and Han, Xintong and Liao, Jing},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  pages={9371--9381},
  year={2021}
}

License

The codes and the pretrained model in this repository are under the MIT license as specified by the LICENSE file. The code is released for academic research use only. For commercial use, please contact kumapower@hnu.edu.cn.

About

The official pytorch code of PD-GAN: Probabilistic Diverse GAN for Image Inpainting (CVPR 2021)

License:MIT License


Languages

Language:Python 100.0%