PD-GAN: Probabilistic Diverse GAN for Image Inpainting
Hongyu Liu, Ziyu Wan, Wei Huang, Yibing Song, Xintong Han and JIng Liao.
In CVPR 2021
Clone this repo.
git clone https://github.com/KumapowerLIU/PD-GAN.git
Prerequisites
- Python3
- Pytorch >=1.0
- Tensorboard
- Torchvision
- pillow
- loguru
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.
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 usingargparse
package. More individuals are dynamically added in other files as well.data/
: process the dataset before passing to the network.
# 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]
.
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.
- upload the test code
- upload the pre-trained model of PDGAN.
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}
}
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 [email protected].