Home

Awesome

Generative Multi-column Convolutional Neural Networks inpainting model in Keras

Keras implementation of GMCNN (Generative Multi-column Convolutional Neural Networks) inpainting model originally proposed at NIPS 2018: Image Inpainting via Generative Multi-column Convolutional Neural Networks

Model architecture

GMCNN model

Installation

Code download:

git clone https://github.com/tlatkowski/inpainting-gmcnn-keras.git
cd inpainting-gmcnn-keras

To install requirements, create Python virtual environment and install dependencies from files:

virtualenv -p /usr/bin/python3.6 .venv
source .venv/bin/activate
pip install -r requirements/requirements.txt

In case of using GPU support:

pip install -r requirements/requirements-gpu.txt

Otherwise (CPU usage):

pip install -r requirements/requirements-cpu.txt

Datasets

Image dataset

Model was trained with usage of high-resolution images from Places365-Standard dataset. It can be found here

Mask dataset

The mask dataset used for model training comes from NVIDIA's paper: Image Inpainting for Irregular Holes Using Partial Convolutions

NVIDIA's mask dataset is available here

Please note that the model training was performed on testing irregular mask dataset containing 12,000 masks.

./samples folder contains exemplary structure of dataset catalogs:

samples
 |-masks
    |-nvidia_masks
 |-images
    |-places365

nvidia_masks catalog contains 5 sample masks from NVIDIA's test set.

places365 catalog contains 5 sample images form Places365 validation set.

Model training

The main configuration file is placed in ./config/main_config.ini. It contains training and model parameters. You can tweak those parameters before model running.

The default configuration looks as follows:

[TRAINING]
WGAN_TRAINING_RATIO = 5
NUM_EPOCHS = 5
BATCH_SIZE = 4
IMG_HEIGHT = 256
IMG_WIDTH = 256
NUM_CHANNELS = 3
LEARNING_RATE = 0.0001
SAVE_MODEL_STEPS_PERIOD = 1000

[MODEL]
ADD_MASK_AS_GENERATOR_INPUT = False
GRADIENT_PENALTY_LOSS_WEIGHT = 10
ID_MRF_LOSS_WEIGHT = 0.05
ADVERSARIAL_LOSS_WEIGHT = 0.001
NN_STRETCH_SIGMA = 0.5
VGG_16_LAYERS = 3,6,10
ID_MRF_STYLE_WEIGHT = 1.0
ID_MRF_CONTENT_WEIGHT = 1.0
NUM_GAUSSIAN_STEPS = 3
GAUSSIAN_KERNEL_SIZE = 32
GAUSSIAN_KERNEL_STD = 40.0

After the dependencies installation you can perform training dry-run using image and mask samples provided in samples directory. To do so, execute the following command:

NOTE: Set BATCH_SIZE to 1 before executing the below command.

python runner.py --train_path ./samples/images --mask_path ./samples/masks --experiment_name "dry-run-test"

If everything goes correct you should be able to see the progress bar logging the basic training metrics.

In order to run GMCNN model training on your training data you have to provide paths to your datasets:

python runner.py --train_path /path/to/training/images --mask_path /path/to/mask/images --experiment_name "experiment_name"

Warm-up generator training

According to the best practices of the usage of GAN frameworks, first we should train the generator model for a while. In order to train the generator only in the first line run the following command (additional flag warm_up_generator is set):

python runner.py --train_path /path/to/training/images --mask_path /path/to/mask/images -warm_up_generator

In this mode the generator will be trained with only confidence-driven reconstruction loss.

Below picture presents GMCNN outcome after 5 epochs training in warm-up generator mode

WGAN-GP training

In order to continue training with full WGAN-GP framework (GMCNN generator, local and global discriminators), execute:

python runner.py --train_path /path/to/training/images --mask_path /path/to/mask/images --experiment_name "experiment_name" -from_weights

Running training with additional from_weights flag will force pipeline to load the latest models checkpoints from ./outputs/weights/ directory.

GMCNN model training in Google Colab notebook

If you don't have an access to workstation with GPU, you can use the below exemplary Google Colab notebook for training your GMCNN model on Places365 validation data and NVIDIA's testing mask with usage of K80 GPU available within Google Colab backend: GMCNN in Google Colab

Pipeline outcomes

During the training procedure the pipeline logs additional results to the outputs directory:

You can track the metrics during the training with usage of TensorBoard:

tensorboard --logdir=./outputs/experiment_name/logs

Implementation differences from original paper

  1. This model is trained using NVIDIA's irregular mask test set whereas the original model is trained using randomly generated rectangle masks.
  2. The current version of pipeline uses the higher-order features extracted from VGG16 model whereas the original model utilizes VGG19.

Visualization of Gaussian blurring masks

Below you can find the visualization of applying Gaussian blur to the training masks for the different number of convolution steps (number of iteration steps over the input raw mask).

Large mask

Original1 step2 steps3 steps4 steps5 steps10 steps

Small mask

Original1 step2 steps3 steps4 steps5 steps10 steps

Rectangle mask

Original1 step2 steps3 steps4 steps5 steps10 steps

Visualization of training losses

After activating TensorBoard you can monitor the following training metrics:

  1. For the generator: confidence reconstruction loss, global wasserstein loss, local wasserstein loss, id mrf loss and total loss
  2. For the local and global discriminators: fake loss, real loss, gradient penalty loss and total loss

Code References

  1. ID-MRF loss function was implemented with usage of original Tensorflow implementation: GMCNN in Tensorflow
  2. Improved Wasserstain GAN was implemented based on: Wasserstein GAN with gradient penalty in Keras
  3. Model architecture diagram was done with usage of PlotNeuralNet: PlotNeuralNet on GitHub