SRGAN_coco_continue.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "SRGAN_coco_continue.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "machine_shape": "hm"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rgWe5LRFQ4DC",
        "colab_type": "text"
      },
      "source": [
        "# SRGAN CONTINUE TRAINING (COCO)\n",
        "This notebook is purely for training on Google Colab where sessions get randomly terminated in 8-12 hours of running.  Truly grateful to Google Colab for being free.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b_N3vLscRF1X",
        "colab_type": "text"
      },
      "source": [
        "### Start with cloning repo\n",
        "I used drive and local to get and save data. Tried GitHub but it was too much trouble managing the commits."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kTUU5Wc2ISH1",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os, shutil\n",
        "from google.colab import files, drive\n",
        "\n",
        "# save locations during training\n",
        "drive.mount('/content/drive')\n",
        "\n",
        "# unzip local file (preprocessed 176*176 images)\n",
        "if not os.path.exists(\"coco\"):\n",
        "    !unzip coco\n",
        "!rm coco.zip\n",
        "\n",
        "# directories\n",
        "data_dir = \"coco/\"                                   # input data directory\n",
        "model_dir = \"drive/My Drive/SRGAN_coco/model/\"       # save model directory\n",
        "output_dir = \"drive/My Drive/SRGAN_coco/output/\"     # save output directory"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-EU9UKLKSniH",
        "colab_type": "text"
      },
      "source": [
        "---\n",
        "### Initialize (Hyper)parameters\n",
        "\n",
        "With better GPU, higher batch sizes and image sizes could be used to improve the results."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YYLYLkauLuzk",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# =================\n",
        "# (HYPER)PARAMETERS\n",
        "# =================\n",
        "\n",
        "# IMAGE\n",
        "down_scale = 4                                       # downsampling scale\n",
        "cropped_width = 176                                  # high res image (width)\n",
        "cropped_height = 176                                 # high res image (height)\n",
        "image_shape = (176, 176, 3)                          # high res image (shape)\n",
        "\n",
        "# TRAIN\n",
        "begin_epoch = 2675                                   # the epoch to begin training at\n",
        "num_images = 2000                                    # total number of images (train & test)\n",
        "split_ratio = 0.8                                    # train-test split ratio\n",
        "epochs = 5000                                        # number of epochs\n",
        "batch_size = 32                                      # number of images for each batch\n",
        "sample_every = 1                                     # number of epochs in between sampling\n",
        "save_every = 25                                      # number of epochs in between saving model"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F-IFV0nJTSuL",
        "colab_type": "text"
      },
      "source": [
        "---\n",
        "### Preprocess Images\n",
        "\n",
        "Crop center, down sample, then normalize. I chose to normalize both LR and HR images to [-1, 1]. Some operations are done with utility functions, not included here.\n",
        "\n",
        "**NOTE**:\\\n",
        "I ended up moving out most of preprocess functions out since in later training stages, I downloaded the outputs of these cells and zipped them for convenience. However, feel free to checkout the `utils.py` file for how I preprocessed the images and how I saved and loaded them."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6C3xE8qAsKnF",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import numpy as np\n",
        "from numpy import random\n",
        "from PIL import Image\n",
        "import matplotlib.pyplot as plt\n",
        "import glob\n",
        "%matplotlib inline\n",
        "\n",
        "def down_sample(img, scale=down_scale):\n",
        "    \"\"\"Convert image to lower resolution.\"\"\"\n",
        "    new_h = img.shape[0]//scale;\n",
        "    new_w = img.shape[1]//scale;\n",
        "    lr_img = np.asarray(Image.fromarray(np.uint8(img)).resize((new_w, new_h), Image.BICUBIC))\n",
        "    return lr_img\n",
        "\n",
        "def normalize(img):\n",
        "    \"\"\"Normalize image to [-1,1].\"\"\"\n",
        "    n_img = np.divide(img.astype(np.float32), 127.5) - np.ones_like(img, dtype=np.float32)\n",
        "    return n_img"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Smc2wM6gVddq",
        "colab_type": "text"
      },
      "source": [
        "Collect preprocessed datasets."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F3LwsphHhyei",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "hr_images = []\n",
        "lr_images = []\n",
        "\n",
        "def get_data(data_dir=data_dir):\n",
        "    \"\"\"Populate 4D arrays of high res and low res images.\"\"\"\n",
        "    for img_name in glob.glob(data_dir + \"*.jpg\"):\n",
        "        hr_img = np.asarray(Image.open(img_name))\n",
        "        lr_img = down_sample(hr_img)\n",
        "        hr_images.append(hr_img)\n",
        "        lr_images.append(lr_img)\n",
        "        \n",
        "get_data()\n",
        "print(\"High resolution image dataset shape:\", np.shape(hr_images))\n",
        "print(\"Low resolution image dataset shape:\", np.shape(lr_images))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fN_O5IayWd_5",
        "colab_type": "text"
      },
      "source": [
        "---\n",
        "### Untility Function for Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OiRouZDyv3HT",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def load_train_data(data_dir=data_dir, num_img=num_images, \n",
        "                    split_ratio=split_ratio, \n",
        "                    hr_images=hr_images, lr_images=lr_images):\n",
        "    \"\"\"Perform train-test split for high and low res images.\"\"\"\n",
        "\n",
        "    num_train = int(num_img * split_ratio)\n",
        "    \n",
        "    hr_images = [normalize(img) for img in hr_images]\n",
        "    lr_images = [normalize(img) for img in lr_images]\n",
        "    \n",
        "    hr_train = hr_images[:num_train]\n",
        "    hr_test = hr_images[num_train:]\n",
        "    lr_train = lr_images[:num_train]\n",
        "    lr_test = lr_images[num_train:]\n",
        "    \n",
        "    return hr_train, hr_test, lr_train, lr_test\n",
        "    \n",
        "# sanity check on dataset shapes\n",
        "hr_train, hr_test, lr_train, lr_test = load_train_data()\n",
        "print(\"HR image training dataset shape:\", np.shape(hr_train), \"\\t\")\n",
        "print(\"LR image training dataset shape:\", np.shape(lr_train), \"\\t\")\n",
        "print(\"HR image testing dataset shape:\", np.shape(hr_test), \"\\t\")\n",
        "print(\"LR image testing dataset shape:\", np.shape(lr_test), \"\\t\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tsd-wpeHAbdm",
        "colab_type": "text"
      },
      "source": [
        "---\n",
        "### Build the SRGAN Neural Network\n",
        "\n",
        "For visual details for on the network architecture, refer to the README.md file of the github repo."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yntra3m1v3Xi",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from keras.layers.core import Activation, Flatten\n",
        "from keras.layers import Input, add, LeakyReLU, PReLU\n",
        "from keras.layers import BatchNormalization, Conv2D, UpSampling2D, Dense\n",
        "from keras.models import Model, load_model\n",
        "from keras.optimizers import Adam"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DC2w1R4Tb_yD",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from keras.applications.vgg19 import VGG19\n",
        "import keras.backend as K\n",
        "\n",
        "def content_loss(y, y_pred):\n",
        "    vgg19_model = VGG19(include_top=False, weights=\"imagenet\", input_shape=image_shape)\n",
        "    # Not trainable\n",
        "    vgg19_model.trainable = False\n",
        "    for layer in vgg19_model.layers:\n",
        "        layer.trainable = False\n",
        "\n",
        "    model = Model(inputs=vgg19_model.input, outputs=vgg19_model.get_layer(\"block5_conv4\").output)\n",
        "    model.trainable = False\n",
        "\n",
        "    return K.mean(K.square(model(y) - model(y_pred)))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nYSMc3neDifI",
        "colab_type": "text"
      },
      "source": [
        "---\n",
        "\n",
        "### Load Model-in-Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XMj2I1Kv5hmN",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "d = load_model('coco_d_model{}.h5'.format(begin_epoch))\n",
        "g = load_model('coco_g_model{}.h5'.format(begin_epoch), custom_objects={'content_loss': content_loss})"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "h6m5kBePd5u7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def GAN_NN(g, d, shape, optimizer, content_loss):\n",
        "    \"\"\"Build and compile the GAN network to connect generator and discriminator.\"\"\"\n",
        "    d.trainable = False\n",
        "    \n",
        "    gan_input = Input(shape=shape)\n",
        "    fake = g(gan_input)\n",
        "    gan_output = d(fake)\n",
        "    \n",
        "    gan_model = Model(inputs=gan_input, outputs=[fake, gan_output])\n",
        "    gan_model.compile(loss=[content_loss, \"binary_crossentropy\"], \n",
        "                loss_weights=[1., 1e-3], optimizer=optimizer)\n",
        "    \n",
        "    return gan_model"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-oII0Rnvv3Yr",
        "colab_type": "text"
      },
      "source": [
        "---\n",
        "### Sample Generator Outputs"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Wy4-_6oED_eN",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def generate_image(e, g, hr_test, lr_test,\n",
        "                    output_dir=output_dir,\n",
        "                    dim=(1, 3), figsize=(15, 5)):\n",
        "    \"\"\"Util function for visualizing results of generator.\"\"\"\n",
        "    hr_batch = np.asarray(hr_test)\n",
        "    lr_batch = np.asarray(lr_test)\n",
        "    sr_batch = np.asarray(g.predict(lr_batch))\n",
        "    \n",
        "    # denormalize\n",
        "    hr_batch = ((hr_batch + 1) * 127.5).astype(np.uint8)\n",
        "    lr_batch = ((lr_batch + 1) * 127.5).astype(np.uint8)\n",
        "    sr_batch = ((sr_batch + 1) * 127.5).astype(np.uint8)\n",
        "    \n",
        "    # random sample\n",
        "    idx = random.randint(0, len(hr_test))\n",
        "\n",
        "    # 1 row of 3 images (low res -> super res -> high res)\n",
        "    plt.figure(figsize=figsize)\n",
        "    plt.subplot(dim[0], dim[1], 1)\n",
        "    plt.imshow(lr_batch[idx], interpolation=\"nearest\")\n",
        "    plt.axis(\"off\")\n",
        "    plt.subplot(dim[0], dim[1], 2)\n",
        "    plt.imshow(sr_batch[idx], interpolation=\"nearest\")\n",
        "    plt.axis(\"off\")\n",
        "    plt.subplot(dim[0], dim[1], 3)\n",
        "    plt.imshow(hr_batch[idx], interpolation=\"nearest\")\n",
        "    plt.axis(\"off\")\n",
        "    plt.tight_layout()\n",
        "    plt.savefig(output_dir + \"result_{}.png\".format(e))\n",
        "    plt.close()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C9r5o4rtwV7Z",
        "colab_type": "text"
      },
      "source": [
        "---\n",
        "### Set up File Structure to Store Training Results"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MJk9bqu7fQBl",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# tqdm gives more info about training\n",
        "from tqdm import tqdm\n",
        "from google.colab import output\n",
        "\n",
        "def setup_training(data_dir=data_dir):\n",
        "    \"\"\"Make model and output directories, also initialize loss file.\"\"\"\n",
        "    if os.path.isdir(model_dir): shutil.rmtree(model_dir)\n",
        "    if os.path.isdir(output_dir): shutil.rmtree(output_dir)\n",
        "    if os.path.exists(model_dir + \"loss.txt\"): os.remove(model_dir + \"loss.txt\")\n",
        "    \n",
        "    os.mkdir(model_dir)\n",
        "    os.mkdir(output_dir)\n",
        "    loss_file = open(model_dir + \"loss.txt\" , \"w+\")\n",
        "    loss_file.close()\n",
        "\n",
        "setup_training()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fUIkuqr5wfrv",
        "colab_type": "text"
      },
      "source": [
        "---\n",
        "### Training\n",
        "For details on the training steps, refer to README.md."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ht2ftfTmi_Vl",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def SRGAN(epochs=epochs, batch_size=batch_size, split_ratio=split_ratio, \n",
        "          sample_every=sample_every, save_every=save_every,\n",
        "          shape=image_shape, scale=down_scale, num_imgs=num_images,\n",
        "          begin_epoch=begin_epoch, g=g, d=d):\n",
        "    \"\"\"The SRGAN training pipline.\"\"\"\n",
        "    \n",
        "    hr_train, hr_test, lr_train, lr_test = load_train_data()\n",
        "    \n",
        "    num_batches = int(len(hr_train) // batch_size)\n",
        "    shape_small = (shape[0]//scale, shape[1]//scale, shape[2])\n",
        "    \n",
        "    g_loss = content_loss\n",
        "    gan = GAN_NN(g, d, shape_small, d.optimizer, g_loss)\n",
        "    \n",
        "    for e in range(begin_epoch+1, begin_epoch+epochs+1):\n",
        "        print(\"\\nEPOCH {}\".format(e))\n",
        "        for _ in tqdm(range(num_batches)):\n",
        "            idxs = random.randint(0, len(hr_train), size=batch_size)\n",
        "            hr_batch = []\n",
        "            lr_batch = []\n",
        "            hr_batch = [hr_train[i] for i in idxs]\n",
        "            hr_batch = np.asarray(hr_batch)\n",
        "            lr_batch = [lr_train[i] for i in idxs]\n",
        "            lr_batch = np.asarray(lr_batch)\n",
        "            sr_batch = g.predict(lr_batch)\n",
        "            \n",
        "            # std = 0.05, mean = 0.9 \n",
        "            real_label = 0.05 * random.randn(batch_size) + 0.9\n",
        "            # std = 0.05, mean = 0.1\n",
        "            fake_label = 0.05 * random.randn(batch_size) + 0.1\n",
        "            \n",
        "            d.trainable = True\n",
        "            d_loss_real = d.train_on_batch(hr_batch, real_label)\n",
        "            d_loss_fake = d.train_on_batch(sr_batch, fake_label)\n",
        "            d_loss = np.add(d_loss_real, d_loss_fake) / 2.0\n",
        "            d.trainable = False\n",
        "\n",
        "            idxs = random.randint(0, len(hr_train), size=batch_size)\n",
        "            hr_batch = []\n",
        "            lr_batch = []\n",
        "            hr_batch = [hr_train[i] for i in idxs]\n",
        "            hr_batch = np.asarray(hr_batch)\n",
        "            lr_batch = [lr_train[i] for i in idxs]\n",
        "            lr_batch = np.asarray(lr_batch)\n",
        "            sr_batch = g.predict(lr_batch)\n",
        "            \n",
        "            # std = 0.05, mean = 0.9 \n",
        "            gan_label = 0.05 * random.randn(batch_size) + 0.9\n",
        "            \n",
        "            gan_loss = gan.train_on_batch(lr_batch, [hr_batch, gan_label])\n",
        "            \n",
        "        output.clear()\n",
        "\n",
        "        loss_file = open(model_dir + \"loss.txt\", \"a\")\n",
        "        loss_file.write(\"EPOCH {}\\td_loss {}\\tgan_loss {}\\n\".format(e, d_loss, str(gan_loss)))\n",
        "        loss_file.close()\n",
        "\n",
        "        if (e == 1) or (e % sample_every == 0):\n",
        "            generate_image(e, g, hr_test, lr_test)\n",
        "        if (e % save_every == 0):\n",
        "            g.save(model_dir + \"coco_g_model{}.h5\".format(e))\n",
        "            d.save(model_dir + \"coco_d_model{}.h5\".format(e))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LD2bxWKkIc1i",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# train, output in loss.txt\n",
        "SRGAN()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VY11xz6ifFNH",
        "colab_type": "text"
      },
      "source": [
        "---\n",
        "### View Some Results\n",
        "\n",
        "Render result every 25 epochs. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "gVcOumIGXbgV",
        "colab": {}
      },
      "source": [
        "# def visualize_result(output_dir=output_dir):\n",
        "#     \"\"\"View All Images.\"\"\"\n",
        "    \n",
        "#     images_shown = []\n",
        "#     for file_name in os.listdir(output_dir):\n",
        "#         name, ext = file_name.split(\".\")\n",
        "#         if int(name) / 25 == 0:\n",
        "#             image = np.asarray(Image.open(file_name))\n",
        "#             images_shown.append(image)\n",
        "\n",
        "#     fig, ax = plt.subplots(1,len(images_shown))\n",
        "#     for img in range(len(image_shown)):\n",
        "#         ax[img].set_title(\"sample \" + str(img+1))\n",
        "#         ax[img].imshow(images_shown[img])"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}