inpainter_notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "eymXTgON6VDw"
   },
   "source": [
    "# Image Inpainting Based on Partial Convolutions in Keras\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 51
    },
    "colab_type": "code",
    "id": "i9XpNMpS6VDx",
    "outputId": "335c7c93-4b3c-4c36-ffbc-874e37e40fbf"
   },
   "outputs": [],
   "source": [
    "from keras.callbacks import ModelCheckpoint, TensorBoard, CSVLogger\n",
    "from keras.preprocessing.image import load_img, img_to_array\n",
    "from inpainter_utils.pconv2d_data import DataGenerator, torch_preprocessing, torch_postprocessing\n",
    "from inpainter_utils.pconv2d_model import pconv_model\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# SETTINGS:\n",
    "IMG_DIR_TRAIN   = \"data/images/train/\"\n",
    "IMG_DIR_VAL     = \"data/images/validation/\"\n",
    "IMG_DIR_TEST    = \"data/images/test/\"\n",
    "VGG16_WEIGHTS   = \"data/vgg16_weights/vgg16_pytorch2keras.h5\"\n",
    "WEIGHTS_DIR     = \"callbacks/weights/\"\n",
    "TB_DIR          = \"callbacks/tensorboard/\"\n",
    "CSV_DIR         = \"callbacks/csvlogger/\"\n",
    "BATCH_SIZE      = 5\n",
    "STEPS_PER_EPOCH = 2500\n",
    "EPOCHS_STAGE1   = 70\n",
    "EPOCHS_STAGE2   = 50\n",
    "LR_STAGE1       = 0.0002\n",
    "LR_STAGE2       = 0.00005\n",
    "STEPS_VAL       = 100\n",
    "BATCH_SIZE_VAL  = 4\n",
    "IMAGE_SIZE      = (512, 512)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hx7JNQEe6VD2"
   },
   "source": [
    "## Data generators"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "jd56DzbnHpym"
   },
   "outputs": [],
   "source": [
    "# DATA GENERATORS:\n",
    "train_datagen   = DataGenerator(preprocessing_function=torch_preprocessing, horizontal_flip=True)\n",
    "train_generator = train_datagen.flow_from_directory(\n",
    "    IMG_DIR_TRAIN,\n",
    "    target_size=IMAGE_SIZE,\n",
    "    batch_size=BATCH_SIZE\n",
    ")\n",
    "val_datagen   = DataGenerator(preprocessing_function=torch_preprocessing)\n",
    "val_generator = val_datagen.flow_from_directory(\n",
    "    IMG_DIR_VAL,\n",
    "    target_size=IMAGE_SIZE,\n",
    "    batch_size=BATCH_SIZE_VAL,\n",
    "    seed=22,\n",
    "    mask_init_seed=1,\n",
    "    total_steps=STEPS_VAL,\n",
    "    shuffle=False\n",
    ")\n",
    "# Create testing generator\n",
    "test_datagen = DataGenerator(preprocessing_function=torch_preprocessing)\n",
    "test_generator = test_datagen.flow_from_directory(\n",
    "    IMG_DIR_TEST,\n",
    "    target_size=IMAGE_SIZE,\n",
    "    batch_size=BATCH_SIZE\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "8SbedHEK6VD3"
   },
   "source": [
    "## Training\n",
    "### Stage 1. Initial training (BN enabled)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1748
    },
    "colab_type": "code",
    "id": "Iw4NWfLk6VD4",
    "outputId": "51ec8ddb-4aa3-463f-b644-e00c445ac346"
   },
   "outputs": [],
   "source": [
    "#LAST_CHECKPOINT = \"callbacks/weights/initial/weights.70-2.02-1.95.hdf5\"\n",
    "model = pconv_model(lr=LR_STAGE1, image_size=IMAGE_SIZE, vgg16_weights=VGG16_WEIGHTS)\n",
    "#model.load_weights(LAST_CHECKPOINT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "XgM2a1XrIcVm"
   },
   "outputs": [],
   "source": [
    "model.fit_generator(\n",
    "    train_generator,\n",
    "    steps_per_epoch=STEPS_PER_EPOCH,\n",
    "    epochs=EPOCHS_STAGE1,\n",
    "    validation_data=val_generator,\n",
    "    validation_steps=STEPS_VAL,\n",
    "    callbacks=[\n",
    "        CSVLogger(CSV_DIR + \"initial/log.csv\", append=True),\n",
    "        TensorBoard(log_dir=TB_DIR + \"initial/\", write_graph=True),\n",
    "        ModelCheckpoint(WEIGHTS_DIR + \"initial/weights.{epoch:02d}-{val_loss:.2f}-{loss:.2f}.hdf5\", monitor=\"val_loss\", verbose=1, save_weights_only=True)\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "IieKWBQ1KQ7v"
   },
   "source": [
    "### Stage 2. Fine-tuning (BN frozen in encoder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "epB7sp_-LT-w"
   },
   "outputs": [],
   "source": [
    "LAST_CHECKPOINT = WEIGHTS_DIR + \"initial/weights.80-1.94-1.83.hdf5\"\n",
    "model = pconv_model(fine_tuning=True, lr=LR_STAGE2, image_size=IMAGE_SIZE, vgg16_weights=VGG16_WEIGHTS)\n",
    "model.load_weights(LAST_CHECKPOINT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "l8Y6_ZTwLpQ8"
   },
   "outputs": [],
   "source": [
    "model.fit_generator(\n",
    "    train_generator,\n",
    "    steps_per_epoch=STEPS_PER_EPOCH,\n",
    "    initial_epoch=EPOCHS_STAGE1,\n",
    "    epochs=EPOCHS_STAGE1 + EPOCHS_STAGE2,\n",
    "    validation_data=val_generator,\n",
    "    validation_steps=STEPS_VAL,\n",
    "    callbacks=[\n",
    "        CSVLogger(CSV_DIR + \"fine_tuning/log.csv\", append=True),\n",
    "        TensorBoard(log_dir=TB_DIR + \"fine_tuning/\", write_graph=True),\n",
    "        ModelCheckpoint(WEIGHTS_DIR + \"fine_tuning/weights.{epoch:02d}-{val_loss:.2f}-{loss:.2f}.hdf5\", monitor=\"val_loss\", verbose=1, save_weights_only=True)\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "9t0s7Jvu6VD6"
   },
   "source": [
    "---\n",
    "## Prediction\n",
    "### Load the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "GuWrTjVUJgF1"
   },
   "outputs": [],
   "source": [
    "LAST_CHECKPOINT = WEIGHTS_DIR + \"fine_tuning/weights.120-1.73-1.78.hdf5\"\n",
    "model = pconv_model(predict_only=True, image_size=IMAGE_SIZE)\n",
    "model.load_weights(LAST_CHECKPOINT)\n",
    "k = 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "h_dAuNxENVcs"
   },
   "source": [
    "### First, try images with random masks from the train set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NFMq9XxdNePl",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Make a prediction for a batch of examples:\n",
    "(input_img, mask), orig_img = next(test_generator)\n",
    "output_img = model.predict([input_img, mask])\n",
    "\n",
    "# Post-processing:\n",
    "orig_img   = torch_postprocessing(orig_img)\n",
    "input_img  = torch_postprocessing(input_img) * mask # the (0,0,0) masked pixels are made grey by post-processing\n",
    "output_img = torch_postprocessing(output_img)\n",
    "output_comp = input_img.copy()\n",
    "output_comp[mask == 0] = output_img[mask == 0]\n",
    "\n",
    "fig, axes = plt.subplots(input_img.shape[0], 2, figsize=(15, 29))\n",
    "for i in range(input_img.shape[0]):\n",
    "    #axes[i,0].imshow(orig_img[i])\n",
    "    axes[i,0].imshow(input_img[i])\n",
    "    axes[i,1].imshow(output_img[i])\n",
    "    #axes[i,2].imshow(output_comp[i])\n",
    "    axes[i,0].tick_params(bottom=False, left=False, labelbottom=False, labelleft=False)\n",
    "    axes[i,1].tick_params(bottom=False, left=False, labelbottom=False, labelleft=False)\n",
    "axes[0,0].set_title('Masked image')\n",
    "axes[0,1].set_title('Prediction')\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"data/examples/{}_result.png\".format(k), bbox_inches='tight', pad_inches=0)\n",
    "plt.show()\n",
    "k += 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "p0JPQo5NJjDy"
   },
   "source": [
    "### Second, try on your own images and masks:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "HULQvllo6VD7",
    "outputId": "dc86af49-c499-4721-f328-f5d002c5a477"
   },
   "outputs": [],
   "source": [
    "img_fname  = \"data/examples/own_image.jpg\"\n",
    "mask_fname = \"data/examples/own_mask.jpg\"\n",
    "# Mask is assumed to have masked pixels in black and valid pixels in white\n",
    "\n",
    "# Loading and pre-processing:\n",
    "orig_img = img_to_array(load_img(img_fname, target_size=IMAGE_SIZE))\n",
    "orig_img = orig_img[None,...] \n",
    "\n",
    "mask = load_img(mask_fname, target_size=IMAGE_SIZE)\n",
    "mask = (img_to_array(mask) == 255).astype(np.float32)\n",
    "mask = mask[None,...] \n",
    "\n",
    "# Prediction:\n",
    "output_img = model.predict([torch_preprocessing(orig_img.copy()) * mask, mask])\n",
    "\n",
    "# Post-processing:\n",
    "output_img  = torch_postprocessing(output_img)\n",
    "input_img   = orig_img * mask\n",
    "output_comp = input_img.copy()\n",
    "output_comp[mask == 0] = output_img[mask == 0]\n",
    "\n",
    "# Plot:\n",
    "fig, axes = plt.subplots(2, 2, figsize=(20,20))\n",
    "axes[0,0].imshow(orig_img[0].astype('uint8'))\n",
    "axes[0,0].set_title('Original image')\n",
    "axes[0,1].imshow(mask[0])\n",
    "axes[0,1].set_title('Mask')\n",
    "axes[1,0].imshow(input_img[0].astype('uint8'))\n",
    "axes[1,0].set_title('Masked image')\n",
    "axes[1,1].imshow(output_img[0])\n",
    "axes[1,1].set_title('Prediction')\n",
    "for ax in axes.flatten():\n",
    "    ax.tick_params(bottom=False, left=False, labelbottom=False, labelleft=False)\n",
    "plt.tight_layout()\n",
    "plt.savefig(\"data/examples/own_image_result.png\", bbox_inches='tight', pad_inches=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "2P7wCOLf6VED"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "inpainter_notebook.ipynb",
   "provenance": [],
   "toc_visible": true,
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}