GAN.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "GAN.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "aDQXAOHS8LZh",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "import sys\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import matplotlib.pyplot as plt\n",
        "from datetime import datetime\n",
        "from tensorflow.keras.layers import (\n",
        "    Activation, AveragePooling2D, BatchNormalization, Conv2D, Conv2DTranspose,\n",
        "    Dense, Dropout, Flatten, Input, LeakyReLU, ReLU, UpSampling2D)\n",
        "from tensorflow.keras.datasets import cifar10\n",
        "from tensorflow.keras.models import Sequential, Model\n",
        "from time import time\n",
        "from skimage.color import rgb2lab, lab2rgb\n",
        "\n",
        "\n",
        "IMAGE_SIZE = 32\n",
        "EPOCHS = 150\n",
        "BATCH_SIZE = 128\n",
        "SHUFFLE_BUFFER_SIZE = 100\n",
        "\n",
        "WORKDIR = \"/content/drive/My Drive/Colab Notebooks/Workspace\""
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fv2attTI8OWz",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "e0f3c53c-192a-42f5-87f5-b371b3793192"
      },
      "source": [
        "def generate_dataset(images, debug=False):\n",
        "    X = []\n",
        "    Y = []\n",
        "\n",
        "    for i in images:\n",
        "        lab_image_array = rgb2lab(i / 255)\n",
        "        x = lab_image_array[:, :, 0]\n",
        "        y = lab_image_array[:, :, 1:]\n",
        "        y /= 128  # normalize\n",
        "\n",
        "        if debug:\n",
        "            fig = plt.figure()\n",
        "            fig.add_subplot(1, 2, 1)\n",
        "            plt.imshow(i / 255)\n",
        "\n",
        "            fig.add_subplot(1, 2, 2)\n",
        "            plt.imshow(lab2rgb(np.dstack((x, y * 128))))\n",
        "            plt.show()\n",
        "\n",
        "        X.append(x.reshape(IMAGE_SIZE, IMAGE_SIZE, 1))\n",
        "        Y.append(y)\n",
        "\n",
        "    X = np.array(X, dtype=np.float32)\n",
        "    Y = np.array(Y, dtype=np.float32)\n",
        "\n",
        "    return X, Y\n",
        "\n",
        "\n",
        "def load_data(force=False):\n",
        "    is_saved_arrays_exist = os.path.isfile(os.path.join(WORKDIR, 'X_train.npy'))\n",
        "\n",
        "    if not is_saved_arrays_exist or force:\n",
        "        (train_images, _), (test_images, _) = cifar10.load_data()\n",
        "        X_train, Y_train = generate_dataset(train_images)\n",
        "        X_test, Y_test = generate_dataset(test_images)\n",
        "        print('Saving processed data to Drive')\n",
        "        np.save(os.path.join(WORKDIR, 'X_train.npy'), X_train)\n",
        "        np.save(os.path.join(WORKDIR, 'Y_train.npy'), Y_train)\n",
        "        np.save(os.path.join(WORKDIR, 'X_test.npy'), X_test)\n",
        "        np.save(os.path.join(WORKDIR, 'Y_test.npy'), Y_test)\n",
        "    else:\n",
        "        print('Loading processed data from Drive')\n",
        "        X_train = np.load(os.path.join(WORKDIR, 'X_train.npy'))\n",
        "        Y_train = np.load(os.path.join(WORKDIR, 'Y_train.npy'))\n",
        "        X_test = np.load(os.path.join(WORKDIR, 'X_test.npy'))\n",
        "        Y_test = np.load(os.path.join(WORKDIR, 'Y_test.npy'))\n",
        "        \n",
        "    \n",
        "    return X_train, Y_train, X_test, Y_test\n",
        "\n",
        "\n",
        "X_train, Y_train, X_test, Y_test = load_data()\n",
        "\n",
        "train_dataset = tf.data.Dataset.from_tensor_slices((X_train, Y_train))\n",
        "test_dataset = tf.data.Dataset.from_tensor_slices((X_test, Y_test))\n",
        "\n",
        "train_dataset = train_dataset.shuffle(SHUFFLE_BUFFER_SIZE).batch(BATCH_SIZE)\n",
        "test_dataset = test_dataset.batch(BATCH_SIZE)"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Loading processed data from Drive\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cJvTBV-p8feY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def make_generator_model():\n",
        "    model = Sequential()\n",
        "    model.add(Conv2D(32, (3, 3), padding='same', activation='relu', input_shape=(IMAGE_SIZE, IMAGE_SIZE, 1)))\n",
        "    model.add(Conv2D(64, (3, 3), padding='same', strides=2, activation='relu'))\n",
        "    model.add(BatchNormalization())\n",
        "    # model = MaxPooling2D(pool_size=(2, 2))(model)\n",
        "\n",
        "    model.add(Conv2D(128, (3, 3), padding='same', activation='relu', strides=2))\n",
        "    model.add(Conv2D(128, (3, 3), padding='same', activation='relu'))\n",
        "    model.add(BatchNormalization())\n",
        "    # model = MaxPooling2D(pool_size=(2, 2))(model)\n",
        "\n",
        "    model.add(Conv2D(256, (3, 3), padding='same'))\n",
        "    model.add(Activation('relu'))\n",
        "    model.add(BatchNormalization())\n",
        "\n",
        "    model.add(UpSampling2D(size=(2,2)))\n",
        "    model.add(Conv2D(128, (3, 3), padding='same', activation='relu'))\n",
        "    model.add(BatchNormalization())\n",
        "\n",
        "    model.add(UpSampling2D(size=(2,2)))\n",
        "    model.add(Conv2D(64, (3, 3), padding='same', activation='relu'))\n",
        "    model.add(BatchNormalization())\n",
        "\n",
        "    model.add(Conv2D(32, (3, 3), padding='same', activation='relu'))\n",
        "    model.add(Conv2D(2, (3, 3), padding='same'))\n",
        "    model.add(BatchNormalization())\n",
        "    model.add(Activation('tanh'))\n",
        "    # self.model = BatchNormalization()(self.model)\n",
        "    # self.model = merge(inputs=[self.g_input, self.model], mode='concat')\n",
        "    # self.model = Activation('linear')(self.model)\n",
        "    return model"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zChqzKKgmFHw",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def downsample(filters, kernel_size, apply_batchnorm=True):\n",
        "    initializer = tf.random_uniform_initializer(0, 0.02)\n",
        "    model = Sequential()\n",
        "    model.add(Conv2D(filters, kernel_size, strides=2, padding='same',\n",
        "                     kernel_initializer=initializer, use_bias=False))\n",
        "    \n",
        "    if apply_batchnorm:\n",
        "        model.add(BatchNormalization())\n",
        "\n",
        "    model.add(LeakyReLU())\n",
        "    return model\n",
        "\n",
        "\n",
        "def upsample(filters, kernel_size, apply_dropout=False):\n",
        "    initializer = tf.random_uniform_initializer(0, 0.02)\n",
        "    model = Sequential()\n",
        "    model.add(Conv2DTranspose(filters, kernel_size, strides=2, padding='same',\n",
        "                              kernel_initializer=initializer, use_bias=False))\n",
        "    model.add(BatchNormalization())\n",
        "\n",
        "    if apply_dropout:\n",
        "        model.add(Dropout(0.5))\n",
        "\n",
        "    model.add(ReLU())\n",
        "    return model\n",
        "\n",
        "\n",
        "def make_autoencoder_generator_model():\n",
        "    inputs = Input(shape=(IMAGE_SIZE, IMAGE_SIZE, 1))\n",
        "\n",
        "    # Downsampling layers\n",
        "    # 1: (BATCH_SIZE, 16, 16, 32)\n",
        "    # 2: (BATCH_SIZE, 8, 8, 64)\n",
        "    # 3: (BATCH_SIZE, 4, 4, 128)\n",
        "    # 4: (BATCH_SIZE, 2, 2, 256)\n",
        "    # 5: (BATCH_SIZE, 1, 1, 256)\n",
        "\n",
        "    downstack = [\n",
        "        downsample(32, 4, apply_batchnorm=False),\n",
        "        downsample(64, 4),\n",
        "        downsample(128, 4),\n",
        "        downsample(256, 4),\n",
        "        downsample(256, 4)\n",
        "    ]\n",
        "\n",
        "    # Upsampling layers\n",
        "    # 1: (BATCH_SIZE, 1, 1, 256)\n",
        "    # 2: (BATCH_SIZE, 1, 1, 128)\n",
        "    # 3: (BATCH_SIZE, 1, 1, 64)\n",
        "    # 4: (BATCH_SIZE, 1, 1, 32)\n",
        "    \n",
        "    upstack = [\n",
        "        upsample(256, 4, apply_dropout=True),\n",
        "        upsample(128, 4),\n",
        "        upsample(64, 4),\n",
        "        upsample(32, 4),\n",
        "    ]\n",
        "\n",
        "    initializer = tf.random_uniform_initializer(0, 0.02)\n",
        "    output_layer = Conv2DTranspose(2, 3, strides=2, padding='same',\n",
        "                                   kernel_initializer=initializer,\n",
        "                                   activation='tanh')\n",
        "    \n",
        "    x = inputs\n",
        "\n",
        "    # Downsampling layers\n",
        "    skips = []\n",
        "    for dm in downstack:\n",
        "        x = dm(x)\n",
        "        skips.append(x)\n",
        "\n",
        "    skips = reversed(skips[:-1])\n",
        "\n",
        "    # Upsampling layers\n",
        "    for um, skip in zip(upstack, skips):\n",
        "        x = um(x)\n",
        "        x = tf.keras.layers.Concatenate()([x, skip])\n",
        "    \n",
        "    x = output_layer(x)\n",
        "\n",
        "    return Model(inputs=inputs, outputs=x)"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1K3rj5IK8kyP",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def make_discriminator_model():\n",
        "    model = Sequential()\n",
        "    model.add(Conv2D(32, (3, 3), padding='same', activation='relu', input_shape=(IMAGE_SIZE, IMAGE_SIZE, 3)))\n",
        "    model.add(Conv2D(32, (3, 3), padding='same', activation='relu'))\n",
        "    model.add(AveragePooling2D(pool_size=(2, 2)))\n",
        "    model.add(Dropout(.25))\n",
        "\n",
        "    model.add(Conv2D(64, (3, 3), padding='same', activation='relu'))\n",
        "    model.add(Conv2D(64, (3, 3), padding='same', activation='relu'))\n",
        "    model.add(AveragePooling2D(pool_size=(2, 2)))\n",
        "    model.add(Dropout(.25))\n",
        "\n",
        "    model.add(Flatten())\n",
        "    model.add(Dense(512))\n",
        "    model.add(LeakyReLU(.2))\n",
        "    model.add(BatchNormalization())\n",
        "    model.add(Dropout(.5))\n",
        "    model.add(Dense(1))\n",
        "    model.add(Activation('sigmoid'))\n",
        "\n",
        "    return model"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GR_ML92o8mP7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "LAMBDA = 100\n",
        "cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)\n",
        "\n",
        "\n",
        "def discriminator_loss(disc_real_output, disc_generated_output):\n",
        "    real_loss = cross_entropy(tf.ones_like(disc_real_output), disc_real_output)\n",
        "    generated_loss = cross_entropy(tf.zeros_like(disc_generated_output), disc_generated_output)\n",
        "    total_disc_loss = real_loss + generated_loss\n",
        "    return total_disc_loss\n",
        "\n",
        "\n",
        "def generator_loss(disc_generated_output, gen_output, target):\n",
        "    gan_loss = cross_entropy(tf.ones_like(disc_generated_output), disc_generated_output)\n",
        "    # mean absolute error\n",
        "    l1_loss = tf.reduce_mean(tf.abs(target - gen_output))\n",
        "    total_gen_loss = gan_loss + (LAMBDA * l1_loss)\n",
        "    return total_gen_loss, gan_loss, l1_loss"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Vj2E6Mu58nhh",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "generator = make_autoencoder_generator_model()\n",
        "discriminator = make_discriminator_model()\n",
        "\n",
        "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 = os.path.join(WORKDIR, '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)\n",
        "manager = tf.train.CheckpointManager(checkpoint, checkpoint_dir, max_to_keep=3)"
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JsH4W1-x8otE",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "summary_log_file = os.path.join(\n",
        "    WORKDIR, 'tf-summary', datetime.now().strftime(\"%Y%m%d-%H%M%S\"))\n",
        "summary_writer = tf.summary.create_file_writer(summary_log_file)\n",
        "\n",
        "@tf.function\n",
        "def train_step(input_image, target, epoch):\n",
        "    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:\n",
        "        gen_output = generator(input_image, training=True)\n",
        "\n",
        "        disc_real_output = discriminator(tf.keras.layers.concatenate([input_image, target]), training=True)\n",
        "        disc_generated_output = discriminator(tf.keras.layers.concatenate([input_image, gen_output]), training=True)\n",
        "\n",
        "        gen_total_loss, gen_gan_loss, gen_l1_loss = generator_loss(\n",
        "            disc_generated_output, gen_output, target)\n",
        "        disc_loss = discriminator_loss(disc_real_output, disc_generated_output)\n",
        "\n",
        "    generator_gradients = gen_tape.gradient(gen_total_loss,\n",
        "                                            generator.trainable_variables)\n",
        "    discriminator_gradients = disc_tape.gradient(disc_loss,\n",
        "                                                 discriminator.trainable_variables)\n",
        "\n",
        "    generator_optimizer.apply_gradients(zip(generator_gradients,\n",
        "                                          generator.trainable_variables))\n",
        "    discriminator_optimizer.apply_gradients(zip(discriminator_gradients,\n",
        "                                              discriminator.trainable_variables))\n",
        "\n",
        "    with summary_writer.as_default():\n",
        "        tf.summary.scalar('gen_total_loss', gen_total_loss, step=epoch)\n",
        "        tf.summary.scalar('gen_gan_loss', gen_gan_loss, step=epoch)\n",
        "        tf.summary.scalar('gen_l1_loss', gen_l1_loss, step=epoch)\n",
        "        tf.summary.scalar('disc_loss', disc_loss, step=epoch)\n",
        "    \n",
        "    return gen_total_loss, disc_loss"
      ],
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "88nS2UE38p_x",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "07f82056-a443-4b26-a13a-8079cec417d1"
      },
      "source": [
        "checkpoint.restore(manager.latest_checkpoint)\n",
        "if manager.latest_checkpoint:\n",
        "    print('Restored from {}'.format(manager.latest_checkpoint))\n",
        "else:\n",
        "    print('Initializing from scratch')"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Restored from /content/drive/My Drive/Colab Notebooks/Workspace/training-checkpoints/ckpt-27\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "O96u8HUY_rSe",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "for e in range(EPOCHS):\n",
        "    start_time = time()\n",
        "    gen_loss_total = disc_loss_total = 0\n",
        "    for input_image, target in train_dataset:\n",
        "        gen_loss, disc_loss = train_step(input_image, target, e)\n",
        "        gen_loss_total += gen_loss\n",
        "        disc_loss_total += disc_loss\n",
        "\n",
        "    time_taken = time() - start_time\n",
        "\n",
        "    if (e + 1) % 10 == 0:\n",
        "        checkpoint.save(file_prefix=checkpoint_prefix)\n",
        "    \n",
        "    # print('Epoch {}: gen loss: {}, disc loss: {}, time: {:.2f}s'.format(\n",
        "    #     e + 1, gen_loss_total / BATCH_SIZE, disc_loss_total / BATCH_SIZE, time_taken))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FAzMrQKo8wlC",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "Y_hat = generator(X_test[:250])\n",
        "total_count = len(Y_hat)\n",
        "\n",
        "for idx, (x, y, y_hat) in enumerate(zip(X_test[:250], Y_test[:250], Y_hat)):\n",
        "\n",
        "    # Original RGB image\n",
        "    orig_lab = np.dstack((x, y * 128))\n",
        "    orig_rgb = lab2rgb(orig_lab)\n",
        "\n",
        "    # Grayscale version of the original image\n",
        "    grayscale_lab = np.dstack((x, np.zeros((IMAGE_SIZE, IMAGE_SIZE, 2))))\n",
        "    grayscale_rgb = lab2rgb(grayscale_lab)\n",
        "\n",
        "    # Colorized image\n",
        "    predicted_lab = np.dstack((x, y_hat * 128))\n",
        "    predicted_rgb = lab2rgb(predicted_lab)\n",
        "\n",
        "    # print(idx)\n",
        "\n",
        "    # fig = plt.figure()\n",
        "    # fig.add_subplot(1, 3, 1)\n",
        "    # plt.axis('off')\n",
        "    # plt.imshow(grayscale_rgb)\n",
        "    # fig.add_subplot(1, 3, 2)\n",
        "    # plt.axis('off')\n",
        "    # plt.imshow(orig_rgb)\n",
        "    # fig.add_subplot(1, 3, 3)\n",
        "    # plt.axis('off')\n",
        "    # plt.imshow(predicted_rgb)\n",
        "    # plt.show()\n",
        "    \n",
        "    plt.axis('off')\n",
        "    plt.imshow(grayscale_rgb)\n",
        "    plt.savefig(os.path.join(WORKDIR, 'results', '{}-bw.png'.format(idx)), dpi=1)\n",
        "\n",
        "    plt.axis('off')\n",
        "    plt.imshow(orig_rgb)\n",
        "    plt.savefig(os.path.join(WORKDIR, 'results', '{}-gt.png'.format(idx)), dpi=1)\n",
        "\n",
        "    plt.axis('off')\n",
        "    plt.imshow(predicted_rgb)\n",
        "    plt.savefig(os.path.join(WORKDIR, 'results', '{}-gan.png'.format(idx)), dpi=1)\n",
        "\n",
        "    sys.stdout.flush()\n",
        "    sys.stdout.write('\\r{} / {}'.format(idx + 1, total_count))\n",
        "\n",
        "None"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zqcecjWOXFaW",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 105
        },
        "outputId": "3bb42a59-e443-4b07-ccf8-1438a1fd220c"
      },
      "source": [
        "tf.saved_model.save(generator, os.path.join(WORKDIR, \"generator-saved-model\"))"
      ],
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/resource_variable_ops.py:1817: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "If using Keras pass *_constraint arguments to layers.\n",
            "INFO:tensorflow:Assets written to: /content/drive/My Drive/Colab Notebooks/Workspace/generator-saved-model/assets\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IibPhDCcXZRy",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "771f5042-0460-49d2-faff-707022e7eb35"
      },
      "source": [
        "tf.saved_model.save(discriminator, os.path.join(WORKDIR, \"disciminator-saved-model\"))"
      ],
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "INFO:tensorflow:Assets written to: /content/drive/My Drive/Colab Notebooks/Workspace/disciminator-saved-model/assets\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cBqVEbm5Xgpj",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}