Conditional_GAN_.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Conditional GAN .ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Kushagratandon12/Satellite-to-Map-Image-Translation/blob/main/Conditional_GAN_.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pC4tPzyz_sn_"
      },
      "source": [
        "import gc \n",
        "import os\n",
        "import cv2\n",
        "import time\n",
        "import glob\n",
        "import imageio\n",
        "import matplotlib\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "from IPython import display\n",
        "from matplotlib import pyplot as plt\n",
        "from tensorflow.keras.preprocessing import image\n",
        "!rm -rf /content/sample_data"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JzgsF90UCmCV"
      },
      "source": [
        "## Load the Dataset - link"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v527FrrYCKpo"
      },
      "source": [
        "# !rm -rf /content/sample_data\n",
        "# If data is not present\n",
        "%%capture\n",
        "!wget https://people.eecs.berkeley.edu/~tinghuiz/projects/pix2pix/datasets/maps.tar.gz\n",
        "!tar -xvf /content/maps.tar.gz\n",
        "\n",
        "Aug_data_path= '/content/maps'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TaAjNzdICWau"
      },
      "source": [
        "train_path= '/content/maps/train' # Aug_data_path\n",
        "train_imgs = glob.glob('/content/maps/train/*.jpg')\n",
        "print(\"training images\",len(train_imgs))\n",
        "\n",
        "\n",
        "def load(image_file,target_shape=(512,512)):\n",
        "    img = image.load_img(image_file,target_size=target_shape)\n",
        "    img = image.img_to_array(img)\n",
        "    w = img.shape[1]\n",
        "    w = w//2\n",
        "    real_image = img[:, w:, :]\n",
        "    input_image = img[:, :w, :]\n",
        "\n",
        "    return input_image, real_image"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Z2OyVrmFJWTv"
      },
      "source": [
        "inp, re = load(train_imgs[np.random.randint(0,len(train_imgs))])\n",
        "# casting to int for matplotlib to show the image\n",
        "plt.figure(figsize=(20,20))\n",
        "plt.title('Input image')\n",
        "plt.imshow(inp/255.0)\n",
        "# real Images \n",
        "plt.figure(figsize=(20,20))\n",
        "plt.title('Unity Image')\n",
        "plt.imshow(re/255.0)\n",
        "print(inp.shape)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "T4U5ChbtKBHb"
      },
      "source": [
        "def resize(input_image, real_image, height, width):\n",
        "  input_image = tf.image.resize(input_image, [height, width],\n",
        "                                method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)\n",
        "  real_image = tf.image.resize(real_image, [height, width],\n",
        "                               method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)\n",
        "  return input_image, real_image\n",
        "\n",
        "def random_crop(input_image, real_image):\n",
        "  stacked_image = tf.stack([input_image, real_image], axis=0)\n",
        "  cropped_image = tf.image.random_crop(\n",
        "      stacked_image, size=[2, IMG_HEIGHT, IMG_WIDTH, 3])\n",
        "\n",
        "  return cropped_image[0], cropped_image[1]\n",
        "\n",
        "#normalizing the images to [-1, 1]\n",
        "\n",
        "def normalize(input_image, real_image):\n",
        "  input_image = (input_image / 255.0) - 1\n",
        "  real_image = (real_image / 255.0) - 1\n",
        "\n",
        "  return input_image, real_image"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "56KXN1H7Mmee"
      },
      "source": [
        "## Pre-process images\n",
        "1. Resize an image to bigger height and width\n",
        "2. Randomly crop to the target size\n",
        "3. Randomly flip the image horizontally\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eZgqYwSnM19k"
      },
      "source": [
        "# Define Hyper Parameters \n",
        "BUFFER_SIZE = 400\n",
        "BATCH_SIZE = 1\n",
        "IMG_SHAPE = (256,256)\n",
        "\n",
        "def load_image_train(image_file):\n",
        "  input_image, real_image = load(image_file,target_shape=(256,512))\n",
        "  input_image, real_image = normalize(input_image, real_image)\n",
        "\n",
        "  return input_image, real_image\n",
        "\n",
        "def load_image_test(image_file):\n",
        "  input_image, real_image = load(image_file)\n",
        "  input_image, real_image = resize(input_image, real_image,\n",
        "                                   IMG_SHAPE[0], IMG_SHAPE[1])\n",
        "  input_image, real_image = normalize(input_image, real_image)\n",
        "\n",
        "  return input_image, real_image\n",
        "\n",
        "def load_dataset(train_imgs,train_path):\n",
        "    src_imgs=list()\n",
        "    tar_imgs =list()\n",
        "    for imgs in train_imgs:\n",
        "        filename =imgs\n",
        "        in_img , rl_img = load_image_train(filename)\n",
        "        src_imgs.append(in_img)\n",
        "        tar_imgs.append(rl_img)\n",
        "    src_imgs = np.asarray(src_imgs)\n",
        "    tar_imgs = np.asarray(tar_imgs)\n",
        "    return [src_imgs,tar_imgs]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "M2KsEbiUTia3"
      },
      "source": [
        "src_images, tar_images = load_dataset(train_imgs,train_path)\n",
        "print('Loaded: source Image {} \\n Target Image Shape {}'.format( src_images.shape, tar_images.shape))\n",
        "\n",
        "dataset = [src_images,tar_images]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vZozK_I5v-cR"
      },
      "source": [
        "## Build the Discriminator\n",
        "\n",
        "1. The Discriminator is a PatchGAN.\n",
        "2. Each block in the discriminator is (Conv -> BatchNorm -> Leaky ReLU)\n",
        "3. The shape of the output after the last layer is (batch_size, 30, 30, 1)\n",
        "4. Each 30x30 patch of the output classifies a 70x70 portion of the input image (such an architecture is called a PatchGAN).\n",
        "5. Discriminator receives 2 inputs.\n",
        "6. Input image and the target image, which it should classify as real.\n",
        "7. Input image and the generated image (output of generator), which it should classify as fake.\n",
        "8. Concatenate these 2 inputs together in the code (tf.concat([inp, tar], axis=-1))\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QOdZgozUdW1f"
      },
      "source": [
        "'''\n",
        "The discriminator is a deep convolutional neural network that performs image classification. \n",
        "Specifically, conditional-image classification. \n",
        "It takes both the source image (e.g. satellite photo) and the target image (e.g. Google maps image) \n",
        "as input and predicts the likelihood of whether target image is real or a fake translation of the source image.\n",
        "'''\n",
        "from tensorflow.keras.layers import Concatenate,Conv2D,BatchNormalization,LeakyReLU,Input,Activation\n",
        "from tensorflow.keras.initializers import RandomNormal\n",
        "from tensorflow.keras import Model\n",
        "from tensorflow.keras.optimizers import Adam\n",
        "gc.collect()\n",
        "def define_discriminator(image_shape):\n",
        "\t# weight initialization\n",
        "\tinit = RandomNormal(stddev=0.02)\n",
        "\t# source image input\n",
        "\tin_src_image = Input(shape=image_shape)\n",
        "\t# target image input\n",
        "\tin_target_image = Input(shape=image_shape)\n",
        "\t# concatenate images channel-wise\n",
        "\tmerged = Concatenate()([in_src_image, in_target_image])\n",
        "\t# C64\n",
        "\td = Conv2D(64, (4,4), strides=(2,2), padding='same', kernel_initializer=init)(merged)\n",
        "\td = LeakyReLU(alpha=0.2)(d)\n",
        "\t# C128\n",
        "\td = Conv2D(128, (4,4), strides=(2,2), padding='same', kernel_initializer=init)(d)\n",
        "\td = BatchNormalization()(d)\n",
        "\td = LeakyReLU(alpha=0.2)(d)\n",
        "\t# C256\n",
        "\td = Conv2D(256, (4,4), strides=(2,2), padding='same', kernel_initializer=init)(d)\n",
        "\td = BatchNormalization()(d)\n",
        "\td = LeakyReLU(alpha=0.2)(d)\n",
        "\t# C512\n",
        "\td = Conv2D(512, (4,4), strides=(2,2), padding='same', kernel_initializer=init)(d)\n",
        "\td = BatchNormalization()(d)\n",
        "\td = LeakyReLU(alpha=0.2)(d)\n",
        "\t# second last output layer\n",
        "\td = Conv2D(512, (4,4), padding='same', kernel_initializer=init)(d)\n",
        "\td = BatchNormalization()(d)\n",
        "\td = LeakyReLU(alpha=0.2)(d)\n",
        "\t# patch output\n",
        "\td = Conv2D(1, (4,4), padding='same', kernel_initializer=init)(d)\n",
        "\tpatch_out = Activation('sigmoid')(d)\n",
        "\t# define model\n",
        "\tmodel = Model([in_src_image, in_target_image], patch_out)\n",
        "\t# compile model\n",
        "\topt = Adam(lr=0.0002, beta_1=0.5)\n",
        "\tmodel.compile(loss='binary_crossentropy', optimizer=opt, loss_weights=[0.5])\n",
        "\treturn model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7M1eTlR8zxGS"
      },
      "source": [
        "discriminator = define_discriminator(image_shape=[IMG_SHAPE[0],IMG_SHAPE[1],3])\n",
        "# tf.keras.utils.plot_model(discriminator, show_shapes=True, dpi=64)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YhbfOMidgAR7"
      },
      "source": [
        "### Build the Generator\n",
        "\n",
        "1. The architecture of generator is a modified U-Net.\n",
        "2. Each block in the encoder is (Conv -> Batchnorm -> Leaky ReLU)\n",
        "3. Each block in the decoder is (Transposed Conv -> Batchnorm -> Dropout(applied to the first 3 blocks) -> ReLU)\n",
        "4. There are skip connections between the encoder and decoder (as in U-Net).\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "S95FiYzEmwaL"
      },
      "source": [
        "from tensorflow.keras.layers import Conv2DTranspose,Dropout\n",
        "def define_encoder_block(layer_in, n_filters, batchnorm=True):\n",
        "\t# weight initialization\n",
        "\tinit = RandomNormal(stddev=0.02)\n",
        "\t# add downsampling layer\n",
        "\tg = Conv2D(n_filters, (4,4), strides=(2,2), padding='same', kernel_initializer=init)(layer_in)\n",
        "\t# conditionally add batch normalization\n",
        "\tif batchnorm:\n",
        "\t\tg = BatchNormalization()(g, training=True)\n",
        "\t# leaky relu activation\n",
        "\tg = LeakyReLU(alpha=0.2)(g)\n",
        "\treturn g\n",
        " \n",
        "# define a decoder block\n",
        "def decoder_block(layer_in, skip_in, n_filters, dropout=True):\n",
        "\t# weight initialization\n",
        "\tinit = RandomNormal(stddev=0.02)\n",
        "\t# add upsampling layer\n",
        "\tg = Conv2DTranspose(n_filters, (4,4), strides=(2,2), padding='same', kernel_initializer=init)(layer_in)\n",
        "\t# add batch normalization\n",
        "\tg = BatchNormalization()(g, training=True)\n",
        "\t# conditionally add dropout\n",
        "\tif dropout:\n",
        "\t\tg = Dropout(0.5)(g, training=True)\n",
        "\t# merge with skip connection\n",
        "\tg = Concatenate()([g, skip_in])\n",
        "\t# relu activation\n",
        "\tg = Activation('relu')(g)\n",
        "\treturn g\n",
        " \n",
        "# define the standalone generator model\n",
        "def define_generator(image_shape=(256,256,3)):\n",
        "\t# weight initialization\n",
        "\tinit = RandomNormal(stddev=0.02)\n",
        "\t# image input\n",
        "\tin_image = Input(shape=image_shape)\n",
        "\t# encoder model\n",
        "\te1 = define_encoder_block(in_image, 64, batchnorm=False)\n",
        "\te2 = define_encoder_block(e1, 128)\n",
        "\te3 = define_encoder_block(e2, 256)\n",
        "\te4 = define_encoder_block(e3, 512)\n",
        "\te5 = define_encoder_block(e4, 512)\n",
        "\te6 = define_encoder_block(e5, 512)\n",
        "\te7 = define_encoder_block(e6, 512)\n",
        "\t# bottleneck, no batch norm and relu\n",
        "\tb = Conv2D(512, (4,4), strides=(2,2), padding='same', kernel_initializer=init)(e7)\n",
        "\tb = Activation('relu')(b)\n",
        "\t# decoder model\n",
        "\td1 = decoder_block(b, e7, 512)\n",
        "\td2 = decoder_block(d1, e6, 512)\n",
        "\td3 = decoder_block(d2, e5, 512)\n",
        "\td4 = decoder_block(d3, e4, 512, dropout=False)\n",
        "\td5 = decoder_block(d4, e3, 256, dropout=False)\n",
        "\td6 = decoder_block(d5, e2, 128, dropout=False)\n",
        "\td7 = decoder_block(d6, e1, 64, dropout=False)\n",
        "\t# output\n",
        "\tg = Conv2DTranspose(3, (4,4), strides=(2,2), padding='same', kernel_initializer=init)(d7)\n",
        "\tout_image = Activation('tanh')(g)\n",
        "\t# define model\n",
        "\tmodel = Model(in_image, out_image)\n",
        "\treturn model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "damwE23Epdih"
      },
      "source": [
        "generator = define_generator(image_shape=(IMG_SHAPE[0],IMG_SHAPE[1],3))\n",
        "# tf.keras.utils.plot_model(generator, show_shapes=True, dpi=64)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X1lDLgjXOtKy"
      },
      "source": [
        "## Define Cyclic GAN Network : Combined Generator and Discriminator models"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Au1LozRS54zq"
      },
      "source": [
        "# define the combined generator and discriminator model, for updating the generator\n",
        "def define_gan(g_model, d_model, image_shape):\n",
        "\t# make weights in the discriminator not trainable\n",
        "\tfor layer in d_model.layers:\n",
        "\t\tif not isinstance(layer, BatchNormalization):\n",
        "\t\t\tlayer.trainable = False\n",
        "\t# define the source image\n",
        "\tin_src = Input(shape=image_shape)\n",
        "\t# connect the source image to the generator input\n",
        "\tgen_out = g_model(in_src)\n",
        "\t# connect the source input and generator output to the discriminator input\n",
        "\tdis_out = d_model([in_src, gen_out])\n",
        "\t# src image as input, generated image and classification output\n",
        "\tmodel = Model(in_src, [dis_out, gen_out])\n",
        "\t# compile model\n",
        "\topt = Adam(lr=0.0002, beta_1=0.5)\n",
        "\tmodel.compile(loss=['binary_crossentropy', 'mae'], optimizer=opt, loss_weights=[1,100])\n",
        "\treturn model\n",
        "\n",
        "gan_model = define_gan(generator, discriminator, (IMG_SHAPE[0],IMG_SHAPE[1],3))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M4ufnDCdP_zw"
      },
      "source": [
        "## Define Optimizers and Checkpoints for generator and Discriminator"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2A2mwBbvP-tF"
      },
      "source": [
        "generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n",
        "discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)\n",
        "\n",
        "checkpoint_dir = './training_checkpoints'\n",
        "checkpoint_prefix = os.path.join(checkpoint_dir, \"ckpt\")\n",
        "checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer,\n",
        "                                 discriminator_optimizer=discriminator_optimizer,\n",
        "                                 generator=generator,\n",
        "                                 discriminator=discriminator)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MEZqsCv2RGvn"
      },
      "source": [
        "## Generate functions for Fake and Real images "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EgeMvH1UREXq"
      },
      "source": [
        "def generate_real_samples(dataset, n_samples, patch_shape):\n",
        "    # unpack dataset\n",
        "    trainA, trainB = dataset\n",
        "    # choose random instances\n",
        "    ix = np.random.randint(0, trainA.shape[0], n_samples)\n",
        "    # retrieve selected images\n",
        "    X1, X2 = trainA[ix], trainB[ix]\n",
        "    # generate 'real' class labels (1)\n",
        "    y = np.ones((n_samples, patch_shape, patch_shape, 1))\n",
        "    return [X1, X2], y\n",
        "\n",
        "# generate a batch of images, returns images and targets\n",
        "def generate_fake_samples(g_model, samples, patch_shape):\n",
        "    # generate fake instance\n",
        "    X = g_model.predict(samples)\n",
        "    # create 'fake' class labels (0)\n",
        "    y = np.zeros((len(X), patch_shape, patch_shape, 1))\n",
        "    return X, y"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D10rBy8HROgU"
      },
      "source": [
        "### Define Training Function for Cyclic GAN Networks"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TBgL6RKtB3_w"
      },
      "source": [
        "def train(d_model, g_model, gan_model, dataset, n_epochs=200, n_batch=1):\n",
        "    # determine the output square shape of the discriminator\n",
        "    n_patch = d_model.output_shape[1]\n",
        "    # unpack dataset\n",
        "    trainA, trainB = dataset\n",
        "    # calculate the number of batches per training epoch\n",
        "    bat_per_epo = int(len(trainA) / n_batch)\n",
        "    # calculate the number of training iterations\n",
        "    n_steps = bat_per_epo * n_epochs\n",
        "    print(\"Running {} steps \".format(n_steps))\n",
        "    # manually enumerate epochs\n",
        "    for i in range(n_steps):\n",
        "        # select a batch of real samples\n",
        "        [X_realA, X_realB], y_real = generate_real_samples(dataset, n_batch, n_patch)\n",
        "        # generate a batch of fake samples\n",
        "        X_fakeB, y_fake = generate_fake_samples(g_model, X_realA, n_patch)\n",
        "   \n",
        "        # update discriminator for real samples\n",
        "        d_loss1 = d_model.train_on_batch([X_realA, X_realB], y_real)\n",
        "        # update discriminator for generated samples\n",
        "        d_loss2 = d_model.train_on_batch([X_realA, X_fakeB], y_fake)\n",
        "        # update the generator\n",
        "        g_loss, _, _ = gan_model.train_on_batch(X_realA, [y_real, X_realB])\n",
        "        # summarize performance and generate images \n",
        "        if i%10 == 0:\n",
        "            #checkpoint.save(file_prefix=checkpoint_prefix)\n",
        "            print(\"Step {} Disc loss real sample {:.3f} Disc Loss predicted {:.3f} generator loss {:.3f}\".format(i,d_loss1,d_loss2,g_loss))\n",
        "            #print('>%d, d1[%.3f] d2[%.3f] g[%.3f]' % (i+1, d_loss1, d_loss2, g_loss))\n",
        "            generate_images(g_model, X_realA, X_realB,iter=i,save_img='/content/sample_imgs')\n",
        "        # summarize model performance\n",
        "        if (i+1) % (bat_per_epo * 10) == 0:\n",
        "            summarize_performance(i, g_model, dataset)\n",
        "\n",
        "def generate_images(model, test_input,tar,iter=0,save_img=None):\n",
        "    prediction = model(test_input, training=True)\n",
        "    display_list = [test_input[0], tar[0], prediction[0]]\n",
        "    title = ['Input Image', 'Ground Truth', 'Predicted Image']\n",
        "\n",
        "    in_img = np.asarray(test_input[0] * 0.5 + 0.5)\n",
        "    gt_img = np.asarray(tar[0]*0.5+0.5)\n",
        "    pd_img = np.asarray(prediction[0]*0.5+0.5)\n",
        "\n",
        "    final_img = cv2.hconcat([in_img,gt_img,pd_img])\n",
        "    final_img = final_img.astype(np.uint8)\n",
        "\n",
        "    if save_img:\n",
        "        imageio.imwrite(os.path.join(save_img,'image_'+str(iter)+'.png'),pd_img)\n",
        "    \n",
        "    for i in range(3):\n",
        "        fig = plt.figure(figsize=(20, 20))\n",
        "        plt.subplot(1, 3, i+1)\n",
        "        plt.title(title[i])\n",
        "        # getting the pixel values between [0, 1] to plot it.\n",
        "        plt.imshow(display_list[i] * 0.5 + 0.5)\n",
        "        plt.axis('off')\n",
        "        plt.show()\n",
        "        img = np.asarray(display_list[2] * 0.5 + 0.5)\n",
        "        img = img.astype(np.uint8)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qksi15N5Tcjz"
      },
      "source": [
        "## Model evaluation "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "N0FTzVLa1XbP"
      },
      "source": [
        "# generate samples and save as a plot and save the model\n",
        "def summarize_performance(step, g_model, dataset, n_samples=3):\n",
        "\t# select a sample of input images\n",
        "\t[X_realA, X_realB], _ = generate_real_samples(dataset, n_samples, 1)\n",
        "\t# generate a batch of fake samples\n",
        "\tX_fakeB, _ = generate_fake_samples(g_model, X_realA, 1)\n",
        "\t# scale all pixels from [-1,1] to [0,1]\n",
        "\tX_realA = (X_realA + 1) / 2.0\n",
        "\tX_realB = (X_realB + 1) / 2.0\n",
        "\tX_fakeB = (X_fakeB + 1) / 2.0\n",
        "\t# plot real source images\n",
        "\tfor i in range(n_samples):\n",
        "\t\tplt.subplot(3, n_samples, 1 + i)\n",
        "\t\tplt.axis('off')\n",
        "\t\tplt.imshow(X_realA[i])\n",
        "\t# plot generated target image\n",
        "\tfor i in range(n_samples):\n",
        "\t\tplt.subplot(3, n_samples, 1 + n_samples + i)\n",
        "\t\tplt.axis('off')\n",
        "\t\tplt.imshow(X_fakeB[i])\n",
        "\t# plot real target image\n",
        "\tfor i in range(n_samples):\n",
        "\t\tplt.subplot(3, n_samples, 1 + n_samples*2 + i)\n",
        "\t\tplt.axis('off')\n",
        "\t\tplt.imshow(X_realB[i])\n",
        "\t# save plot to file\n",
        "\tfilename1 = 'plot_%06d.png' % (step+1)\n",
        "\tplt.savefig(filename1)\n",
        "\tplt.close()\n",
        "\t# save the generator model\n",
        "\tfilename2 = 'model_%06d.h5' % (step+1)\n",
        "\tg_model.save(filename2)\n",
        "\tprint('>Saved: %s and %s' % (filename1, filename2))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V0Ui2LdLVp5s"
      },
      "source": [
        "## Summary writer to view logs "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-ussBHOEClqt"
      },
      "source": [
        "# remove the checkpoints if created previously \n",
        "! rm -rf '/content/training_checkpoints'\n",
        "! rm -rf '/content/logs'\n",
        "\n",
        "print('Loaded training Data {} Target Data{}'.format(dataset[0].shape,dataset[1].shape))\n",
        "# train model\n",
        "train(discriminator, generator, gan_model, dataset,n_batch=10)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lttq2Epen0n8"
      },
      "source": [
        "sample_images = glob.glob('/content/sample_imgs/*.png')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}