StackGAN.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "StackGAN.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "Fz27a-5GAbUl"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "U_ngLlWRxqAf"
      },
      "source": [
        "import tensorflow as tf\n",
        "from tensorflow.keras import backend as K\n",
        "import numpy as np\n",
        "import os\n",
        "from matplotlib import pyplot as plt\n",
        "import pickle\n",
        "import cv2\n",
        "import time"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UkrXDsWPwTF4"
      },
      "source": [
        "def KL_loss(y_true, y_pred):\n",
        "  mean = y_pred[:, :128]\n",
        "  logsigma = y_pred[:, 128:]\n",
        "  loss = -logsigma + 0.5*(-1 + K.exp(2.0*logsigma) + K.square(mean))\n",
        "  loss = K.mean(loss)\n",
        "  return loss"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VC2vJ-YBL9U9"
      },
      "source": [
        "# Data Pipeline"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ajlV6kDMIYf8"
      },
      "source": [
        "WORKERS = tf.data.experimental.AUTOTUNE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "w_UaZT_wHej6"
      },
      "source": [
        "class Dataset():\n",
        "  def __init__(self, image_size, batch_size = 64):\n",
        "    data_dir = \"/content/drive/My Drive/Text to Image/dataset/birds\"\n",
        "    train_dir = data_dir + \"/train\"\n",
        "    test_dir = data_dir + \"/test\"\n",
        "    train_embeddings_path = train_dir + \"/char-CNN-RNN-embeddings.pickle\"\n",
        "    test_embeddings_path = test_dir + \"/char-CNN-RNN-embeddings.pickle\"\n",
        "    train_filenames_path = train_dir + \"/filenames.pickle\"\n",
        "    test_filenames_path = test_dir + \"/filenames.pickle\"\n",
        "    cub_dataset_dir = \"/content/drive/My Drive/Text to Image/dataset/CUB_200_2011/CUB_200_2011\"\n",
        "    bounding_boxes_path = cub_dataset_dir + \"/bounding_boxes.txt\"\n",
        "    image_ids_path = cub_dataset_dir + \"/images.txt\"\n",
        "    images_path = cub_dataset_dir + \"/images\"\n",
        "    \n",
        "    self.image_width = image_size[0]\n",
        "    self.image_height = image_size[1]\n",
        "    self.batch_size = batch_size\n",
        "\n",
        "    with open(train_filenames_path, 'rb') as f:\n",
        "      self.train_filenames = pickle.load(f, encoding='latin1')\n",
        "      self.train_filenames = [os.path.join(images_path, filename) + '.jpg' for filename in self.train_filenames]\n",
        "    \n",
        "    with open(test_filenames_path, 'rb') as f:\n",
        "      self.test_filenames = pickle.load(f, encoding='latin1')\n",
        "      self.test_filenames = [os.path.join(images_path, filename) + '.jpg' for filename in self.test_filenames]\n",
        "\n",
        "    with open(train_embeddings_path, 'rb') as f:\n",
        "      self.train_embeddings = pickle.load(f, encoding = 'latin1')\n",
        "\n",
        "    with open(test_embeddings_path, 'rb') as f:\n",
        "      self.test_embeddings = pickle.load(f, encoding = 'latin1')\n",
        "\n",
        "    bounding_boxes = {}\n",
        "    with open(bounding_boxes_path, 'rb') as f:\n",
        "      box_coordinates = f.read()\n",
        "      box_coordinates = box_coordinates.splitlines()\n",
        "      box_coordinates = [box_coordinate.decode('utf-8') for box_coordinate in box_coordinates]\n",
        "      for i in range(len(box_coordinates)):\n",
        "        bounding_box = box_coordinates[i].split()\n",
        "        bounding_boxes[bounding_box[0]] = [int(float(c)) for c in box_coordinates[i].split()][1:]\n",
        "\n",
        "    image_ids_mapping = {}\n",
        "    with open(image_ids_path, 'rb') as f:\n",
        "      image_ids = f.read()\n",
        "      image_ids = image_ids.splitlines()\n",
        "      image_ids = [image_id.decode('utf-8') for image_id in image_ids]\n",
        "      for i in range(len(image_ids)):\n",
        "        image_id = image_ids[i].split()\n",
        "        image_ids_mapping[image_id[0]] = image_id[1]\n",
        "\n",
        "    bounding_boxes_mapping = {}\n",
        "    for image_id in bounding_boxes.keys():\n",
        "      bounding_boxes_mapping[images_path + \"/\" + image_ids_mapping[image_id]] = bounding_boxes[image_id]\n",
        "\n",
        "    self.train_bounding_boxes = []\n",
        "    self.test_bounding_boxes = []\n",
        "    for i in range(len(self.train_filenames)):\n",
        "      self.train_bounding_boxes.append(bounding_boxes_mapping[self.train_filenames[i]])\n",
        "    for i in range(len(self.test_filenames)):\n",
        "      self.test_bounding_boxes.append(bounding_boxes_mapping[self.test_filenames[i]])\n",
        "\n",
        "  def crop(self, image, bounding_box):\n",
        "    image = image.numpy()\n",
        "    if bounding_box is not None:\n",
        "      x, y, width, height = bounding_box\n",
        "      image = image[y:(y + height), x:(x + width)]\n",
        "      image = cv2.resize(image, (self.image_width, self.image_height))\n",
        "    return image\n",
        "\n",
        "  def parse_function(self, image_path, embeddings, bounding_box):\n",
        "    image = tf.io.read_file(image_path)\n",
        "    image = tf.image.decode_jpeg(image, channels = 3)\n",
        "    image = tf.py_function(func = self.crop, inp = [image, bounding_box], Tout = tf.float32)\n",
        "    image.set_shape([self.image_width, self.image_height, 3])\n",
        "    image = (image - 127.5) / 127.5\n",
        "\n",
        "    embedding_index = np.random.randint(0, embeddings.shape[0] - 1)\n",
        "    embedding = embeddings[embedding_index]\n",
        "    return image, embedding\n",
        "  \n",
        "  def get_train_ds(self):\n",
        "    BUFFER_SIZE = len(self.train_filenames)\n",
        "    ds = tf.data.Dataset.from_tensor_slices((self.train_filenames, self.train_embeddings, self.train_bounding_boxes))\n",
        "    ds = ds.shuffle(BUFFER_SIZE)\n",
        "    ds = ds.repeat()\n",
        "    ds = ds.map(self.parse_function, num_parallel_calls = WORKERS)\n",
        "    ds = ds.batch(self.batch_size, drop_remainder = True)\n",
        "    ds = ds.prefetch(1)\n",
        "    return ds\n",
        "  \n",
        "  def get_test_ds(self):\n",
        "    BUFFER_SIZE = len(self.test_filenames)\n",
        "    ds = tf.data.Dataset.from_tensor_slices((self.test_filenames, self.test_embeddings, self.test_bounding_boxes))\n",
        "    ds = ds.shuffle(BUFFER_SIZE)\n",
        "    ds = ds.repeat(1)\n",
        "    ds = ds.map(self.parse_function, num_parallel_calls = WORKERS)\n",
        "    ds = ds.batch(self.batch_size, drop_remainder = True)\n",
        "    ds = ds.prefetch(1)\n",
        "    return ds"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZLPBCNUeLygV"
      },
      "source": [
        "#Stage 1 StackGAN"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wRcUy6S1Ckx5"
      },
      "source": [
        "class ConditioningAugmentation(tf.keras.Model):\n",
        "  def __init__(self):\n",
        "    super(ConditioningAugmentation, self).__init__()\n",
        "    self.dense = tf.keras.layers.Dense(units = 256)\n",
        "\n",
        "  def call(self, E):\n",
        "    X = self.dense(E)\n",
        "    phi = tf.nn.leaky_relu(X)\n",
        "    mean = phi[:, :128]\n",
        "    std = K.exp(phi[:, 128:])\n",
        "    epsilon = K.random_normal(shape = K.constant((mean.shape[1], ), dtype = 'int32'))\n",
        "    C = mean + epsilon*std\n",
        "    return C, phi"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jhVBzTUnDMmT"
      },
      "source": [
        "class EmbeddingCompressor(tf.keras.Model):\n",
        "  def __init__(self):\n",
        "    super(EmbeddingCompressor, self).__init__()\n",
        "    self.dense = tf.keras.layers.Dense(units = 128)\n",
        "\n",
        "  def call(self, E):\n",
        "    X = self.dense(E)\n",
        "    return tf.nn.relu(X)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1cn9qNR721K-"
      },
      "source": [
        "class Stage1Generator(tf.keras.Model):\n",
        "  def __init__(self):\n",
        "    super(Stage1Generator, self).__init__()\n",
        "    self.canet = ConditioningAugmentation()\n",
        "    self.concat = tf.keras.layers.Concatenate(axis = 1)\n",
        "    self.dense = tf.keras.layers.Dense(units = 128*8*4*4, kernel_initializer = tf.random_normal_initializer(stddev = 0.02))\n",
        "    self.reshape = tf.keras.layers.Reshape(target_shape = (4, 4, 128*8), input_shape = (128*8*4*4, ))\n",
        "    self.batchnorm1 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.deconv1 = tf.keras.layers.Conv2DTranspose(filters = 512, kernel_size = 4, strides = (2, 2), padding = \"same\", kernel_initializer = tf.random_normal_initializer(stddev = 0.02))\n",
        "    self.batchnorm2 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.deconv2 = tf.keras.layers.Conv2DTranspose(filters = 256, kernel_size = 4, strides = (2, 2), padding = \"same\", kernel_initializer = tf.random_normal_initializer(stddev = 0.02))\n",
        "    self.batchnorm3 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.deconv3 = tf.keras.layers.Conv2DTranspose(filters = 128, kernel_size = 4, strides = (2, 2), padding = \"same\", kernel_initializer = tf.random_normal_initializer(stddev = 0.02))\n",
        "    self.batchnorm4 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.deconv4 = tf.keras.layers.Conv2DTranspose(filters = 3, kernel_size = 4, strides = (2, 2), padding = \"same\", kernel_initializer = tf.random_normal_initializer(stddev = 0.02))\n",
        "\n",
        "  def call(self, inputs):\n",
        "    E, Z = inputs\n",
        "    C, phi = self.canet(E)\n",
        "\n",
        "    gen_input = self.concat([C, Z])\n",
        "    X = self.dense(gen_input)\n",
        "    X = self.reshape(X)\n",
        "    X = self.batchnorm1(X)\n",
        "    X = tf.nn.relu(X)\n",
        "\n",
        "    X = self.deconv1(X)\n",
        "    X = self.batchnorm1(X)\n",
        "    X = tf.nn.relu(X)\n",
        "\n",
        "    X = self.deconv2(X)\n",
        "    X = self.batchnorm2(X)\n",
        "    X = tf.nn.relu(X)\n",
        "\n",
        "    X = self.deconv3(X)\n",
        "    X = self.batchnorm3(X)\n",
        "    X = tf.nn.relu(X)\n",
        "\n",
        "    X = self.deconv4(X)\n",
        "    return tf.nn.tanh(X), phi"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F89MagHt4BHc"
      },
      "source": [
        "class Stage1Discriminator(tf.keras.Model):\n",
        "  def __init__(self):\n",
        "    super(Stage1Discriminator, self).__init__()\n",
        "    self.conv1 = tf.keras.layers.Conv2D(filters = 64, kernel_size = 4, strides = 2, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.conv2 = tf.keras.layers.Conv2D(filters = 128, kernel_size = 4, strides = 2, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.batchnorm1 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.conv3 = tf.keras.layers.Conv2D(filters = 256, kernel_size = 4, strides = 2, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.batchnorm2 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.conv4 = tf.keras.layers.Conv2D(filters = 512, kernel_size = 4, strides = 2, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.batchnorm3 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.embed = EmbeddingCompressor()\n",
        "    self.reshape = tf.keras.layers.Reshape(target_shape = (1, 1, 128))\n",
        "    self.concat = tf.keras.layers.Concatenate()\n",
        "    self.conv5 = tf.keras.layers.Conv2D(filters = 64*8, kernel_size = 1, strides = 1, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.batchnorm4 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.conv6 = tf.keras.layers.Conv2D(filters = 1, kernel_size = 4, strides = 1, padding = \"valid\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "\n",
        "  def call(self, inputs):\n",
        "    I, E = inputs\n",
        "    X = self.conv1(I)\n",
        "    X = tf.nn.leaky_relu(X)\n",
        "\n",
        "    X = self.conv2(X)\n",
        "    X = self.batchnorm1(X)\n",
        "    X = tf.nn.leaky_relu(X)\n",
        "\n",
        "    X = self.conv3(X)\n",
        "    X = self.batchnorm2(X)\n",
        "    X = tf.nn.leaky_relu(X)\n",
        "\n",
        "    X = self.conv4(X)\n",
        "    X = self.batchnorm3(X)\n",
        "    X = tf.nn.leaky_relu(X)\n",
        "\n",
        "    T = self.embed(E)\n",
        "    T = self.reshape(T)\n",
        "    T = tf.tile(T, (1, 4, 4, 1))\n",
        "    merged_input = self.concat([X, T])\n",
        "\n",
        "    Y = self.conv5(merged_input)\n",
        "    Y = self.batchnorm4(Y)\n",
        "    Y = tf.nn.leaky_relu(Y)\n",
        "\n",
        "    Y = self.conv6(Y)\n",
        "    return tf.squeeze(Y)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2BnqGJG54A9O"
      },
      "source": [
        "class Stage1Model(tf.keras.Model):\n",
        "  def __init__(self):\n",
        "    super(Stage1Model, self).__init__()\n",
        "    self.stage1_generator = Stage1Generator()\n",
        "    self.stage1_discriminator = Stage1Discriminator()\n",
        "\n",
        "  def train(self, train_ds, batch_size = 64, num_epochs = 600, z_dim = 100, c_dim = 128, stage1_generator_lr = 0.0004, stage1_discriminator_lr = 0.0004):\n",
        "    generator_optimizer = tf.keras.optimizers.Adam(lr = stage1_generator_lr, beta_1 = 0.5, beta_2 = 0.999)\n",
        "    discriminator_optimizer = tf.keras.optimizers.Adam(lr = stage1_discriminator_lr, beta_1 = 0.5, beta_2 = 0.999)\n",
        "    \n",
        "    for epoch in range(num_epochs):\n",
        "      print(\"Epoch %d/%d:\\n [\"%(epoch + 1, num_epochs), end = \"\")\n",
        "      start_time = time.time()\n",
        "      if epoch % 100 == 0:\n",
        "        K.set_value(generator_optimizer.learning_rate, generator_optimizer.learning_rate / 2)\n",
        "        K.set_value(discriminator_optimizer.learning_rate, discriminator_optimizer.learning_rate / 2)\n",
        "    \n",
        "      generator_loss_log = []\n",
        "      discriminator_loss_log = []\n",
        "      steps_per_epoch = 125\n",
        "      batch_iter = iter(train_ds)\n",
        "      for i in range(steps_per_epoch):\n",
        "        if i % 5 == 0:\n",
        "          print(\"=\", end = \"\")\n",
        "        image_batch, embedding_batch = next(batch_iter)\n",
        "        z_noise = tf.random.normal((batch_size, z_dim))\n",
        "\n",
        "        mismatched_images = tf.roll(image_batch, shift = 1, axis = 0)\n",
        "\n",
        "        real_labels = tf.random.uniform(shape = (batch_size, ), minval = 0.9, maxval = 1.0)\n",
        "        fake_labels = tf.random.uniform(shape = (batch_size, ), minval = 0.0, maxval = 0.1)\n",
        "        mismatched_labels = tf.random.uniform(shape = (batch_size, ), minval = 0.0, maxval = 0.1)\n",
        "\n",
        "        with tf.GradientTape() as generator_tape, tf.GradientTape() as discriminator_tape:\n",
        "          fake_images, phi = self.stage1_generator([embedding_batch, z_noise])\n",
        "          \n",
        "          real_logits = self.stage1_discriminator([image_batch, embedding_batch])\n",
        "          fake_logits = self.stage1_discriminator([fake_images, embedding_batch])\n",
        "          mismatched_logits = self.stage1_discriminator([mismatched_images, embedding_batch])\n",
        "          \n",
        "          l_sup = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(real_labels, fake_logits))\n",
        "          l_klreg = KL_loss(tf.random.normal((phi.shape[0], phi.shape[1])), phi)\n",
        "          generator_loss = l_sup + 2.0*l_klreg\n",
        "          \n",
        "          l_real = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(real_labels, real_logits))\n",
        "          l_fake = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(fake_labels, fake_logits))\n",
        "          l_mismatched = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(mismatched_labels, mismatched_logits))\n",
        "          discriminator_loss = 0.5*tf.add(l_real, 0.5*tf.add(l_fake, l_mismatched))\n",
        "        \n",
        "        generator_gradients = generator_tape.gradient(generator_loss, self.stage1_generator.trainable_variables)\n",
        "        discriminator_gradients = discriminator_tape.gradient(discriminator_loss, self.stage1_discriminator.trainable_variables)\n",
        "        \n",
        "        generator_optimizer.apply_gradients(zip(generator_gradients, self.stage1_generator.trainable_variables))\n",
        "        discriminator_optimizer.apply_gradients(zip(discriminator_gradients, self.stage1_discriminator.trainable_variables))\n",
        "        \n",
        "        generator_loss_log.append(generator_loss)\n",
        "        discriminator_loss_log.append(discriminator_loss)\n",
        "\n",
        "      end_time = time.time()\n",
        "\n",
        "      if epoch % 1 == 0:\n",
        "        epoch_time = end_time - start_time\n",
        "        template = \"] - generator_loss: {:.4f} - discriminator_loss: {:.4f} - epoch_time: {:.2f} s\"\n",
        "        print(template.format(tf.reduce_mean(generator_loss_log), tf.reduce_mean(discriminator_loss_log), epoch_time))\n",
        "\n",
        "      if (epoch + 1) % 10 == 0 or epoch == num_epochs - 1:\n",
        "        save_path = \"/content/drive/My Drive/Text to Image/lr_results/epoch_\" + str(epoch + 1)\n",
        "        temp_embeddings = None\n",
        "        for _, embeddings in train_ds:\n",
        "          temp_embeddings = embeddings.numpy()\n",
        "          break\n",
        "        if os.path.exists(save_path) == False:\n",
        "          os.makedirs(save_path)\n",
        "        temp_batch_size = 10\n",
        "        temp_z_noise = tf.random.normal((temp_batch_size, z_dim))\n",
        "        temp_embedding_batch = temp_embeddings[0:temp_batch_size]\n",
        "        fake_images, _ = self.stage1_generator([temp_embedding_batch, temp_z_noise])\n",
        "        for i, image in enumerate(fake_images):\n",
        "          image = 127.5*image + 127.5\n",
        "          image = image.numpy().astype('uint8')\n",
        "          cv2.imwrite(save_path + \"/gen_%d.png\"%(i), image)\n",
        "      \n",
        "        self.stage1_generator.save_weights(\"/content/drive/My Drive/Text to Image/lr_model_checkpoints/stage1_generator_\" + str(epoch + 1) + \".ckpt\")\n",
        "        self.stage1_discriminator.save_weights(\"/content/drive/My Drive/Text to Image/lr_model_checkpoints/stage1_discriminator_\" + str(epoch + 1) + \".ckpt\")\n",
        "\n",
        "    def generate_image(self, embedding):\n",
        "      self.stage1_generator.compile(loss = \"mse\", optimizer = \"adam\")\n",
        "      self.stage1_generator.load_weights(\"/content/drive/My Drive/Text to Image/lr_model_checkpoints/stage1_generator_600.ckpt\").expect_partial()\n",
        "      z_noise = tf.random.normal((batch_size, z_dim))\n",
        "      generated_image = self.stage1_generator([embedding, z_noise])\n",
        "      return generated_image"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ia8-xUroxrAp"
      },
      "source": [
        "dataset = Dataset(image_size = (64, 64))\n",
        "train_ds = dataset.get_train_ds()\n",
        "test_ds = dataset.get_test_ds()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8Ac-PgbtyVVy"
      },
      "source": [
        "model = Stage1Model()\n",
        "model.train(train_ds)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dtvKCODqLitp"
      },
      "source": [
        "#Stage 2 StackGAN"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dhzzZriSwwyx"
      },
      "source": [
        "class ResidualBlock(tf.keras.layers.Layer):\n",
        "  def __init__(self):\n",
        "    super(ResidualBlock, self).__init__()\n",
        "    self.conv1 = tf.keras.layers.Conv2D(filters = 128*4, kernel_size = 3, strides = 1, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.batchnorm1 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.conv2 = tf.keras.layers.Conv2D(filters = 128*4, kernel_size = 3, strides = 1, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.batchnorm2 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "\n",
        "    def call(self, I):\n",
        "      X = self.conv1(I)\n",
        "      X = self.batchnorm1(X)\n",
        "      X = tf.nn.relu(X)\n",
        "\n",
        "      X = self.conv2(X)\n",
        "      X = self.batchnorm2(X)\n",
        "      X = tf.nn.relu(X)\n",
        "      X = tf.keras.layers.Add()([X, I])\n",
        "      X = tf.nn.relu(X)\n",
        "      return X"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sGeTGuO-SGeG"
      },
      "source": [
        "class Stage2Generator(tf.keras.Model):\n",
        "  def __init__(self):\n",
        "    super(Stage2Generator, self).__init__()\n",
        "    self.canet = ConditioningAugmentation()\n",
        "    self.conv1 = tf.keras.layers.Conv2D(128, kernel_size = 3, strides = 1, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.conv2 = tf.keras.layers.Conv2D(256, kernel_size = 4, strides = 2, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.batchnorm1 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.conv3 = tf.keras.layers.Conv2D(512, kernel_size = 4, strides = 2, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.batchnorm2 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.conv4 = tf.keras.layers.Conv2D(512, kernel_size = 3, strides = 1, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.batchnorm3 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.resblock1 = ResidualBlock()\n",
        "    self.resblock2 = ResidualBlock()\n",
        "    self.resblock3 = ResidualBlock()\n",
        "    self.resblock4 = ResidualBlock()\n",
        "    self.upsamp1 = tf.keras.layers.UpSampling2D(size = (2, 2))\n",
        "    self.conv5 = tf.keras.layers.Conv2D(256, kernel_size = 3, strides = 1, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.batchnorm4 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.upsamp2 = tf.keras.layers.UpSampling2D(size = (2, 2))\n",
        "    self.conv6 = tf.keras.layers.Conv2D(filters = 128, kernel_size = 3, strides = 1, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.batchnorm5 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.upsamp3 = tf.keras.layers.UpSampling2D(size = (2, 2))\n",
        "    self.conv7 = tf.keras.layers.Conv2D(filters = 64, kernel_size = 3, strides = 1, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.batchnorm6 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.upsamp4 = tf.keras.layers.UpSampling2D(size = (2, 2))\n",
        "    self.conv8 = tf.keras.layers.Conv2D(filters = 32, kernel_size = 3, strides = 1, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.batchnorm7 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.conv9 = tf.keras.layers.Conv2D(filters = 3, kernel_size = 3, strides = 1, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "  \n",
        "  def call(self, inputs):\n",
        "    E, I = inputs\n",
        "    C, phi = self.canet(E)\n",
        "\n",
        "    X = self.conv1(I)\n",
        "    X = tf.nn.relu(X)\n",
        "    \n",
        "    X = self.conv2(X)\n",
        "    X = self.batchnorm1(X)\n",
        "    X = tf.nn.relu(X)\n",
        "\n",
        "    X = self.conv3(X)\n",
        "    X = self.batchnorm2(X)\n",
        "    X = tf.nn.relu(X)\n",
        "\n",
        "    C = K.expand_dims(C, axis = 1)\n",
        "    C = K.expand_dims(C, axis = 1)\n",
        "    C = K.tile(C, [1, 16, 16, 1])\n",
        "    J = K.concatenate([C, X], axis = 3)\n",
        "\n",
        "    X = self.conv4(X)\n",
        "    X = self.batchnorm3(X)\n",
        "    X = tf.nn.relu(X)\n",
        "\n",
        "    X = self.resblock1(X)\n",
        "    X = self.resblock2(X)\n",
        "    X = self.resblock3(X)\n",
        "    X = self.resblock4(X)\n",
        "\n",
        "    X = self.upsamp1(X)\n",
        "    X = self.conv5(X)\n",
        "    X = self.batchnorm4(X)\n",
        "    X = tf.nn.relu(X)\n",
        "    \n",
        "    X = self.upsamp2(X)\n",
        "    X = self.conv6(X)\n",
        "    X = self.batchnorm5(X)\n",
        "    X = tf.nn.relu(X)\n",
        "    \n",
        "    X = self.upsamp3(X)\n",
        "    X = self.conv7(X)\n",
        "    X = self.batchnorm6(X)\n",
        "    X = tf.nn.relu(X)\n",
        "    \n",
        "    X = self.upsamp4(X)\n",
        "    X = self.conv8(X)\n",
        "    X = self.batchnorm7(X)\n",
        "    X = tf.nn.relu(X)\n",
        "    \n",
        "    X = self.conv9(X)\n",
        "    return tf.nn.tanh(X), phi"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "T2ODvpdkSGY-"
      },
      "source": [
        "class Stage2Discriminator(tf.keras.Model):\n",
        "  def __init__(self):\n",
        "    super(Stage2Discriminator, self).__init__()\n",
        "    self.embed = EmbeddingCompressor()\n",
        "    self.reshape = tf.keras.layers.Reshape(target_shape = (1, 1, 128))\n",
        "    self.conv1 = tf.keras.layers.Conv2D(filters = 64, kernel_size = 4, strides = 2, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.conv2 = tf.keras.layers.Conv2D(filters = 128, kernel_size = 4, strides = 2, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.batchnorm1 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.conv3 = tf.keras.layers.Conv2D(filters = 256, kernel_size = 4, strides = 2, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.batchnorm2 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.conv4 = tf.keras.layers.Conv2D(filters = 512, kernel_size = 4, strides = 2, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.batchnorm3 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.conv5 = tf.keras.layers.Conv2D(filters = 1024, kernel_size = 4, strides = 2, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.batchnorm4 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.conv6 = tf.keras.layers.Conv2D(filters = 2048, kernel_size = 4, strides = 2, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.batchnorm5 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.conv7 = tf.keras.layers.Conv2D(filters = 1024, kernel_size = 1, strides = 1, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.batchnorm6 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.conv8 = tf.keras.layers.Conv2D(filters = 512, kernel_size = 1, strides = 1, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.batchnorm7 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.conv9 = tf.keras.layers.Conv2D(filters = 128, kernel_size = 1, strides = 1, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.batchnorm8 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.conv10 = tf.keras.layers.Conv2D(filters = 128, kernel_size = 3, strides = 1, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.batchnorm9 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.conv11 = tf.keras.layers.Conv2D(filters = 512, kernel_size = 3, strides = 1, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.batchnorm10 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.conv12 = tf.keras.layers.Conv2D(filters = 64*8, kernel_size = 1, strides = 1, padding = \"same\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "    self.batchnorm11 = tf.keras.layers.BatchNormalization(axis = -1, momentum = 0.99)\n",
        "    self.conv13 = tf.keras.layers.Conv2D(filters = 1, kernel_size = 4, strides = 1, padding = \"valid\", kernel_initializer = tf.keras.initializers.TruncatedNormal(stddev = 0.02))\n",
        "\n",
        "  def call(self, inputs):\n",
        "    I, E = inputs\n",
        "    T = self.embed(E)\n",
        "    T = self.reshape(T)\n",
        "    T = tf.tile(T, (1, 4, 4, 1))\n",
        "\n",
        "    X = self.conv1(I)\n",
        "    X = tf.nn.leaky_relu(X)\n",
        "\n",
        "    X = self.conv2(X)\n",
        "    X = self.batchnorm1(X)\n",
        "    X = tf.nn.leaky_relu(X)\n",
        "    \n",
        "    X = self.conv3(X)\n",
        "    X = self.batchnorm2(X)\n",
        "    X = tf.nn.leaky_relu(X)\n",
        "    \n",
        "    X = self.conv4(X)\n",
        "    X = self.batchnorm3(X)\n",
        "    X = tf.nn.leaky_relu(X)\n",
        "    \n",
        "    X = self.conv5(X)\n",
        "    X = self.batchnorm4(X)\n",
        "    X = tf.nn.leaky_relu(X)\n",
        "   \n",
        "    X = self.conv6(X)\n",
        "    X = self.batchnorm5(X)\n",
        "    X = tf.nn.leaky_relu(X)\n",
        "    \n",
        "    X = self.conv7(X)\n",
        "    X = self.batchnorm6(X)\n",
        "    X = tf.nn.leaky_relu(X)\n",
        "    \n",
        "    X = self.conv8(X)\n",
        "    X = self.batchnorm7(X)\n",
        "\n",
        "    Y = self.conv9(X)\n",
        "    Y = self.batchnorm8(Y)\n",
        "    Y = tf.nn.leaky_relu(Y)\n",
        "\n",
        "    Y = self.conv10(Y)\n",
        "    Y = self.batchnorm9(Y)\n",
        "    Y = tf.nn.leaky_relu(Y)\n",
        "\n",
        "    Y = self.conv11(Y)\n",
        "    Y = self.batchnorm10(Y)\n",
        "\n",
        "    A = tf.keras.layers.Add()([X, Y])\n",
        "    A = tf.nn.leaky_relu(A)\n",
        "\n",
        "    merged_input = tf.keras.layers.concatenate([A, T])\n",
        "\n",
        "    Z = self.conv12(merged_input)\n",
        "    Z = self.batchnorm11(Z)\n",
        "    Z = tf.nn.leaky_relu(Z)\n",
        "    \n",
        "    Z = self.conv13(Z)\n",
        "    return tf.squeeze(Z)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "p9aDTZqZ5RVT"
      },
      "source": [
        "class Stage2Model(tf.keras.Model):\n",
        "  def __init__(self):\n",
        "    super(Stage2Model, self).__init__()\n",
        "    self.stage1_generator = Stage1Generator()\n",
        "    self.stage1_generator.compile(loss = \"mse\", optimizer = \"adam\")\n",
        "    self.stage1_generator.load_weights(\"/content/drive/My Drive/Text to Image/lr_model_checkpoints/stage1_generator_600.ckpt\").expect_partial()\n",
        "    \n",
        "    self.stage2_generator = Stage2Generator()\n",
        "    self.stage2_discriminator = Stage2Discriminator()\n",
        "    self.stage2_generator.compile(loss = \"mse\", optimizer = \"adam\")\n",
        "    self.stage2_generator.load_weights(\"/content/drive/My Drive/Text to Image/hr_model_checkpoints/stage2_generator_170.ckpt\").expect_partial()\n",
        "    self.stage2_discriminator.compile(loss = \"mse\", optimizer = \"adam\")\n",
        "    self.stage2_discriminator.load_weights(\"/content/drive/My Drive/Text to Image/hr_model_checkpoints/stage2_discriminator_170.ckpt\").expect_partial()\n",
        "    \n",
        "  def train(self, train_ds, batch_size = 64, num_epochs = 1200, z_dim = 100, stage1_generator_lr = 0.0001, stage1_discriminator_lr = 0.0001):\n",
        "    generator_optimizer = tf.keras.optimizers.Adam(lr = stage1_generator_lr, beta_1 = 0.5, beta_2 = 0.999)\n",
        "    discriminator_optimizer = tf.keras.optimizers.Adam(lr = stage1_discriminator_lr, beta_1 = 0.5, beta_2 = 0.999)\n",
        "    \n",
        "    for epoch in range(num_epochs):\n",
        "      print(\"Epoch %d/%d:\\n [\"%(epoch + 1, num_epochs), end = \"\")\n",
        "      start_time = time.time()\n",
        "      if epoch % 100 == 0:\n",
        "        K.set_value(generator_optimizer.learning_rate, generator_optimizer.learning_rate / 2)\n",
        "        K.set_value(discriminator_optimizer.learning_rate, discriminator_optimizer.learning_rate / 2)\n",
        "    \n",
        "      generator_loss_log = []\n",
        "      discriminator_loss_log = []\n",
        "      steps_per_epoch = 125\n",
        "      batch_iter = iter(train_ds)\n",
        "      for i in range(steps_per_epoch):\n",
        "        if i % 5 == 0:\n",
        "          print(\"=\", end = \"\")\n",
        "        hr_image_batch, embedding_batch = next(batch_iter)\n",
        "        z_noise = tf.random.normal((batch_size, z_dim))\n",
        "\n",
        "        mismatched_images = tf.roll(hr_image_batch, shift = 1, axis = 0)\n",
        "\n",
        "        real_labels = tf.random.uniform(shape = (batch_size, ), minval = 0.9, maxval = 1.0)\n",
        "        fake_labels = tf.random.uniform(shape = (batch_size, ), minval = 0.0, maxval = 0.1)\n",
        "        mismatched_labels = tf.random.uniform(shape = (batch_size, ), minval = 0.0, maxval = 0.1)\n",
        "\n",
        "        with tf.GradientTape() as generator_tape, tf.GradientTape() as discriminator_tape:\n",
        "          lr_fake_images, _ = self.stage1_generator([embedding_batch, z_noise])\n",
        "          hr_fake_images, phi = self.stage2_generator([embedding_batch, lr_fake_images])\n",
        "          real_logits = self.stage2_discriminator([hr_image_batch, embedding_batch])\n",
        "          fake_logits = self.stage2_discriminator([hr_fake_images, embedding_batch])\n",
        "          mismatched_logits = self.stage2_discriminator([mismatched_images, embedding_batch])\n",
        "          \n",
        "          l_sup = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(real_labels, fake_logits))\n",
        "          l_klreg = KL_loss(tf.random.normal((phi.shape[0], phi.shape[1])), phi)\n",
        "          generator_loss = l_sup + 2.0*l_klreg\n",
        "          \n",
        "          l_real = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(real_labels, real_logits))\n",
        "          l_fake = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(fake_labels, fake_logits))\n",
        "          l_mismatched = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(mismatched_labels, mismatched_logits))\n",
        "          discriminator_loss = 0.5*tf.add(l_real, 0.5*tf.add(l_fake, l_mismatched))\n",
        "        \n",
        "        generator_gradients = generator_tape.gradient(generator_loss, self.stage2_generator.trainable_variables)\n",
        "        discriminator_gradients = discriminator_tape.gradient(discriminator_loss, self.stage2_discriminator.trainable_variables)\n",
        "        \n",
        "        generator_optimizer.apply_gradients(zip(generator_gradients, self.stage2_generator.trainable_variables))\n",
        "        discriminator_optimizer.apply_gradients(zip(discriminator_gradients, self.stage2_discriminator.trainable_variables))\n",
        "        \n",
        "        generator_loss_log.append(generator_loss)\n",
        "        discriminator_loss_log.append(discriminator_loss)\n",
        "        \n",
        "      end_time = time.time()\n",
        "\n",
        "      if epoch % 1 == 0:\n",
        "        epoch_time = end_time - start_time\n",
        "        template = \"] - generator_loss: {:.4f} - discriminator_loss: {:.4f} - epoch_time: {:.2f} s\"\n",
        "        print(template.format(tf.reduce_mean(generator_loss_log), tf.reduce_mean(discriminator_loss_log), epoch_time))\n",
        "\n",
        "      if (epoch + 1) % 10 == 0 or epoch == num_epochs - 1:\n",
        "        save_path = \"/content/drive/My Drive/Text to Image/hr_results/epoch_\" + str(epoch + 1)\n",
        "        temp_embeddings = None\n",
        "        for _, embeddings in train_ds:\n",
        "          temp_embeddings = embeddings.numpy()\n",
        "          break\n",
        "        if os.path.exists(save_path) == False:\n",
        "          os.makedirs(save_path)\n",
        "        temp_batch_size = 10\n",
        "        temp_z_noise = tf.random.normal((temp_batch_size, z_dim))\n",
        "        temp_embedding_batch = temp_embeddings[0:temp_batch_size]\n",
        "        fake_images, _ = self.stage1_generator([temp_embedding_batch, temp_z_noise])\n",
        "        for i, image in enumerate(fake_images):\n",
        "          image = 127.5*image + 127.5\n",
        "          image = image.numpy().astype('uint8')\n",
        "          cv2.imwrite(save_path + \"/gen_%d.png\"%(i), image)\n",
        "        self.stage2_generator.save_weights(\"/content/drive/My Drive/Text to Image/hr_model_checkpoints/stage2_generator_\" + str(epoch + 1) + \".ckpt\")\n",
        "        self.stage2_discriminator.save_weights(\"/content/drive/My Drive/Text to Image/hr_model_checkpoints/stage2_discriminator_\" + str(epoch + 1) + \".ckpt\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gN16Hw8gHohV"
      },
      "source": [
        "hr_dataset = Dataset(image_size = (256, 256), batch_size = 64)\n",
        "train_ds = hr_dataset.get_train_ds()\n",
        "test_ds = hr_dataset.get_test_ds()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zHSCSFVu9RqS"
      },
      "source": [
        "stage2_model = Stage2Model()\n",
        "stage2_model.train(train_ds)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}