DeblurGAN.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "-QgPPbUoz1nE"
   },
   "source": [
    "# DeblurGAN: Blind Motion Deblurring Using Conditional Adversarial Networks\n",
    "\n",
    "- [Paper](https://arxiv.org/pdf/1711.07064.pdf)\n",
    "- [Dataset - GoPro_Large](https://drive.google.com/file/d/1H0PIXvJH4c40pk7ou6nAwoxuR4Qh_Sa2/view)\n",
    "\n",
    "## Authors\n",
    "- [Saujas Adarkar](https://github.com/Saujas)\n",
    "- [Aryan Mehra](https://github.com/aryanmehra1999)\n",
    "- [Siddhant Khandelwal](https://github.com/siddhantkhandelwal)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "D8R8cZqA6Znb"
   },
   "source": [
    "# Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 80
    },
    "colab_type": "code",
    "id": "xvTz2tBc3wfe",
    "outputId": "38419bc7-40ce-4e63-d697-710c78e420c5"
   },
   "outputs": [],
   "source": [
    "import random\n",
    "import keras\n",
    "import cv2\n",
    "import tqdm\n",
    "import click\n",
    "import datetime\n",
    "import math\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "import keras.backend as K\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from matplotlib.pyplot import imshow\n",
    "%matplotlib inline\n",
    "\n",
    "from PIL import Image, ImageEnhance\n",
    "\n",
    "from keras import losses\n",
    "from keras.backend.common import normalize_data_format\n",
    "from keras.engine import InputSpec\n",
    "from keras.engine.topology import Layer\n",
    "from keras.utils import conv_utils\n",
    "from keras.models import Model\n",
    "from keras.applications.vgg16 import VGG16\n",
    "from keras.callbacks import TensorBoard\n",
    "from keras.optimizers import Adam\n",
    "from keras.layers import Input, Activation, Add, UpSampling2D, BatchNormalization, Conv2D, Dense, concatenate, GlobalAveragePooling2D, AveragePooling2D, Flatten, Lambda\n",
    "from keras.layers.merge import Add\n",
    "from keras.layers.core import Dense, Dropout, Activation, Flatten, Reshape\n",
    "from keras.layers.convolutional import Convolution2D, MaxPooling2D, Conv2DTranspose, UpSampling2D, Conv2D\n",
    "from keras.layers.advanced_activations import LeakyReLU\n",
    "from keras.layers.normalization import BatchNormalization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 122
    },
    "colab_type": "code",
    "id": "dBzvaqyJb7b9",
    "outputId": "7841014b-b5aa-4f06-9385-e9f2560a52e0"
   },
   "outputs": [],
   "source": [
    "# If running the notebook on Google Colab, else change the paths in the following cells accordingly\n",
    "# Please add the dataset to your Google Drive, mount your Drive by running this cell.\n",
    "from google.colab import drive\n",
    "drive.mount('/content/drive')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "H0JRTQf36T7R"
   },
   "source": [
    "# Model Architecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "DO1sMITKTNvY"
   },
   "outputs": [],
   "source": [
    "# Parameters as defined in the paper\n",
    "\n",
    "channel_rate = 64\n",
    "image_shape = (256, 256, 3)\n",
    "patch_shape = (channel_rate, channel_rate, 3)\n",
    "RESHAPE = (256, 256)\n",
    "ngf = 64\n",
    "ndf = 64\n",
    "input_nc = 3\n",
    "output_nc = 3\n",
    "input_shape_generator = (256, 256, input_nc)\n",
    "input_shape_discriminator = (256, 256, output_nc)\n",
    "n_blocks_gen = 9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "cZFsw4sT_Ae7"
   },
   "outputs": [],
   "source": [
    "def res_block(input, filters, kernel_size=(3, 3), strides=(1, 1), use_dropout=False):\n",
    "    \"\"\"\n",
    "    Instantiates a Resnet Block using sequential API.\n",
    "\n",
    "    Parameters\n",
    "    input: Input tensor\n",
    "    filters: Number of filters to use\n",
    "    kernel_size: Shape of the kernel for the convolution\n",
    "    strides: Shape of the strides for the convolution\n",
    "    use_dropout: Boolean value to determine the use of dropout\n",
    "    \n",
    "    Returns\n",
    "    Keras Model\n",
    "    \"\"\"\n",
    "    block = ReflectionPadding2D((1, 1))(input)\n",
    "    block = Conv2D(filters=filters,\n",
    "               kernel_size=kernel_size,\n",
    "               strides=strides,)(block)\n",
    "    block = BatchNormalization()(block)\n",
    "    block = Activation('relu')(block)\n",
    "\n",
    "    if use_dropout:\n",
    "        block = Dropout(0.5)(block)\n",
    "\n",
    "    block = ReflectionPadding2D((1, 1))(block)\n",
    "    block = Conv2D(filters=filters,\n",
    "               kernel_size=kernel_size,\n",
    "               strides=strides,)(block)\n",
    "    block = BatchNormalization()(block)\n",
    "\n",
    "    merged = Add()([input, block])\n",
    "    return merged"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "nnC_aPFt_O62"
   },
   "outputs": [],
   "source": [
    "def spatial_reflection_2d_padding(x, padding=((1, 1), (1, 1)), data_format=None):\n",
    "    \"\"\"\n",
    "    Pad the 2nd and 3rd dimensions of a 4D tensor.\n",
    "\n",
    "    Parameters\n",
    "    x: Input tensor\n",
    "    padding: Shape of padding to use\n",
    "    data_format: Tensorflow vs Theano convention ('channels_last', 'channels_first')\n",
    "    \n",
    "    Returns\n",
    "    Tensorflow tensor\n",
    "    \"\"\"\n",
    "    assert len(padding) == 2\n",
    "    assert len(padding[0]) == 2\n",
    "    assert len(padding[1]) == 2\n",
    "\n",
    "    if data_format is None:\n",
    "        data_format = image_data_format()\n",
    "    \n",
    "    if data_format not in {'channels_first', 'channels_last'}:\n",
    "        raise ValueError('Unknown data_format ' + str(data_format))\n",
    "\n",
    "    if data_format == 'channels_first':\n",
    "        pattern = [[0, 0],\n",
    "                   [0, 0],\n",
    "                   list(padding[0]),\n",
    "                   list(padding[1])]\n",
    "    else:\n",
    "        pattern = [[0, 0],\n",
    "                   list(padding[0]), list(padding[1]),\n",
    "                   [0, 0]]\n",
    "    return tf.pad(x, pattern, \"REFLECT\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "n3zQM4tX_Tts"
   },
   "outputs": [],
   "source": [
    "class ReflectionPadding2D(Layer):\n",
    "\n",
    "    def __init__(self,\n",
    "                 padding=(1, 1),\n",
    "                 data_format=None,\n",
    "                 **kwargs):\n",
    "        super(ReflectionPadding2D, self).__init__(**kwargs)\n",
    "        self.data_format = normalize_data_format(data_format)\n",
    "        if isinstance(padding, int):\n",
    "            self.padding = ((padding, padding), (padding, padding))\n",
    "        elif hasattr(padding, '__len__'):\n",
    "            if len(padding) != 2:\n",
    "                raise ValueError('`padding` should have two elements. '\n",
    "                                 'Found: ' + str(padding))\n",
    "            height_padding = conv_utils.normalize_tuple(padding[0], 2,\n",
    "                                                        '1st entry of padding')\n",
    "            width_padding = conv_utils.normalize_tuple(padding[1], 2,\n",
    "                                                       '2nd entry of padding')\n",
    "            self.padding = (height_padding, width_padding)\n",
    "        else:\n",
    "            raise ValueError('`padding` should be either an int, '\n",
    "                             'a tuple of 2 ints '\n",
    "                             '(symmetric_height_pad, symmetric_width_pad), '\n",
    "                             'or a tuple of 2 tuples of 2 ints '\n",
    "                             '((top_pad, bottom_pad), (left_pad, right_pad)). '\n",
    "                             'Found: ' + str(padding))\n",
    "        self.input_spec = InputSpec(ndim=4)\n",
    "\n",
    "    def compute_output_shape(self, input_shape):\n",
    "        if self.data_format == 'channels_first':\n",
    "            if input_shape[2] is not None:\n",
    "                rows = input_shape[2] + self.padding[0][0] + self.padding[0][1]\n",
    "            else:\n",
    "                rows = None\n",
    "            if input_shape[3] is not None:\n",
    "                cols = input_shape[3] + self.padding[1][0] + self.padding[1][1]\n",
    "            else:\n",
    "                cols = None\n",
    "            return (input_shape[0],\n",
    "                    input_shape[1],\n",
    "                    rows,\n",
    "                    cols)\n",
    "        elif self.data_format == 'channels_last':\n",
    "            if input_shape[1] is not None:\n",
    "                rows = input_shape[1] + self.padding[0][0] + self.padding[0][1]\n",
    "            else:\n",
    "                rows = None\n",
    "            if input_shape[2] is not None:\n",
    "                cols = input_shape[2] + self.padding[1][0] + self.padding[1][1]\n",
    "            else:\n",
    "                cols = None\n",
    "            return (input_shape[0],\n",
    "                    rows,\n",
    "                    cols,\n",
    "                    input_shape[3])\n",
    "\n",
    "    def call(self, inputs):\n",
    "        return spatial_reflection_2d_padding(inputs,\n",
    "                                             padding=self.padding,\n",
    "                                             data_format=self.data_format)\n",
    "\n",
    "    def get_config(self):\n",
    "        config = {'padding': self.padding,\n",
    "                  'data_format': self.data_format}\n",
    "        base_config = super(ReflectionPadding2D, self).get_config()\n",
    "        return dict(list(base_config.items()) + list(config.items()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "oR_uUoLy_4JE"
   },
   "outputs": [],
   "source": [
    "def generator_model():\n",
    "    \"\"\"Build generator architecture.\"\"\"\n",
    "\n",
    "    inputs = Input(shape=image_shape)\n",
    "\n",
    "    x = ReflectionPadding2D((3, 3))(inputs)\n",
    "    x = Conv2D(filters=ngf, kernel_size=(7, 7), padding='valid')(x)\n",
    "    x = BatchNormalization()(x)\n",
    "    x = Activation('relu')(x)\n",
    "\n",
    "    n_downsampling = 2\n",
    "    for i in range(n_downsampling):\n",
    "        mult = 2**i\n",
    "        x = Conv2D(filters=ngf*mult*2, kernel_size=(3, 3), strides=2, padding='same')(x)\n",
    "        x = BatchNormalization()(x)\n",
    "        x = Activation('relu')(x)\n",
    "\n",
    "    mult = 2**n_downsampling\n",
    "    for i in range(n_blocks_gen):\n",
    "        x = res_block(x, ngf*mult, use_dropout=True)\n",
    "\n",
    "    for i in range(n_downsampling):\n",
    "        mult = 2**(n_downsampling - i)\n",
    "        # x = Conv2DTranspose(filters=int(ngf * mult / 2), kernel_size=(3, 3), strides=2, padding='same')(x)\n",
    "        x = UpSampling2D()(x)\n",
    "        x = Conv2D(filters=int(ngf * mult / 2), kernel_size=(3, 3), padding='same')(x)\n",
    "        x = BatchNormalization()(x)\n",
    "        x = Activation('relu')(x)\n",
    "\n",
    "    x = ReflectionPadding2D((3, 3))(x)\n",
    "    x = Conv2D(filters=output_nc, kernel_size=(7, 7), padding='valid')(x)\n",
    "    x = Activation('tanh')(x)\n",
    "\n",
    "    outputs = Add()([x, inputs])\n",
    "    # outputs = Lambda(lambda z: K.clip(z, -1, 1))(x)\n",
    "    outputs = Lambda(lambda z: z/2)(outputs)\n",
    "\n",
    "    model = Model(inputs=inputs, outputs=outputs, name='Generator')\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "YEz2GodkAWAx"
   },
   "outputs": [],
   "source": [
    "def discriminator_model():\n",
    "    \"\"\"Build discriminator architecture.\"\"\"\n",
    "    \n",
    "    n_layers, use_sigmoid = 3, False\n",
    "    inputs = Input(shape=input_shape_discriminator)\n",
    "\n",
    "    x = Conv2D(filters=ndf, kernel_size=(4, 4), strides=2, padding='same')(inputs)\n",
    "    x = LeakyReLU(0.2)(x)\n",
    "\n",
    "    nf_mult, nf_mult_prev = 1, 1\n",
    "    for n in range(n_layers):\n",
    "        nf_mult_prev, nf_mult = nf_mult, min(2**n, 8)\n",
    "        x = Conv2D(filters=ndf*nf_mult, kernel_size=(4, 4), strides=2, padding='same')(x)\n",
    "        x = BatchNormalization()(x)\n",
    "        x = LeakyReLU(0.2)(x)\n",
    "\n",
    "    nf_mult_prev, nf_mult = nf_mult, min(2**n_layers, 8)\n",
    "    x = Conv2D(filters=ndf*nf_mult, kernel_size=(4, 4), strides=1, padding='same')(x)\n",
    "    x = BatchNormalization()(x)\n",
    "    x = LeakyReLU(0.2)(x)\n",
    "\n",
    "    x = Conv2D(filters=1, kernel_size=(4, 4), strides=1, padding='same')(x)\n",
    "    if use_sigmoid:\n",
    "        x = Activation('sigmoid')(x)\n",
    "\n",
    "    x = Flatten()(x)\n",
    "    x = Dense(1024, activation='tanh')(x)\n",
    "    x = Dense(1, activation='sigmoid')(x)\n",
    "\n",
    "    model = Model(inputs=inputs, outputs=x, name='Discriminator')\n",
    "    return model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "aBJ9Oga3AgSt"
   },
   "outputs": [],
   "source": [
    "def generator_containing_discriminator(generator, discriminator):\n",
    "    inputs = Input(shape=image_shape)\n",
    "    generated_image = generator(inputs)\n",
    "    outputs = discriminator(generated_image)\n",
    "    model = Model(inputs=inputs, outputs=outputs)\n",
    "    return model\n",
    "\n",
    "def generator_containing_discriminator_multiple_outputs(generator, discriminator):\n",
    "    inputs = Input(shape=image_shape)\n",
    "    generated_image = generator(inputs)\n",
    "    outputs = discriminator(generated_image)\n",
    "    model = Model(inputs=inputs, outputs=[generated_image, outputs])\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NJDOJw7OAl3e"
   },
   "outputs": [],
   "source": [
    "# Defining losses as mentioned in the paper\n",
    "\n",
    "def perceptual_loss(y_true, y_pred):\n",
    "    vgg = VGG16(include_top=False, weights='imagenet', input_shape=image_shape)\n",
    "    loss_model = Model(inputs=vgg.input, outputs=vgg.get_layer('block3_conv3').output)\n",
    "    loss_model.trainable = False\n",
    "    return K.mean(K.square(loss_model(y_true) - loss_model(y_pred)))\n",
    "\n",
    "\n",
    "def wasserstein_loss(y_true, y_pred):\n",
    "    return K.mean(y_true*y_pred)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "YOjNvYPH6EqG"
   },
   "source": [
    "# Utility Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "kg69B2JDfp1X"
   },
   "outputs": [],
   "source": [
    "def is_an_npy_file(filename):\n",
    "    # Checks if the given file is a .npy file\n",
    "    NPY_EXTENSIONS = ['.npy']\n",
    "    for ext in NPY_EXTENSIONS:\n",
    "        if ext in filename:\n",
    "            return True\n",
    "    return False\n",
    "\n",
    "\n",
    "def list_npy_files(directory):\n",
    "    # Lists all the .npy files in a directory\n",
    "\n",
    "    files = sorted(os.listdir(directory))\n",
    "    return [os.path.join(directory, f) for f in files if is_an_npy_file(f)]\n",
    "\n",
    "\n",
    "def load_npy(path):\n",
    "    # Loads a .npy file from a given path\n",
    "\n",
    "    npy = np.load(path)\n",
    "    return npy\n",
    "\n",
    "\n",
    "def preprocess_npy(npy):\n",
    "    # Preprocess the .npy file\n",
    "\n",
    "    npy = (npy - 127.5) / 127.5\n",
    "    return npy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "BVqWp8-Kf0X8"
   },
   "outputs": [],
   "source": [
    "def load_npys(path, n_npy):\n",
    "    # Load n_npy number of .npy files from a given path randomly\n",
    "\n",
    "    if n_npy < 0:\n",
    "        n_npy = float(\"inf\")\n",
    "    A_paths, B_paths = os.path.join(path, 'A'), os.path.join(path, 'B')\n",
    "    all_A_paths, all_B_paths = list_npy_files(\n",
    "        A_paths), list_npy_files(B_paths)\n",
    "    \n",
    "    zipped_lists = list(zip(all_A_paths, all_B_paths))\n",
    "    random.shuffle(zipped_lists)\n",
    "\n",
    "    all_A_paths, all_B_paths = zip(*zipped_lists)\n",
    "\n",
    "    npy_A, npy_B = [], []\n",
    "    npy_A_paths, npy_B_paths = [], []\n",
    "    for path_A, path_B in zip(all_A_paths, all_B_paths):\n",
    "        np_A, np_B = load_npy(path_A), load_npy(path_B)\n",
    "        npy_A.append(preprocess_npy(np_A))\n",
    "        npy_B.append(preprocess_npy(np_B))\n",
    "        npy_A_paths.append(path_A)\n",
    "        npy_B_paths.append(path_B)\n",
    "        if len(npy_A) > n_npy - 1:\n",
    "            break\n",
    "\n",
    "    return {\n",
    "        'A': np.array(npy_A),\n",
    "        'A_paths': np.array(npy_A_paths),\n",
    "        'B': np.array(npy_B),\n",
    "        'B_paths': np.array(npy_B_paths)\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "KPNWoFwmDDsM"
   },
   "outputs": [],
   "source": [
    "def deprocess_image(img):\n",
    "    # Deprocess the image file\n",
    "\n",
    "    img = img * 127.5 + 127.5\n",
    "    return img.astype('uint8')\n",
    "\n",
    "\n",
    "def save_image(np_arr, path):\n",
    "    # Save the image from the corresponsing numpy array\n",
    "\n",
    "    img = np_arr * 127.5 + 127.5\n",
    "    im = Image.fromarray(img)\n",
    "    im.save(path)\n",
    "\n",
    "\n",
    "def save_all_weights(d, g, epoch_number, current_loss):\n",
    "    # Save all weights of the discriminator and generator models on given epoch number and loss value\n",
    "    \n",
    "    g.save(os.path.join('/content/drive/My Drive/ProcessedDeblurGAN_Dataset', 'generator_{}_{}.h5'.format(epoch_number+100, current_loss)))\n",
    "    d.save(os.path.join('/content/drive/My Drive/ProcessedDeblurGAN_Dataset', 'discriminator_{}.h5'.format(epoch_number+100)))\n",
    "\n",
    "def show_images(images, cols = 1, titles = None):\n",
    "    \"\"\"Display a list of images in a single figure with matplotlib.\"\"\"\n",
    "    assert((titles is None)or (len(images) == len(titles)))\n",
    "    n_images = len(images)\n",
    "    if titles is None: titles = ['Image (%d)' % i for i in range(1,n_images + 1)]\n",
    "    fig = plt.figure()\n",
    "    for n, (image, title) in enumerate(zip(images, titles)):\n",
    "        a = fig.add_subplot(cols, np.ceil(n_images/float(cols)), n + 1)\n",
    "        plt.imshow(image)\n",
    "        a.set_title(title)\n",
    "    fig.set_size_inches(np.array(fig.get_size_inches()) * n_images)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "hwmugDyL57Vl"
   },
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "XvdUxBAtEoDJ"
   },
   "outputs": [],
   "source": [
    "def train_multiple_outputs(n_images, batch_size, epoch_num, critic_updates=5, train=False):\n",
    "    if train==True:\n",
    "        data = load_npys('/content/drive/My Drive/ProcessedDeblurGAN_Dataset/ProcessedImages/train', n_images)\n",
    "        y_train, x_train = data['B'], data['A']\n",
    "\n",
    "    g = generator_model()\n",
    "    d = discriminator_model()\n",
    "    d_on_g = generator_containing_discriminator_multiple_outputs(g, d)\n",
    "\n",
    "    d_opt = Adam(lr=1E-4, beta_1=0.9, beta_2=0.999, epsilon=1e-08)\n",
    "    d_on_g_opt = Adam(lr=1E-4, beta_1=0.9, beta_2=0.999, epsilon=1e-08)\n",
    "\n",
    "    # To resume training from a particular epoch, load weights here\n",
    "    # g.load_weights(\"/content/drive/My Drive/ProcessedDeblurGAN_Dataset/generator_95_257.h5\")\n",
    "    # d.load_weights(\"/content/drive/My Drive/ProcessedDeblurGAN_Dataset/discriminator_95.h5\")\n",
    "\n",
    "    d.trainable = True\n",
    "    d.compile(optimizer=d_opt, loss=wasserstein_loss)\n",
    "    d.trainable = False\n",
    "    loss = [perceptual_loss, wasserstein_loss]\n",
    "    loss_weights = [100, 1]\n",
    "    d_on_g.compile(optimizer=d_on_g_opt, loss=loss, loss_weights=loss_weights)\n",
    "    d.trainable = True\n",
    "\n",
    "    output_true_batch, output_false_batch = np.ones((batch_size, 1)), -np.ones((batch_size, 1))\n",
    "\n",
    "    if train==True:\n",
    "        for epoch in tqdm.tqdm(range(epoch_num)):\n",
    "            permutated_indexes = np.random.permutation(x_train.shape[0])\n",
    "\n",
    "            d_losses = []\n",
    "            d_on_g_losses = []\n",
    "            for index in range(int(x_train.shape[0] / batch_size)):\n",
    "                batch_indexes = permutated_indexes[index*batch_size:(index+1)*batch_size]\n",
    "                image_blur_batch = x_train[batch_indexes]\n",
    "                image_full_batch = y_train[batch_indexes]\n",
    "\n",
    "                generated_images = g.predict(x=image_blur_batch, batch_size=batch_size)\n",
    "\n",
    "                for _ in range(critic_updates):\n",
    "                    d_loss_real = d.train_on_batch(image_full_batch, output_true_batch)\n",
    "                    d_loss_fake = d.train_on_batch(generated_images, output_false_batch)\n",
    "                    d_loss = 0.5 * np.add(d_loss_fake, d_loss_real)\n",
    "                    d_losses.append(d_loss)\n",
    "\n",
    "                d.trainable = False\n",
    "\n",
    "                d_on_g_loss = d_on_g.train_on_batch(image_blur_batch, [image_full_batch, output_true_batch])\n",
    "                d_on_g_losses.append(d_on_g_loss)\n",
    "\n",
    "                d.trainable = True\n",
    "\n",
    "            print(np.mean(d_losses), np.mean(d_on_g_losses))\n",
    "\n",
    "            if epoch%5==0:\n",
    "              save_all_weights(d, g, epoch, int(np.mean(d_on_g_losses)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 598
    },
    "colab_type": "code",
    "id": "_VGJhfUvEzUH",
    "outputId": "0d687d22-71d4-40c3-e998-5724d88ae6a7"
   },
   "outputs": [],
   "source": [
    "train_multiple_outputs(n_images=4, batch_size=1, epoch_num=1, train=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "mH53AnUG5-AG"
   },
   "source": [
    "# Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "KfalhhdgFgc6"
   },
   "outputs": [],
   "source": [
    "def test(batch_size):\n",
    "    data = load_npys('/content/drive/My Drive/ProcessedDeblurGAN_Dataset/ProcessedImages/test', batch_size)\n",
    "    y_test, x_test = data['B'], data['A']\n",
    "    g = generator_model()\n",
    "    g.load_weights('/content/drive/My Drive/ProcessedDeblurGAN_Dataset/generator_195_204.h5')\n",
    "    generated_images = g.predict(x=x_test, batch_size=batch_size)\n",
    "    generated = np.array([deprocess_image(img) for img in generated_images])\n",
    "    x_test = deprocess_image(x_test)\n",
    "    y_test = deprocess_image(y_test)\n",
    "\n",
    "    for i in range(generated_images.shape[0]):\n",
    "        y = y_test[i, :, :, :]\n",
    "        x = x_test[i, :, :, :]\n",
    "        img = generated[i, :, :, :]\n",
    "        \n",
    "        cv2.imwrite(\"/content/drive/My Drive/ProcessedDeblurGAN_Dataset/ImgEnh/deblurred/\" + str(i) + \".png\", img)\n",
    "        cv2.imwrite(\"/content/drive/My Drive/ProcessedDeblurGAN_Dataset/ImgEnh/blurred/\" + str(i) + \".png\", x)\n",
    "        cv2.imwrite(\"/content/drive/My Drive/ProcessedDeblurGAN_Dataset/ImgEnh/target/\" + str(i) + \".png\", y)\n",
    "\n",
    "        im = Image.open(\"/content/drive/My Drive/ProcessedDeblurGAN_Dataset/ImgEnh/deblurred/\" + str(i) + \".png\")\n",
    "        enhancer = ImageEnhance.Brightness(im)\n",
    "        enhanced_im = enhancer.enhance(1.05)\n",
    "        enhanced_im.save(\"/content/drive/My Drive/ProcessedDeblurGAN_Dataset/ImgEnh/deblurredFixed/\" + str(i) + \".png\")\n",
    "\n",
    "        output = np.concatenate((y, x, img), axis=1)\n",
    "        im = Image.fromarray(output.astype(np.uint8))\n",
    "        im.save('results{}.png'.format(i))\n",
    "    return generated_images\n",
    "\n",
    "def display_test_results(generated_images):\n",
    "    random_index = random.randint(0, generated_images.shape[0])\n",
    "    deblurred = Image.open(\"/content/drive/My Drive/ProcessedDeblurGAN_Dataset/ImgEnh/deblurred/\" + str(random_index) + \".png\")\n",
    "    target = Image.open(\"/content/drive/My Drive/ProcessedDeblurGAN_Dataset/ImgEnh/target/\" + str(random_index) + \".png\")\n",
    "    blurred = Image.open(\"/content/drive/My Drive/ProcessedDeblurGAN_Dataset/ImgEnh/blurred/\" + str(random_index) + \".png\")\n",
    "    show_images(images=[blurred, deblurred, target], cols=3, titles=[\"blurred\", \"deblurred\", \"target\"])\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ALKv58MEtH4w"
   },
   "outputs": [],
   "source": [
    "generated_images = test(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Tow-0lX9hj-r"
   },
   "outputs": [],
   "source": [
    "display_test_results(generated_images)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "QQp44bzc5orm"
   },
   "source": [
    "# Evaluation Metrics\n",
    "\n",
    "- PSNR\n",
    "- SSIM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "6eCgCFTaEXQN",
    "outputId": "ee0a490c-2722-4033-cdf3-8f23a4b74d53"
   },
   "outputs": [],
   "source": [
    "sess = tf.InteractiveSession()\n",
    "avg_psnr1 = 0\n",
    "avg_psnr2 = 0\n",
    "for i in range(100):\n",
    "    deblurred = Image.open(\"/content/drive/My Drive/ProcessedDeblurGAN_Dataset/ImgEnh/predicted/\" + str(i) + \".png\")\n",
    "    target = Image.open(\"/content/drive/My Drive/ProcessedDeblurGAN_Dataset/ImgEnh/target/\" + str(i) + \".png\")\n",
    "    blurred = Image.open(\"/content/drive/My Drive/ProcessedDeblurGAN_Dataset/ImgEnh/blurred/\" + str(i) + \".png\")\n",
    "\n",
    "    deblurred = np.array(deblurred)\n",
    "    target = np.array(target)\n",
    "    blurred = np.array(blurred)\n",
    "\n",
    "    deblurred = deblurred.reshape(1,256,256,3)\n",
    "    target = target.reshape(1,256,256,3)\n",
    "    blurred = blurred.reshape(1,256,256,3)\n",
    "\n",
    "    psnr1= tf.image.psnr(tf.convert_to_tensor(target,dtype=tf.float32), tf.convert_to_tensor(deblurred,dtype=tf.float32), max_val=255)\n",
    "    psnr2= tf.image.psnr(tf.convert_to_tensor(target,dtype=tf.float32), tf.convert_to_tensor(blurred,dtype=tf.float32), max_val=255)\n",
    "    \n",
    "    avg_psnr1 += psnr1\n",
    "    avg_psnr2 += psnr2\n",
    "\n",
    "    print(i)\n",
    "    print(\"PSNR P-T: \" + str(psnr1.eval()) + \" PSNR H-T: \" + str(psnr2.eval()))\n",
    "    print(\"\\n\")\n",
    "\n",
    "avg_psnr1 = avg_psnr1.eval()\n",
    "avg_psnr2 = avg_psnr2.eval()\n",
    "print(avg_psnr1/100)\n",
    "print(avg_psnr2/100)\n",
    "sess.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "id": "Mx4FTOjeaT0e",
    "outputId": "8b026ea4-377e-49fd-fc39-ae82e071072c"
   },
   "outputs": [],
   "source": [
    "sess = tf.InteractiveSession()\n",
    "avg_ssim1 = 0\n",
    "avg_ssim2 = 0\n",
    "for i in range(100):\n",
    "    deblurred = Image.open(\"/content/drive/My Drive/ProcessedDeblurGAN_Dataset/ImgEnh/predicted/\" + str(i) + \".png\")\n",
    "    target = Image.open(\"/content/drive/My Drive/ProcessedDeblurGAN_Dataset/ImgEnh/target/\" + str(i) + \".png\")\n",
    "    blurred = Image.open(\"/content/drive/My Drive/ProcessedDeblurGAN_Dataset/ImgEnh/blurred/\" + str(i) + \".png\")\n",
    "\n",
    "    deblurred = np.array(deblurred)\n",
    "    target = np.array(target)\n",
    "    blurred = np.array(blurred)\n",
    "\n",
    "    deblurred = deblurred.reshape(1,256,256,3)\n",
    "    target = target.reshape(1,256,256,3)\n",
    "    blurred = blurred.reshape(1,256,256,3)\n",
    "    \n",
    "    ssim1= tf.image.ssim(tf.convert_to_tensor(target,dtype=tf.float32), tf.convert_to_tensor(deblurred,dtype=tf.float32), max_val=255)\n",
    "    ssim2= tf.image.ssim(tf.convert_to_tensor(target,dtype=tf.float32), tf.convert_to_tensor(hazy,dtype=tf.float32), max_val=255)\n",
    "    \n",
    "    avg_ssim1 += ssim1\n",
    "    avg_ssim2 += ssim2\n",
    "    \n",
    "    print(i)\n",
    "    print(\"SSIM DT: \" + str(ssim1.eval()) + \" SSIM DFT: \" + str(ssim2.eval()))\n",
    "    print(\"\\n\")\n",
    "\n",
    "avg_ssim1 = avg_ssim1.eval()\n",
    "avg_ssim2 = avg_ssim2.eval()\n",
    "print(avg_ssim1/100)\n",
    "print(avg_ssim2/100)\n",
    "sess.close()"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [
    "H0JRTQf36T7R",
    "YOjNvYPH6EqG",
    "hwmugDyL57Vl",
    "mH53AnUG5-AG",
    "QQp44bzc5orm"
   ],
   "machine_shape": "hm",
   "name": "DeBlurGAN.ipynb",
   "provenance": []
  },
  "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}