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

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

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

The checkpoints and logs will be saved to checkpoints。

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

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

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