pix2pix.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import cv2\n",
    "import keras\n",
    "import tensorflow as tf\n",
    "import h5py\n",
    "\n",
    "from os import listdir\n",
    "\n",
    "from numpy import load\n",
    "from numpy import zeros\n",
    "from numpy import ones\n",
    "from numpy import expand_dims\n",
    "from numpy.random import randint\n",
    "from numpy import vstack\n",
    "from numpy import asarray\n",
    "from numpy import savez_compressed\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "from keras import Sequential\n",
    "from keras.preprocessing.image import load_img\n",
    "from keras.preprocessing.image import img_to_array\n",
    "from keras.initializers import RandomNormal\n",
    "from keras.models import Model\n",
    "from keras.models import Input\n",
    "from keras.models import load_model\n",
    "from keras.optimizers import Adam\n",
    "from keras.layers import Conv2D\n",
    "from keras.layers import Conv2DTranspose\n",
    "from keras.layers import LeakyReLU\n",
    "from keras.layers import Activation\n",
    "from keras.layers import Concatenate\n",
    "from keras.layers import Dropout\n",
    "from keras.layers import BatchNormalization\n",
    "from keras.layers import ReLU"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data Preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_images(path, size=(256,512)):\n",
    "    \n",
    "    # path : Path for the training images directory\n",
    "    # size : Size of each image. For our case one image contains two separate images of size 256*256.\n",
    "\n",
    "    # load all images in a directory into memory and returns them as a combined array\n",
    "\n",
    "    src_list, tar_list = list(), list()\n",
    "\n",
    "    # enumerate all images in directory\n",
    "    for filename in listdir(path):\n",
    "        # load and resize the image\n",
    "        pixels = load_img(path + filename, target_size=size)\n",
    "        # convert to numpy array\n",
    "        pixels = img_to_array(pixels)\n",
    "        # split into source and target\n",
    "        s_img, t_img = pixels[:, :256], pixels[:, 256:]\n",
    "        src_list.append(s_img)\n",
    "        tar_list.append(t_img)\n",
    "        print(filename)\n",
    "        \n",
    "    print('Done')\n",
    "    return [asarray(src_list), asarray(tar_list)]\n",
    "\n",
    "def compress_images(src_images, tar_images, path_dest):\n",
    "    \n",
    "    # src_images : source image array\n",
    "    # tar_images : target image array\n",
    "    # path_dest : path to store the compressed .npz file\n",
    "    \n",
    "    # converts source and targets arrays into a single .npz file\n",
    "\n",
    "    # [src_images, tar_images] = load_images(path_src)\n",
    "    print(\"Loaded: \", src_images.shape, tar_images.shape)\n",
    "    savez_compressed(path_dest, src_images, tar_images)\n",
    "    print('Compressed and saved successfully!')\n",
    "\n",
    "def plot_figures(src_images, tar_images, plot_dest, n_samples=3):\n",
    "\n",
    "    # n_samples : number of samples to be plotted\n",
    "    # plot_dest : path to save the plot\n",
    "\n",
    "    # function to plot figures of dataset\n",
    "    \n",
    "    # [src_images, tar_images] = load_images(path_src)\n",
    "\n",
    "    # Plots first n_samples number of images\n",
    "    for i in range(n_samples):\n",
    "        plt.subplot(2, n_samples, 1+i)\n",
    "        plt.axis('off')\n",
    "        plt.imshow(src_images[i].astype('uint8'))\n",
    "    \n",
    "    for i in range(n_samples):\n",
    "        plt.subplot(2, n_samples, 1+n_samples+i)\n",
    "        plt.axis('off')\n",
    "        plt.imshow(tar_images[i].astype('uint8'))\n",
    "        \n",
    "    plt.savefig(plot_dest)\n",
    "\n",
    "    print('Images saved successfully')\n",
    "\n",
    "def convert_to_hdf5(src_images, tar_images, path_dest):\n",
    "    \n",
    "    # src_images : source image array\n",
    "    # tar_images : target image array\n",
    "    # path_dest : path to store the compressed .hdf5 file\n",
    "\n",
    "    # function to compress arrays to hdf5 file format\n",
    "\n",
    "    # [src_images, tar_images] = load_images(path_src)\n",
    "    print('Loaded: ', src_images.shape, tar_images.shape)\n",
    "\n",
    "    h5f = h5py.File(path_dest, 'w')\n",
    "    h5f.create_dataset('dataset_1', data=[src_images, tar_images])\n",
    "    h5f.close()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# USAGE\n",
    "\n",
    "path_src = './assets/datasets/maps/train/'\n",
    "path_dest = 'file.npz'\n",
    "plot_dest = 'img.png'\n",
    "\n",
    "[src_images, tar_images] = load_images(path_src)\n",
    "compress_images(src_images, tar_images, path_dest)\n",
    "plot_figures(src_images, tar_images, plot_dest)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def encoder(ip_layer, n_filters, batchNorm=True):\n",
    "    \n",
    "    # batchNorm : boolean value that determines whether to apply batch normalization\n",
    "    # ip_layer : input feature vector\n",
    "    # n_filters : number of filters after concatenation\n",
    "\n",
    "    # Encoder architecture : C64-C128-C256-C512-C512-C512-C512-C512\n",
    "    # batchNorm is only applied to first layer of the encoder i.e. C64\n",
    "\n",
    "    # Weights initialized from a Gaussian distribution with mean 0 and standard deviation 0.02\n",
    "    initialized_weights = RandomNormal(stddev=0.02)\n",
    "\n",
    "    # All layers have kernel size 4*4 and stride 2*2\n",
    "    kernel_size = (4, 4)\n",
    "    stride = (2, 2)\n",
    "\n",
    "    layer = Conv2D(n_filters, kernel_size, strides=stride, padding='same', kernel_initializer=initialized_weights)(ip_layer)\n",
    "\n",
    "    if batchNorm:\n",
    "        layer = BatchNormalization()(layer, training=True)\n",
    "\n",
    "    layer = LeakyReLU(alpha=0.2)(layer)\n",
    "\n",
    "    return layer\n",
    "\n",
    "def decoder(ip_layer, skip_connection, n_filters, dropout=True):\n",
    "\n",
    "    # ip_layer : input feature vector\n",
    "    # n_filters : number of filters after concatenation\n",
    "    # dropout : boolean value that determines whether to apply dropout to the layer\n",
    "    # skip_connection : the encoder layer from which skip connections will be applied\n",
    "\n",
    "    # Decoder architecture : C512-C512-C512-C512-C256-C128-C64\n",
    "\n",
    "    # Weights initialized from a Gaussian distribution with mean 0 and standard deviation 0.02\n",
    "    initialized_weights = RandomNormal(stddev=0.02)\n",
    "\n",
    "    # All layers have kernel size 4*4 and stride 2*2\n",
    "    kernel_size = (4, 4)\n",
    "    stride = (2, 2)\n",
    "\n",
    "    layer = Conv2DTranspose(n_filters, kernel_size, strides=stride, padding='same', kernel_initializer=initialized_weights)(ip_layer)\n",
    "    layer = BatchNormalization()(layer, training=True)\n",
    "    if dropout:\n",
    "        layer = Dropout(0.5)(layer, training=True)\n",
    "\n",
    "    # Applying the skip connection\n",
    "    layer = Concatenate()([layer, skip_connection])\n",
    "    # ReLUs in the decoder are not leaky - Ref. 6.1.1.\n",
    "    layer = Activation('relu')(layer)\n",
    "\n",
    "    return layer\n",
    "\n",
    "def generator(image_size = (256, 256, 3)):\n",
    "\n",
    "    # image_size : shape of input image\n",
    "\n",
    "    # Generator Architecture : CCD512-CD1024-CD1024-CD1024-CD1024-CD512-CD256-CD128\n",
    "    # U-Net architecture with skip connections between each layer i in the encoder and the layer n-i in the decoder\n",
    "\n",
    "    # Weights initialized from a Gaussian distribution with mean 0 and standard deviation 0.02\n",
    "    kernel_init = RandomNormal(stddev = 0.02)\n",
    "\n",
    "    ip_image = Input(shape = image_size)\n",
    "    \n",
    "    # All layers have kernel size 4*4 and stride 2*2\n",
    "    kernel_size = (4,4)\n",
    "    stride = (2,2)\n",
    "\n",
    "    e1 = encoder(ip_image, 64, batchNorm=False) # No batchNorm for first encoder layer\n",
    "    e2 = encoder(e1, 128)\n",
    "    e3 = encoder(e2, 256)\n",
    "    e4 = encoder(e3, 512)\n",
    "    e5 = encoder(e4, 512)\n",
    "    e6 = encoder(e5, 512)\n",
    "    e7 = encoder(e6, 512)\n",
    "\n",
    "    # Bottleneck layer, connecting encoder and decoder\n",
    "    bottle_neck = Conv2D(512, kernel_size, strides=stride, padding='same', kernel_initializer=kernel_init, activation='relu')(e7)\n",
    "\n",
    "    d1 = decoder(bottle_neck, e7, 512)\n",
    "    d2 = decoder(d1, e6, 512)\n",
    "    d3 = decoder(d2, e5, 512)\n",
    "    d4 = decoder(d3, e4, 512, dropout=False)\n",
    "    d5 = decoder(d4, e3, 256, dropout=False)\n",
    "    d6 = decoder(d5, e2, 128, dropout=False)\n",
    "    d7 = decoder(d6, e1, 64, dropout=False)\n",
    "\n",
    "    # Convolution is applied to map to the number of output channels, followed by a tanh function - Ref. 6.1.1.\n",
    "    op_image = Conv2DTranspose(3, kernel_size, strides=stride, padding='same', kernel_initializer=kernel_init, activation='tanh')(d7)\n",
    "\n",
    "    # compile model\n",
    "    model = Model(ip_image, op_image)\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = generator()\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Discriminator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def discriminator(image_shape=(256, 256, 3)):\n",
    "\n",
    "    # image_shape : specifies the shape of the input image. By default- 256*256\n",
    "    \n",
    "    # Discriminator architecture details:\n",
    "        # Ck denotes a Convulution-BatchNorm-LeakyReLU layer with k filters (Ref 6.1.)\n",
    "        # Discriminator architecture: C64-C128-C256-C512 (Ref 6.1.2.)\n",
    "        # Receptive field size for above architecture : 70*70\n",
    "\n",
    "        # The depth of the discriminator architecture is responsible for the receptive field size\n",
    "\n",
    "    # Weights initialized from a Gaussian distribution with mean 0 and standard deviation 0.02 (Ref 6.2.)\n",
    "    curv_val = RandomNormal(stddev=0.02)\n",
    "\n",
    "    # Discriminator uses two inputs, the source input image and the target input image\n",
    "    # source image input\n",
    "    # For maps and cityscapes dataset, input image shape is 256*256\n",
    "    src_image_inp = Input(shape=image_shape)\n",
    "    # target image input\n",
    "    target_image_inp = Input(shape=image_shape)\n",
    "\n",
    "    # concatenate images channel-wise\n",
    "    merged_input = Concatenate()([src_image_inp, target_image_inp])\n",
    "\n",
    "    kernel_size = (4,4)\n",
    "    stride = (2,2)\n",
    "\n",
    "    # First layer : C64\n",
    "    layer = Conv2D( 64, kernel_size, strides = stride, padding ='same', kernel_initializer=curv_val)(merged_input)\n",
    "    # No batch-normalization for the first layer\n",
    "    layer = LeakyReLU(alpha=0.2)(layer)\n",
    "\n",
    "    # Subsequent layers : C128-C256-C512\n",
    "    filter_size = [128, 256, 512]\n",
    "    for i in range(len(filter_size)):\n",
    "        layer = Conv2D( filter_size[i], kernel_size, strides = stride, padding ='same', kernel_initializer=curv_val)(layer)\n",
    "        layer = BatchNormalization()(layer)\n",
    "        # Applying Leaky ReLU activation with slope 0.2\n",
    "        layer = LeakyReLU(alpha=0.2)(layer)\n",
    "        \n",
    "    layer = Conv2D(512, (4,4), padding='same', kernel_initializer=curv_val)(layer)\n",
    "    layer = BatchNormalization()(layer)\n",
    "    layer= LeakyReLU(alpha=0.2)(layer)\n",
    "\n",
    "    # After last layer a convolution is applied to map to a 1-dimensional output\n",
    "    layer = Conv2D(1, (4,4), padding='same', kernel_initializer=curv_val)(layer)\n",
    "    patch_out = Activation('sigmoid')(layer)\n",
    "\n",
    "    # define model\n",
    "    model = Model([src_image_inp, target_image_inp], patch_out)\n",
    "    # compile model\n",
    "    opt = Adam(lr=0.0002, beta_1=0.5)\n",
    "    model.compile(loss='binary_crossentropy', optimizer=opt, loss_weights=[0.5])\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = discriminator()\n",
    "print(model.summary())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Composite Model - GAN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def GAN(g_model, d_model, input_shape=(256, 256, 3)):\n",
    "\n",
    "    # g_model : input generator model\n",
    "    # d_model : input discriminator model\n",
    "    # input_shape : input image shape\n",
    "    \n",
    "    # This is a composite model to connect the output of the generator model to the input of the discrimintor model\n",
    "\n",
    "    ip_image = Input(shape=input_shape)\n",
    "\n",
    "    # Discriminator trainable attribute set to false to train the generator.\n",
    "    # Discriminator state remains the same while the generator trains. \n",
    "    d_model.trainable = False\n",
    "    g_output = g_model(ip_image)\n",
    "    d_output = d_model([ip_image, g_output])\n",
    "\n",
    "    GAN_model = Model(ip_image, [d_output, g_output])\n",
    "    # Defining loss function\n",
    "    # loss = Adversarial_loss + lambda * L1_loss\n",
    "    loss = ['binary_crossentropy', 'mae']\n",
    "    # weighting the loss contributions of the different model outputs\n",
    "    loss_weights = [1, 100]\n",
    "\n",
    "    # optimizer parameters specified in  3.3.\n",
    "    optimizer = Adam(lr=0.0002, beta_1=0.5, beta_2=0.999)\n",
    "    GAN_model.compile(loss=loss, optimizer=optimizer, loss_weights=loss_weights)\n",
    "\n",
    "    return GAN_model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g_model = generator()\n",
    "d_model = discriminator()\n",
    "GAN_model = GAN(g_model, d_model)\n",
    "GAN_model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Utility functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_random_training_samples(data, n_samples, n_patch):\n",
    "\n",
    "    # data : input dataset in .npz format\n",
    "    # n_samples : number of samples required\n",
    "    # n_patch : output feature map size (16*16 in our case)\n",
    "\n",
    "    # this function generates a batch of random samples and returns source images and target\n",
    "\n",
    "    train_A, train_B = data\n",
    "    n = randint(0, train_A.shape[0], n_samples)\n",
    "    X1, X2 = train_A[n], train_B[n]\n",
    "\n",
    "    # generate the target array of ones\n",
    "    y = ones((n_samples, n_patch, n_patch, 1))\n",
    "\n",
    "    return [X1, X2], y\n",
    "\n",
    "def generate_fake_samples(generator_model, samples, n_patch):\n",
    "   \n",
    "    # generator_model : input the generator model\n",
    "    # input sample for prediction\n",
    "    # n_patch : output feature map size (16*16 in our case)\n",
    "    \n",
    "    #  generates a batch of fake images through the generator model and the associated target\n",
    "\n",
    "    print(samples.shape)\n",
    "    X = generator_model.predict(samples)\n",
    "\n",
    "    # generate the target array of zeros\n",
    "    y = zeros((len(X), n_patch, n_patch, 1))\n",
    "    return X, y\n",
    "\n",
    "def load_real_samples(filename):\n",
    "\n",
    "    # filename : input .npz filename\n",
    "\n",
    "    # function loads and preprocesses image array before training\n",
    "    \n",
    "    # load compressed numpy arrays (.npz)\n",
    "    data = load(filename)\n",
    "    # unpack arrays\n",
    "    X1, X2 = data['arr_0'], data['arr_1']\n",
    "    \n",
    "    # scale from [0,255] to [-1,1]\n",
    "    X1 = (X1 - 127.5) / 127.5\n",
    "    X2 = (X2 - 127.5) / 127.5\n",
    "\n",
    "    return [X1, X2]\n",
    "\n",
    "def save_model(step, g_model, d_model, gan_model, model_dest):\n",
    "\n",
    "    # step : step at which model is being saved\n",
    "    # g_model, d_model, gan_model : models\n",
    "    # model_dest : destination to save the models\n",
    "\n",
    "    # function saves the models at the given step for further training later\n",
    "\n",
    "    filename1 = model_dest + ('model_g_%06d.h5' % (step+1))\n",
    "    g_model.save(filename1)\n",
    " \n",
    "    filename2 = model_dest + ('model_d_%06d.h5' % (step+1))\n",
    "    d_model.save(filename2)\n",
    " \n",
    "    filename3 = model_dest + ('model_gan_%06d.h5' % (step+1))\n",
    "    gan_model.save(filename3)\n",
    "    print('Models successfully saved at step: %d' % (step))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(discriminator_model, generator_model, gan_model, data, model_dest, n_epochs=200, n_batch=1, n_patch=16, random_jitter=False, current_step=0):\n",
    "\n",
    "    # discriminator_model : input discriminator model\n",
    "    # generator_model : input generator model\n",
    "    # gan_model : input composite gan model\n",
    "    # data : input dataset as an array of images\n",
    "    # model_dest : destination for saving model\n",
    "    # n_epochs : number of epochs\n",
    "    # n_batch : batch size\n",
    "    # n_patch : output feature map size\n",
    "    # random_jitter : boolean value that determines whether to apply random jitter to an image before training\n",
    "    # current_step : in case of resuming training from a checkpoint, current_step indicates the point from where to restart the training\n",
    "\n",
    "    train_A, train_B = data\n",
    "\n",
    "    # calculating total number of steps required in training\n",
    "    batches_per_epoch = int((len(train_A)) / n_batch)\n",
    "    n_steps = batches_per_epoch*n_epochs\n",
    "\n",
    "    print(n_steps, batches_per_epoch)\n",
    "\n",
    "    # Looping over all the steps\n",
    "    for i in range(current_step, n_steps):\n",
    "\n",
    "        # Get a batch of real images\n",
    "        [X_real_A, X_real_B], y_real = generate_random_training_samples(data, n_batch, n_patch)\n",
    "\n",
    "        # Adding random jitter\n",
    "        if random_jitter==True:\n",
    "\n",
    "            # Upsample input images from 256*256 to 286*286\n",
    "            input_image = tf.image.resize(X_real_A, [286, 286], method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)\n",
    "            real_image = tf.image.resize(X_real_B, [286, 286], method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)\n",
    "\n",
    "            stacked_image = tf.stack([input_image, real_image], axis=0)\n",
    "\n",
    "            # Randomly crop the images back to 256*256\n",
    "            cropped_image = tf.image.random_crop( stacked_image, size=[2, 1, 256, 256, 3])\n",
    "\n",
    "            X_real_A, X_real_B = cropped_image[0], cropped_image[1]\n",
    "            \n",
    "            # convert from tensor to numpy\n",
    "            X_real_A = keras.backend.eval(X_real_A)\n",
    "            X_real_B = keras.backend.eval(X_real_B)\n",
    "\n",
    "        # Generate a batch of fake images\n",
    "        X_fake, y_fake = generate_fake_samples(generator_model, X_real_A, n_patch)\n",
    "\n",
    "        # Calculate the discriminator losses\n",
    "        discriminator_loss_real = discriminator_model.train_on_batch([X_real_A, X_real_B], y_real)\n",
    "        discriminator_loss_generated = discriminator_model.train_on_batch([X_real_A, X_fake], y_fake)\n",
    "\n",
    "        # Calculate the generator loss\n",
    "        generator_loss, a, b = gan_model.train_on_batch(X_real_A, [y_real, X_real_B])\n",
    "\n",
    "        print('%d, d1[%.3f] d2[%.3f] g[%.3f]' % (i+1, discriminator_loss_real, discriminator_loss_generated, generator_loss))\n",
    "\n",
    "        # Save model state every 10 epochs\n",
    "        if (i+1) % (batches_per_epoch * 10) == 0:\n",
    "            save_model(i, generator_model, discriminator_model, gan_model, model_dest)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def start_training(dataset_url, model_dest):\n",
    "\n",
    "    # dataset_url : path to compressed dataset\n",
    "    # model_dest : destination path to save models\n",
    "\n",
    "    # this function loads the dataset and starts the training\n",
    "    \n",
    "    train_generator = load_real_samples(dataset_url)\n",
    "    print('Dataset Loaded', train_generator[0].shape, train_generator[1].shape)\n",
    "    # define input shape based on the loaded dataset\n",
    "    image_shape = train_generator[0].shape[1:]\n",
    "    # define the models\n",
    "    d_model = discriminator(image_shape)\n",
    "    g_model = generator(image_shape)\n",
    "    # define the composite GAN model\n",
    "    gan_model = GAN(g_model, d_model, image_shape)\n",
    "    # train model\n",
    "    train(d_model, g_model, gan_model, train_generator, model_dest) \n",
    "\n",
    "def resume_training(step, dataset_url, d_model_src, g_model_src, gan_model_src, model_dest):\n",
    "\n",
    "    # step : step from which training has to be resumed\n",
    "    # d_model_src, g_model_src, gan_model_src : path where models have been saved\n",
    "\n",
    "    # this function resumes the training from the mentioned step and the already saved models\n",
    "\n",
    "    d_model = load_model(d_model_src)\n",
    "    g_model = load_model(g_model_src)\n",
    "    gan_model = load_model(gan_model_src)\n",
    "\n",
    "    dataset = load_real_samples(dataset_url)\n",
    "    train(d_model, g_model, gan_model, dataset, model_dest, current_step=step)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# USAGE\n",
    "\n",
    "# Call to start training\n",
    "dataset_url = './assets/datasets/maps/compressed/maps_256.npz'\n",
    "model_dest = './models/maps/'\n",
    "start_training(dataset_url, model_dest)\n",
    "\n",
    "# Call to resume training from step 153400 ( 140 epochs on the maps dataset)\n",
    "step = 153440\n",
    "d_model_src = model_dest + 'model_d_' + str(step) + '.h5'\n",
    "g_model_src = model_dest + 'model_g_' + str(step) + '.h5'\n",
    "gan_model_src = model_dest + 'model_gan_' + str(step) + '.h5'\n",
    "resume_training(step, dataset_url, d_model_src, g_model_src, gan_model_src, model_dest)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_images(src_img, gen_img, tar_img, dest):\n",
    "\n",
    "    # Plots all the input images in the destination specified by dest\n",
    "\n",
    "    images = vstack((src_img, gen_img, tar_img))\n",
    "    # scaling from [-1,1] to [0,1]\n",
    "    images = (images + 1) / 2.0\n",
    "    titles = ['Source', 'Generated', 'Expected']\n",
    "\n",
    "    fig = plt.figure(figsize=(10, 20))\n",
    "\n",
    "    for i in range(len(images)):\n",
    "        plt.subplot(1, 3, 1 + i)\n",
    "        plt.axis('off')\n",
    "        plt.imshow(images[i])\n",
    "        plt.title(titles[i])\n",
    "    \n",
    "    print('Figure saved successfully.')\n",
    "    plt.savefig(dest)\n",
    " \n",
    "def generate_prediction(models, dataset, dest):\n",
    "\n",
    "    # models : input an array of models for prediction\n",
    "    # dataset : path to compressed dataset (.npz format)\n",
    "    # dest : destination path to store the plot\n",
    "\n",
    "    [X1, X2] = dataset\n",
    "    # selects a random image from the dataset\n",
    "    ix = randint(0, len(X1), 1)\n",
    "\n",
    "    dest = dest + 'train_' + str(ix) + '.jpg'\n",
    "\n",
    "    # generate an image for every model\n",
    "    for model in models:\n",
    "        src_image, tar_image = X1[ix], X2[ix]\n",
    "        gen_image = model.predict(src_image)\n",
    "        plot_images(src_image, gen_image, tar_image, dest)\n",
    "\n",
    "def prediction(model, img_url, size=(256, 512)):\n",
    "\n",
    "    # model : accepts a model\n",
    "    # img_url : input path to image\n",
    "    # size : size of input image\n",
    "\n",
    "    # load image and convert to numpy array\n",
    "    pix = load_img(img_url, target_size=size)\n",
    "    pix = img_to_array(pix)\n",
    "\n",
    "    # split the image into source image and target image \n",
    "    s_img, t_img = pix[:, :256], pix[:, 256:]\n",
    "    # scale from [0,255] to [-1,1]\n",
    "    s = (s_img-127.5) / 127.5\n",
    "    s = expand_dims(s, 0)\n",
    "\n",
    "    # generate an image\n",
    "    gen = model.predict(s)\n",
    "    gen = (gen + 1) / 2\n",
    "    \n",
    "    # returns generated image, target image\n",
    "    return gen[0], t_img\n",
    "\n",
    "def load_and_plot(img_url, models, dest, size=(256, 512)):\n",
    "\n",
    "    # img_url : path to image\n",
    "    # models : an array of models for image generation\n",
    "    # dest : destination path to save images\n",
    "    # size : size of input images\n",
    "\n",
    "    # Function used for comparison of images generated by models trained for different number of epochs\n",
    "\n",
    "    # load and convert the image into numpy array\n",
    "    pix = load_img(img_url, target_size=size)\n",
    "    pix = img_to_array(pix)\n",
    "    \n",
    "    s_img, t_img = pix[:, :256], pix[:, 256:]\n",
    "    s = (s_img-127.5) / 127.5\n",
    "    s = expand_dims(sat, 0)\n",
    "\n",
    "    fig = plt.figure(figsize=(10, 20))\n",
    "\n",
    "    for i in range(len(models)):\n",
    "        gen = models[i].predict(s)\n",
    "        gen = (gen + 1) / 2\n",
    "\n",
    "        plt.subplot(len(models), 3, i*3 + 1)\n",
    "        plt.title('Source')\n",
    "        plt.axis('off')\n",
    "        plt.imshow(s_img.astype('uint8'))\n",
    "        plt.subplot(len(models), 3, i*3 + 2)\n",
    "        plt.title('Generated')\n",
    "        plt.imshow(gen[0])\n",
    "        plt.axis('off')\n",
    "        plt.subplot(len(models), 3, i*3 + 3)\n",
    "        plt.title('Expected')\n",
    "        plt.axis('off')\n",
    "        plt.imshow(t_img.astype('uint8'))\n",
    "\n",
    "    print('Figure saved successfully at destination: %s' % dest)\n",
    "\n",
    "    plt.savefig(dest)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# USAGE\n",
    "\n",
    "# Test on maps dataset\n",
    "model_1 = load_model('./models/maps/model_g_153440.h5')\n",
    "model_2 = load_model('./models/maps/model_g_021920.h5')\n",
    "\n",
    "for i in range(10):\n",
    "    generate_random_map_image = randint(1, 1099)\n",
    "    img_url = './assets/datasets/maps/val/' + str(generate_random_map_image) + '.jpg'\n",
    "    dest_val = './assets/plots/maps/maps_' + str(generate_random_map_image) + '_val.jpg'\n",
    "\n",
    "    load_and_plot(img_url, models=[model_2, model_1], dest=dest_val)\n",
    "\n",
    "# To test on cityscapes dataset\n",
    "dest_dir = './assets/plots/maps/'\n",
    "dataset = load_real_samples('./assets/datasets/maps/compressed/maps_256.npz')\n",
    "generate_prediction(models=[model], dataset=dataset, dest=dest_dir)\n",
    "\n",
    "model_1 = load_model('./models/cityscapes/model_g_059500.h5')\n",
    "model_2 = load_model('./models/cityscapes/model_g_297500.h5')\n",
    "\n",
    "for i in range(10):\n",
    "    generate_random_map_image = randint(1, 501)\n",
    "    img_url = './assets/datasets/cityscapes/val/' + str(generate_random_map_image) + '.jpg'\n",
    "    dest_val_1 = './assets/plots/cityscapes/cityscapes_' + str(generate_random_map_image) + '_val.jpg'\n",
    "\n",
    "    load_and_plot(img_url, models=[model_1, model_2], dest=dest_val_1)\n",
    "\n",
    "# To test on trained images\n",
    "dest_dir = './assets/plots/cityscapes/'\n",
    "dataset = load_real_samples('./assets/datasets/cityscapes/compressed/cityscapes_256.npz')\n",
    "generate_prediction(models=[model], dataset=dataset, dest=dest_dir)\n"
   ]
  }
 ],
 "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.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}