SRResNet-VGG22.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using TensorFlow backend.\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "import keras\n",
    "import tensorflow as tf \n",
    "from keras.layers import Conv2D, Dense\n",
    "from keras.layers import concatenate, add\n",
    "from keras.layers import BatchNormalization, Activation, PReLU\n",
    "from keras.layers import Input, Flatten, UpSampling2D\n",
    "from keras.models import Model\n",
    "from keras.optimizers import Adam\n",
    "from keras.callbacks import ModelCheckpoint, EarlyStopping\n",
    "import keras.backend as K\n",
    "from keras.applications.vgg19 import VGG19\n",
    "\n",
    "from shutil import rmtree\n",
    "\n",
    "from batch_generator import COCOBatchGenerator\n",
    "from bsd100_callback import BSD100_Evaluator\n",
    "from utils import print_available_devices, deprocess_HR, deprocess_LR\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "if tf.test.is_built_with_cuda():\n",
    "    data_format = 'channels_first'\n",
    "    keras.backend.set_image_data_format('channels_first')\n",
    "else:\n",
    "    data_format = 'channels_last'\n",
    "    keras.backend.set_image_data_format('channels_last')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Keras :  2.1.6\n",
      "\t data_format :  channels_first\n",
      "Tensorflow :  1.8.0\n",
      "Device : /device:CPU:0\n",
      "\t type : CPU\n",
      "\t desc :\n",
      "\n",
      "Device : /device:GPU:0\n",
      "\t type : GPU\n",
      "\t desc :device: 0, name: GeForce GTX 1070, pci bus id: 0000:02:00.0, compute capability: 6.1\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(\"Keras : \", keras.__version__)\n",
    "print(\"\\t data_format : \", keras.backend.image_data_format())\n",
    "print(\"Tensorflow : \", tf.__version__)\n",
    "print_available_devices()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epochs =  20\n"
     ]
    }
   ],
   "source": [
    "B = 16 # number of residual block\n",
    "\n",
    "batch_size = 16\n",
    "target_size = (96,96)\n",
    "downscale_factor = 4\n",
    "\n",
    "num_steps = 1e5\n",
    "steps_per_epoch = 5000\n",
    "epochs = int(num_steps//steps_per_epoch)\n",
    "\n",
    "# axis used in Parametric ReLU !\n",
    "shared_axis = [1,2] if data_format == 'channels_last' else [2,3]\n",
    "\n",
    "# axis for Batch Normalization\n",
    "axis = -1 if data_format == 'channels_last' else 1\n",
    "\n",
    "print('epochs = ', epochs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build the Generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# build a residual block\n",
    "def res_block(inputs):\n",
    "    x = Conv2D(64, kernel_size=(3,3), strides=(1,1), padding='same', activation=None, use_bias=False)(inputs)\n",
    "    x = BatchNormalization(axis=axis)(x)\n",
    "    x = PReLU(alpha_initializer='zeros', alpha_regularizer=None, alpha_constraint=None, shared_axes=shared_axis)(x)\n",
    "    x = Conv2D(64, kernel_size=(3,3), strides=(1,1), padding='same', activation=None, use_bias=False)(x)\n",
    "    x = BatchNormalization(axis=axis)(x)\n",
    "    return add([x, inputs])\n",
    "\n",
    "# build an upscale block\n",
    "# PixelShuffler is replaced by an UpSampling2D layer (nearest upsampling)\n",
    "\n",
    "def up_block(x):\n",
    "    x = Conv2D(256, kernel_size=(3,3), strides=(1,1) , padding='same', activation=None, use_bias=False)(x)\n",
    "    x = UpSampling2D(size=(2, 2))(x)\n",
    "    x = PReLU(alpha_initializer='zeros', alpha_regularizer=None, alpha_constraint=None, shared_axes=shared_axis)(x)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_generator = Input(shape=(None, None, 3) if data_format=='channels_last' else (3, None, None), \n",
    "                        name='input_generator')\n",
    "\n",
    "x = Conv2D(filters=64, kernel_size=(9,9),\n",
    "           strides=(1,1), padding='same',\n",
    "           activation=None)(input_generator)\n",
    "\n",
    "x_input_res_block = PReLU(alpha_initializer='zeros',\n",
    "                          alpha_regularizer=None,\n",
    "                          alpha_constraint=None,\n",
    "                          shared_axes=shared_axis)(x)\n",
    "\n",
    "x = x_input_res_block\n",
    "\n",
    "# add B residual blocks \n",
    "for i in range(B):\n",
    "    x = res_block(x)\n",
    "\n",
    "x = Conv2D(64, kernel_size=(3,3), strides=(1,1), padding='same', activation=None, use_bias=False)(x)\n",
    "x = BatchNormalization(axis=axis)(x)\n",
    "\n",
    "# skip connection\n",
    "x = add([x, x_input_res_block])\n",
    "\n",
    "# two upscale blocks\n",
    "x = up_block(x)\n",
    "x = up_block(x)\n",
    "\n",
    "# final conv layer : activated with tanh -> pixels in [-1, 1]\n",
    "output_generator = Conv2D(3, kernel_size=(9,9), \n",
    "                          strides=(1,1), activation='tanh',\n",
    "                          use_bias=False, padding='same')(x)\n",
    "\n",
    "generator = Model(inputs=input_generator, outputs=output_generator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "__________________________________________________________________________________________________\n",
      "Layer (type)                    Output Shape         Param #     Connected to                     \n",
      "==================================================================================================\n",
      "input_generator (InputLayer)    (None, 3, None, None 0                                            \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_1 (Conv2D)               (None, 64, None, Non 15616       input_generator[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "p_re_lu_1 (PReLU)               (None, 64, None, Non 64          conv2d_1[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_2 (Conv2D)               (None, 64, None, Non 36864       p_re_lu_1[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_1 (BatchNor (None, 64, None, Non 256         conv2d_2[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "p_re_lu_2 (PReLU)               (None, 64, None, Non 64          batch_normalization_1[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_3 (Conv2D)               (None, 64, None, Non 36864       p_re_lu_2[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_2 (BatchNor (None, 64, None, Non 256         conv2d_3[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "add_1 (Add)                     (None, 64, None, Non 0           batch_normalization_2[0][0]      \n",
      "                                                                 p_re_lu_1[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_4 (Conv2D)               (None, 64, None, Non 36864       add_1[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_3 (BatchNor (None, 64, None, Non 256         conv2d_4[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "p_re_lu_3 (PReLU)               (None, 64, None, Non 64          batch_normalization_3[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_5 (Conv2D)               (None, 64, None, Non 36864       p_re_lu_3[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_4 (BatchNor (None, 64, None, Non 256         conv2d_5[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "add_2 (Add)                     (None, 64, None, Non 0           batch_normalization_4[0][0]      \n",
      "                                                                 add_1[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_6 (Conv2D)               (None, 64, None, Non 36864       add_2[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_5 (BatchNor (None, 64, None, Non 256         conv2d_6[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "p_re_lu_4 (PReLU)               (None, 64, None, Non 64          batch_normalization_5[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_7 (Conv2D)               (None, 64, None, Non 36864       p_re_lu_4[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_6 (BatchNor (None, 64, None, Non 256         conv2d_7[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "add_3 (Add)                     (None, 64, None, Non 0           batch_normalization_6[0][0]      \n",
      "                                                                 add_2[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_8 (Conv2D)               (None, 64, None, Non 36864       add_3[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_7 (BatchNor (None, 64, None, Non 256         conv2d_8[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "p_re_lu_5 (PReLU)               (None, 64, None, Non 64          batch_normalization_7[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_9 (Conv2D)               (None, 64, None, Non 36864       p_re_lu_5[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_8 (BatchNor (None, 64, None, Non 256         conv2d_9[0][0]                   \n",
      "__________________________________________________________________________________________________\n",
      "add_4 (Add)                     (None, 64, None, Non 0           batch_normalization_8[0][0]      \n",
      "                                                                 add_3[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_10 (Conv2D)              (None, 64, None, Non 36864       add_4[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_9 (BatchNor (None, 64, None, Non 256         conv2d_10[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "p_re_lu_6 (PReLU)               (None, 64, None, Non 64          batch_normalization_9[0][0]      \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_11 (Conv2D)              (None, 64, None, Non 36864       p_re_lu_6[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_10 (BatchNo (None, 64, None, Non 256         conv2d_11[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "add_5 (Add)                     (None, 64, None, Non 0           batch_normalization_10[0][0]     \n",
      "                                                                 add_4[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_12 (Conv2D)              (None, 64, None, Non 36864       add_5[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_11 (BatchNo (None, 64, None, Non 256         conv2d_12[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "p_re_lu_7 (PReLU)               (None, 64, None, Non 64          batch_normalization_11[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_13 (Conv2D)              (None, 64, None, Non 36864       p_re_lu_7[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_12 (BatchNo (None, 64, None, Non 256         conv2d_13[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "add_6 (Add)                     (None, 64, None, Non 0           batch_normalization_12[0][0]     \n",
      "                                                                 add_5[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_14 (Conv2D)              (None, 64, None, Non 36864       add_6[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_13 (BatchNo (None, 64, None, Non 256         conv2d_14[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "p_re_lu_8 (PReLU)               (None, 64, None, Non 64          batch_normalization_13[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_15 (Conv2D)              (None, 64, None, Non 36864       p_re_lu_8[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_14 (BatchNo (None, 64, None, Non 256         conv2d_15[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "add_7 (Add)                     (None, 64, None, Non 0           batch_normalization_14[0][0]     \n",
      "                                                                 add_6[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_16 (Conv2D)              (None, 64, None, Non 36864       add_7[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_15 (BatchNo (None, 64, None, Non 256         conv2d_16[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "p_re_lu_9 (PReLU)               (None, 64, None, Non 64          batch_normalization_15[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_17 (Conv2D)              (None, 64, None, Non 36864       p_re_lu_9[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_16 (BatchNo (None, 64, None, Non 256         conv2d_17[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "add_8 (Add)                     (None, 64, None, Non 0           batch_normalization_16[0][0]     \n",
      "                                                                 add_7[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_18 (Conv2D)              (None, 64, None, Non 36864       add_8[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_17 (BatchNo (None, 64, None, Non 256         conv2d_18[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "p_re_lu_10 (PReLU)              (None, 64, None, Non 64          batch_normalization_17[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_19 (Conv2D)              (None, 64, None, Non 36864       p_re_lu_10[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_18 (BatchNo (None, 64, None, Non 256         conv2d_19[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "add_9 (Add)                     (None, 64, None, Non 0           batch_normalization_18[0][0]     \n",
      "                                                                 add_8[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_20 (Conv2D)              (None, 64, None, Non 36864       add_9[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_19 (BatchNo (None, 64, None, Non 256         conv2d_20[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "p_re_lu_11 (PReLU)              (None, 64, None, Non 64          batch_normalization_19[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_21 (Conv2D)              (None, 64, None, Non 36864       p_re_lu_11[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_20 (BatchNo (None, 64, None, Non 256         conv2d_21[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "add_10 (Add)                    (None, 64, None, Non 0           batch_normalization_20[0][0]     \n",
      "                                                                 add_9[0][0]                      \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_22 (Conv2D)              (None, 64, None, Non 36864       add_10[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_21 (BatchNo (None, 64, None, Non 256         conv2d_22[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "p_re_lu_12 (PReLU)              (None, 64, None, Non 64          batch_normalization_21[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_23 (Conv2D)              (None, 64, None, Non 36864       p_re_lu_12[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_22 (BatchNo (None, 64, None, Non 256         conv2d_23[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "add_11 (Add)                    (None, 64, None, Non 0           batch_normalization_22[0][0]     \n",
      "                                                                 add_10[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_24 (Conv2D)              (None, 64, None, Non 36864       add_11[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_23 (BatchNo (None, 64, None, Non 256         conv2d_24[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "p_re_lu_13 (PReLU)              (None, 64, None, Non 64          batch_normalization_23[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_25 (Conv2D)              (None, 64, None, Non 36864       p_re_lu_13[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_24 (BatchNo (None, 64, None, Non 256         conv2d_25[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "add_12 (Add)                    (None, 64, None, Non 0           batch_normalization_24[0][0]     \n",
      "                                                                 add_11[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_26 (Conv2D)              (None, 64, None, Non 36864       add_12[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_25 (BatchNo (None, 64, None, Non 256         conv2d_26[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "p_re_lu_14 (PReLU)              (None, 64, None, Non 64          batch_normalization_25[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_27 (Conv2D)              (None, 64, None, Non 36864       p_re_lu_14[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_26 (BatchNo (None, 64, None, Non 256         conv2d_27[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "add_13 (Add)                    (None, 64, None, Non 0           batch_normalization_26[0][0]     \n",
      "                                                                 add_12[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_28 (Conv2D)              (None, 64, None, Non 36864       add_13[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_27 (BatchNo (None, 64, None, Non 256         conv2d_28[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "p_re_lu_15 (PReLU)              (None, 64, None, Non 64          batch_normalization_27[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_29 (Conv2D)              (None, 64, None, Non 36864       p_re_lu_15[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_28 (BatchNo (None, 64, None, Non 256         conv2d_29[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "add_14 (Add)                    (None, 64, None, Non 0           batch_normalization_28[0][0]     \n",
      "                                                                 add_13[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_30 (Conv2D)              (None, 64, None, Non 36864       add_14[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_29 (BatchNo (None, 64, None, Non 256         conv2d_30[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "p_re_lu_16 (PReLU)              (None, 64, None, Non 64          batch_normalization_29[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_31 (Conv2D)              (None, 64, None, Non 36864       p_re_lu_16[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_30 (BatchNo (None, 64, None, Non 256         conv2d_31[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "add_15 (Add)                    (None, 64, None, Non 0           batch_normalization_30[0][0]     \n",
      "                                                                 add_14[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_32 (Conv2D)              (None, 64, None, Non 36864       add_15[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_31 (BatchNo (None, 64, None, Non 256         conv2d_32[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "p_re_lu_17 (PReLU)              (None, 64, None, Non 64          batch_normalization_31[0][0]     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_33 (Conv2D)              (None, 64, None, Non 36864       p_re_lu_17[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_32 (BatchNo (None, 64, None, Non 256         conv2d_33[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "add_16 (Add)                    (None, 64, None, Non 0           batch_normalization_32[0][0]     \n",
      "                                                                 add_15[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_34 (Conv2D)              (None, 64, None, Non 36864       add_16[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "batch_normalization_33 (BatchNo (None, 64, None, Non 256         conv2d_34[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "add_17 (Add)                    (None, 64, None, Non 0           batch_normalization_33[0][0]     \n",
      "                                                                 p_re_lu_1[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_35 (Conv2D)              (None, 256, None, No 147456      add_17[0][0]                     \n",
      "__________________________________________________________________________________________________\n",
      "up_sampling2d_1 (UpSampling2D)  (None, 256, None, No 0           conv2d_35[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "p_re_lu_18 (PReLU)              (None, 256, None, No 256         up_sampling2d_1[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_36 (Conv2D)              (None, 256, None, No 589824      p_re_lu_18[0][0]                 \n",
      "__________________________________________________________________________________________________\n",
      "up_sampling2d_2 (UpSampling2D)  (None, 256, None, No 0           conv2d_36[0][0]                  \n",
      "__________________________________________________________________________________________________\n",
      "p_re_lu_19 (PReLU)              (None, 256, None, No 256         up_sampling2d_2[0][0]            \n",
      "__________________________________________________________________________________________________\n",
      "conv2d_37 (Conv2D)              (None, 3, None, None 62208       p_re_lu_19[0][0]                 \n",
      "==================================================================================================\n",
      "Total params: 2,041,664\n",
      "Trainable params: 2,037,440\n",
      "Non-trainable params: 4,224\n",
      "__________________________________________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "generator.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Add VGG19 as Perceptual Loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preproces_vgg(x):\n",
    "    # scale from [-1,1] to [0, 255]\n",
    "    x += 1.\n",
    "    x *= 127.5\n",
    "    \n",
    "    # RGB -> BGR\n",
    "    if data_format == 'channels_last':\n",
    "        x = x[..., ::-1]\n",
    "    else:\n",
    "        x = x[:,::-1,:,:]\n",
    "    \n",
    "    # apply Imagenet preprocessing : BGR mean\n",
    "    mean = [103.939, 116.778, 123.68]\n",
    "    _IMAGENET_MEAN = K.constant(-np.array(mean))\n",
    "    x = K.bias_add(x, K.cast(_IMAGENET_MEAN, K.dtype(x)))\n",
    "    \n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def vgg_loss(y_true, y_pred):\n",
    "    # load pretrained VGG\n",
    "    vgg19 = VGG19(include_top=False,\n",
    "                  input_shape=target_size+(3,) if data_format=='channels_last' else (3,)+target_size, \n",
    "                  weights='imagenet')\n",
    "    vgg19.trainable = False\n",
    "    for l in vgg19.layers:\n",
    "        l.trainable = False\n",
    "    \n",
    "    # create a model that ouputs the features from level 'block2_conv2'\n",
    "    features_extractor = Model(inputs=vgg19.input, outputs=vgg19.get_layer(\"block2_conv2\").output)\n",
    "    \n",
    "    # compute the features, y_pred and y_true are tensors scaled in [-1, 1]\n",
    "    # it needs te be rescaled and shifted to respect VGG preprocessing strategy\n",
    "    features_pred = features_extractor(preproces_vgg(y_pred))\n",
    "    features_true = features_extractor(preproces_vgg(y_true))\n",
    "    \n",
    "    # adding the scaling factor (to have similar values as with MSE within image space)\n",
    "    return 0.006*K.mean(K.square(features_pred - features_true), axis=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/carnd-term1/lib/python3.5/site-packages/keras/applications/vgg19.py:191: UserWarning: You are using the TensorFlow backend, yet you are using the Theano image data format convention (`image_data_format=\"channels_first\"`). For best performance, set `image_data_format=\"channels_last\"` in your Keras config at ~/.keras/keras.json.\n",
      "  warnings.warn('You are using the TensorFlow backend, yet you '\n"
     ]
    }
   ],
   "source": [
    "opt = Adam(lr=1e-4, beta_1=0.9)\n",
    "generator.compile(loss=vgg_loss, optimizer=opt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 0 images\n"
     ]
    }
   ],
   "source": [
    "batch_gen = COCOBatchGenerator(directory='data/COCO/train2014/',\n",
    "                               batch_size=batch_size,\n",
    "                               downscale_factor=4,\n",
    "                               target_size=target_size,\n",
    "                               shuffle=True,\n",
    "                               crop_mode='fixed_size',\n",
    "                               color_mode='rgb',\n",
    "                               data_format=data_format)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_LR, batch_HR = batch_gen.next()\n",
    "print(batch_LR.shape)\n",
    "print(batch_HR.shape)\n",
    "\n",
    "if data_format == 'channels_first':\n",
    "    batch_LR = np.transpose(batch_LR, (0,2,3,1))\n",
    "    batch_HR = np.transpose(batch_HR, (0,2,3,1))\n",
    "\n",
    "# LR| HR \n",
    "fig, axes = plt.subplots(4, 2, figsize=(7, 15))\n",
    "for i in range(4):\n",
    "    axes[i, 0].imshow(deprocess_LR(batch_LR[i]).astype(np.uint8))\n",
    "    axes[i, 1].imshow(deprocess_HR(batch_HR[i]).astype(np.uint8))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train 'generator' with MSE Loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if os.path.isdir('checkpoints/SRResNet-VGG22/'):\n",
    "    rmtree('checkpoints/SRResNet-VGG22/')\n",
    "os.mkdir('checkpoints/SRResNet-VGG22/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = ModelCheckpoint(filepath='checkpoints/SRResNet-VGG22/weights.{epoch:02d}-{loss:.2f}.hdf5',\n",
    "                             monitor='loss',\n",
    "                             save_weights_only=True,\n",
    "                             period=10,\n",
    "                             verbose=2)\n",
    "\n",
    "best_checkpoint = ModelCheckpoint(filepath='checkpoints/SRResNet-VGG22/best_weights.hdf5',\n",
    "                                  monitor='loss',\n",
    "                                  save_weights_only=True,\n",
    "                                  save_best_only=True, \n",
    "                                  period=1,\n",
    "                                  verbose=2)\n",
    "\n",
    "bsd100_callback = BSD100_Evaluator(directory='data/BSD100/',\n",
    "                                   filepath='checkpoints/SRResNet-VGG22/BSD100/',\n",
    "                                   downscale_factor=4,\n",
    "                                   color_mode='rgb',\n",
    "                                   verbose=1,\n",
    "                                   data_format=data_format,\n",
    "                                   margin=5)\n",
    "\n",
    "\n",
    "early_stop = EarlyStopping(monitor='loss',min_delta=0, patience=3, verbose=1, mode='min')\n",
    "\n",
    "\n",
    "callbacks = [checkpoint, best_checkpoint, bsd100_callback, early_stop]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "history = generator.fit_generator(batch_gen,\n",
    "                                  epochs=epochs,\n",
    "                                  callbacks=callbacks,\n",
    "                                  steps_per_epoch=steps_per_epoch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(16, 4))\n",
    "plt.plot(history.history['loss'])\n",
    "plt.title('training loss')\n",
    "plt.ylabel('mse [-1, 1]')\n",
    "plt.xlabel('epoch')\n",
    "plt.show()\n",
    "\n",
    "plt.figure(figsize=(16, 4))\n",
    "plt.plot(bsd100_callback.logs['mse'])\n",
    "plt.title('validation loss')\n",
    "plt.ylabel('mse [0, 255]')\n",
    "plt.xlabel('epoch')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(16, 4))\n",
    "plt.plot(bsd100_callback.logs['psnr'])\n",
    "plt.title('validation PSNR')\n",
    "plt.ylabel('PSNR (dB)')\n",
    "plt.xlabel('epoch')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.5.2"
  },
  "widgets": {
   "state": {},
   "version": "1.1.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}