WGAN-GP.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import glob\n",
    "import time\n",
    "import cv2\n",
    "\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import imageio\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import display, Image\n",
    "\n",
    "import tensorflow as tf\n",
    "from tensorflow.keras import Model, Sequential\n",
    "from tensorflow.keras.preprocessing import image\n",
    "from tensorflow.keras.optimizers import Adam, RMSprop\n",
    "from tensorflow.keras.losses import BinaryCrossentropy\n",
    "from tensorflow.keras.layers import (Dense, Conv2DTranspose, Conv2D,\n",
    "                                     BatchNormalization,MaxPool2D,AveragePooling2D,\n",
    "                                     LeakyReLU, Dropout, Reshape, Flatten,\n",
    "                                     ELU,Activation,Input,UpSampling2D)\n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "from functools import partial\n",
    "from tqdm import tqdm\n",
    "from typing import List , Union , Tuple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_img(path : str ,\n",
    "             width : int = 1200) -> tf.Tensor:\n",
    "    \"\"\" loading an image with resize into shape of (1200,1200) \"\"\"\n",
    "    img = cv2.imread(path).astype(np.float32)\n",
    "    \n",
    "    img = tf.constant(img/127.5 - 1.) \n",
    "    shape_dst = int(np.min(img.shape[:2])*0.8)\n",
    "    oh = (img.shape[0] - shape_dst) // 2\n",
    "    ow = (img.shape[1] - shape_dst) // 2\n",
    "    new_size = (width , width)\n",
    "\n",
    "    img = tf.expand_dims(img[oh:oh + shape_dst, ow:ow + shape_dst] , axis = 0)\n",
    "    img = tf.image.resize(img, new_size, method = tf.image.ResizeMethod.BILINEAR)\n",
    "    return img[0]\n",
    "\n",
    "def load_all_img(ImageNamePath : List[str]):\n",
    "    for idx in range(len(ImageNamePath)):\n",
    "        support = load_img(ImageNamePath[idx])\n",
    "        yield support\n",
    "        \n",
    "        \n",
    "def rnd_crop_img(img : tf.Tensor ,\n",
    "                 win_size : int = 256) -> tf.Tensor :\n",
    "    \"\"\" Random Cropping Image into size (256 , 256) \"\"\"\n",
    "    CH = 3 # color_channels\n",
    "    new_size = (win_size , win_size , CH)\n",
    "    return tf.image.random_crop(img , new_size)\n",
    "\n",
    "\n",
    "def gen_task(img_loader : tf.data.Dataset , \n",
    "             win_size : int = 256) -> (tf.Tensor , tf.Tensor):\n",
    "    ''' Creating Dataset '''\n",
    "    for img in img_loader.skip(64).take(1).repeat(BATCH_SIZE*20):\n",
    "        '''Support processing'''\n",
    "        support_= rnd_crop_img(img,win_size)\n",
    "        word = tf.convert_to_tensor(word_embedding['WordEmbedding'][64])\n",
    "        yield (support_,word)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Generator(Model):\n",
    "    def __init__(self, z_dim ):\n",
    "        super(Generator, self).__init__()\n",
    "\n",
    "        self.model = Sequential()\n",
    "\n",
    "        # [z_dim] => [8,8,64]\n",
    "        self.model.add(Dense(8*8*64,\n",
    "                             activation = tf.nn.leaky_relu ,\n",
    "                             use_bias = False , \n",
    "                             input_shape = (z_dim ,)))\n",
    "        self.model.add(Reshape((8, 8,64)))\n",
    "\n",
    "        # [8,8,64] => [128,128,64]\n",
    "        for ii in range(int(np.log2(128)-np.log2(8))):\n",
    "            self.model.add(UpSampling2D(size = (2, 2), \n",
    "                                        interpolation = 'bilinear'))\n",
    "            self.model.add(Conv2D(filters = 64 ,\n",
    "                                  kernel_size = 3 ,\n",
    "                                  padding = 'same' ,\n",
    "                                  activation = tf.nn.leaky_relu ,\n",
    "                                  use_bias = False))\n",
    "        \n",
    "        # [128,128,64] => [256,256,3]\n",
    "        self.model.add(UpSampling2D(size = (2, 2) ,\n",
    "                                    interpolation = 'bilinear'))\n",
    "        self.model.add(Conv2D(filters = CH ,\n",
    "                              kernel_size = 5 ,\n",
    "                              padding = 'same' , \n",
    "                              activation = 'tanh' ,\n",
    "                              use_bias = False))\n",
    "\n",
    "    def call(self, x):\n",
    "        return self.model(x)\n",
    "\n",
    "class Discriminator(Model):\n",
    "    def __init__(self):\n",
    "        super(Discriminator, self).__init__()\n",
    "        # [256,256,3] => [128,128, 64]\n",
    "        self.conv0 = Conv2D(filters = 64 ,\n",
    "                            kernel_size = 5 ,\n",
    "                            padding = \"same\" , \n",
    "                            activation = tf.nn.leaky_relu ,\n",
    "                            input_shape = (W, W, CH))\n",
    "        \n",
    "        self.conv1dx = [Conv2D(filters = 64 ,\n",
    "                               kernel_size = 1 ,\n",
    "                               padding = \"same\" , \n",
    "                               activation = tf.nn.leaky_relu) for _ in range(4)]\n",
    "        self.convx = [Conv2D(filters = 64 ,\n",
    "                             kernel_size = 3,\n",
    "                             padding = \"same\" ,\n",
    "                             activation = tf.nn.leaky_relu ,\n",
    "                             use_bias = False) for _ in range(4)]\n",
    "        # [8,8, 64] => [1]\n",
    "        self.dense = Dense(1)\n",
    "\n",
    "    def call(self, x):\n",
    "        xc = self.conv0(x)\n",
    "        xl = MaxPool2D(pool_size = (2, 2))(xc)\n",
    "        xr = AveragePooling2D(pool_size = (2, 2))(x)\n",
    "        x1 = tf.concat((xl,xr),axis = -1)\n",
    "        for ii in range(4):\n",
    "            x1 = self.conv1dx[ii](x1)\n",
    "            xc = self.convx[ii](x1)\n",
    "            xl = MaxPool2D(pool_size = (2, 2))(xc)\n",
    "            xr = AveragePooling2D(pool_size = (2, 2))(x1)\n",
    "            x1 = tf.concat((xl,xr),axis = -1)\n",
    "        outs = self.dense(Flatten()(x1))\n",
    "        return outs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gradient_penalty(real_images : tf.Tensor , \n",
    "                     fake_images : tf.Tensor) -> tf.Tensor :\n",
    "    '''\n",
    "    @params : BATCH_SIZE : int \n",
    "    @params : W : int , width of image\n",
    "    @params : CH : int = 3 , color_channels\n",
    "    '''\n",
    "    epsilon = tf.random.uniform(shape = [BATCH_SIZE, W, W, CH], minval = 0.0, maxval = 1.0)\n",
    "    x_hat = epsilon * real_images + (1 - epsilon) * fake_images\n",
    "\n",
    "    with tf.GradientTape() as t:\n",
    "        t.watch(x_hat)\n",
    "        d_hat = discriminator(x_hat)\n",
    "    gradients = t.gradient(d_hat, x_hat)\n",
    "\n",
    "    g_norm = tf.sqrt(tf.reduce_sum(gradients ** 2, axis = [1, 2]))\n",
    "    gradient_penalty = tf.reduce_mean((g_norm - 1.0) ** 2)\n",
    "    return gradient_penalty\n",
    "\n",
    "\n",
    "def gan_loss(d_real_output, d_fake_output, train_loader, fake_images):\n",
    "    \"\"\"\n",
    "    d_loss -> loss of discriminator \n",
    "    g_loss -> loss of generator\n",
    "    \"\"\"\n",
    "    d_loss = tf.reduce_mean(d_fake_output) - tf.reduce_mean(d_real_output) + \\\n",
    "             gradient_penalty(train_loader[0], fake_images) * gp_weight\n",
    "    g_loss = tf.reduce_mean(-d_fake_output)\n",
    "    return d_loss, g_loss\n",
    "\n",
    "\n",
    "@tf.function\n",
    "def train_step(train_loader, generator, discriminator, g_optimizer, d_optimizer):\n",
    "    noise = train_loader\n",
    "    with tf.GradientTape() as g_tape, tf.GradientTape() as d_tape:\n",
    "        fake_images = generator(noise[1], training=True)\n",
    "        d_real_logits = discriminator(noise[0])\n",
    "        d_fake_logits = discriminator(fake_images)\n",
    "\n",
    "        d_loss, g_loss = gan_loss(d_real_logits, d_fake_logits, train_loader, fake_images)\n",
    "\n",
    "    g_gradients = g_tape.gradient(g_loss, generator.trainable_variables)\n",
    "    d_gradients = d_tape.gradient(d_loss, discriminator.trainable_variables)\n",
    "\n",
    "    g_optimizer.apply_gradients(zip(g_gradients, generator.trainable_variables))\n",
    "    d_optimizer.apply_gradients(zip(d_gradients, discriminator.trainable_variables))\n",
    "    return d_loss, g_loss\n",
    "\n",
    "\n",
    "def train(train_loader, epochs):\n",
    "    for epoch in range(epochs):\n",
    "        start = time.time()\n",
    "\n",
    "        for image_batch in train_loader:\n",
    "            d_loss, g_loss = train_step(image_batch, generator, discriminator, g_optimizer, d_optimizer)\n",
    "\n",
    "        # Produce images\n",
    "        print('Time for epoch {} is {} sec'.format(epoch + 1, time.time() - start))\n",
    "        print('discriminator loss: %.5f' % d_loss)\n",
    "        print('generator loss: %.5f' % g_loss)\n",
    "        generate_and_save_images(generator, epoch + 1, seed, save_dir)\n",
    "\n",
    "        # Save the model every 25 epochs\n",
    "        if (epoch + 1) % 25 == 0:\n",
    "            checkpoint.save(file_prefix=checkpoint_prefix)\n",
    "    # generating / saving after the final epoch\n",
    "    generate_and_save_images(generator, epochs, seed, save_dir)\n",
    "    checkpoint.save(file_prefix = checkpoint_prefix)\n",
    "    \n",
    "    \n",
    "def generate_and_save_images(model, epoch, train_loader, save_path):\n",
    "    predictions = model(train_loader)\n",
    "\n",
    "    fig = plt.figure(figsize=(4, 4))\n",
    "    for i in range(predictions.shape[0]):\n",
    "        plt.subplot(4, 4, i + 1)\n",
    "        plt.imshow((predictions[i, :, :, :].numpy()+1)*0.5)\n",
    "        plt.axis('off')\n",
    "\n",
    "    if (epoch + 1) % 10 == 0:\n",
    "        if epoch<=100:\n",
    "            plt.savefig(os.path.join(save_path, 'image_at_epoch_{:04d}.png'.format(epoch+1)))\n",
    "        elif (epoch + 1) % 100 == 0:\n",
    "            if epoch<=1000:\n",
    "                plt.savefig(os.path.join(save_path, 'image_at_epoch_{:04d}.png'.format(epoch+1)))\n",
    "            elif (epoch + 1) % 1000 == 0:\n",
    "                plt.savefig(os.path.join(save_path, 'image_at_epoch_{:04d}.png'.format(epoch+1)))\n",
    "    plt.show()\n",
    "    \n",
    "    \n",
    "def calculate_fid(real_embeddings, generated_embeddings):\n",
    "    from scipy.linalg import sqrtm\n",
    "    mu1, sigma1 = real_embeddings.mean(axis=0), np.cov(real_embeddings, rowvar=False)\n",
    "    mu2, sigma2 = generated_embeddings.mean(axis=0), np.cov(generated_embeddings,  rowvar=False)\n",
    "    ssdiff = np.sum((mu1 - mu2)**2.0)\n",
    "    covmean = sqrtm(sigma1.dot(sigma2))\n",
    "    if np.iscomplexobj(covmean):\n",
    "        covmean = covmean.real\n",
    "    fid = ssdiff + np.trace(sigma1 + sigma2 - 2.0 * covmean)\n",
    "    return fid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}