justHungryMan / generative-inpainting-pytorch

A PyTorch reimplementation for paper Generative Image Inpainting with Contextual Attention (https://arxiv.org/abs/1801.07892)

generative-inpainting-pytorch

A PyTorch reimplementation for the paper Generative Image Inpainting with Contextual Attention according to the author's TensorFlow implementation.

Prerequisites

This code has been tested on Ubuntu 14.04 and the following are the main components that need to be installed:

• Python3
• PyTorch 1.0+
• torchvision 0.2.0+
• tensorboardX
• pyyaml

Train the model

python train.py --config configs/config.yaml

The checkpoints and logs will be saved to checkpoints。

Test with the trained model

By default, it will load the latest saved model in the checkpoints. You can also use --iter to choose the saved models by iteration.

Trained PyTorch model: [Google Drive] [Baidu Wangpan]

python test_single.py \
	--image examples/imagenet/imagenet_patches_ILSVRC2012_val_00008210_input.png \
	--mask examples/center_mask_256.png \
	--output examples/output.png

Test with the converted TF model:

Converted TF model: [Google Drive]

python test_tf_model.py \
	--image examples/imagenet/imagenet_patches_ILSVRC2012_val_00008210_input.png \
	--mask examples/center_mask_256.png \
	--output examples/output.png \
	--model-path torch_model.p

Test results on ImageNet validation set patches

With PyTorch, the model was trained on ImageNet for 430k iterations to converge (with batch_size 48, about 150h). Here are some test results on the patches from ImageNet validation set.

Input	Inpainted