GANGANStyle_celeb-A_dataset.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "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.6.6"
    },
    "colab": {
      "name": "GANGANStyle_celeb-A_dataset.ipynb",
      "private_outputs": true,
      "provenance": [],
      "include_colab_link": true
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/consequencesunintended/GANGANstyle/blob/master/GANGANStyle_celeb-A_dataset.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Bg_BVJZ8USy7"
      },
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import PIL\n",
        "import os\n",
        "import time\n",
        "import imageio\n",
        "import glob\n",
        "from IPython import display\n",
        "print (tf.__version__)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uDMwA3CBUSzC"
      },
      "source": [
        "from tensorflow.keras import layers"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6FZxrt5JUSzD"
      },
      "source": [
        "from tensorflow.keras.preprocessing.image import ImageDataGenerator"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MRaeGDr_USzE"
      },
      "source": [
        "DIMENSION = 64\n",
        "initial_g_dimension = 8\n",
        "noise_dim = 16*16\n",
        "BATCH_SIZE = 128\n",
        "EPOCHS = 500"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lHWIDBUgUSzE"
      },
      "source": [
        "# At the time of writing this notebook there is an issue with celeb_a tensorflow dataset which already has been reported\n",
        "# https://github.com/tensorflow/datasets/issues/1482\n",
        "# you can manually download the dataset and copied it in the right folder for the next line to work\n",
        "\n",
        "#dataset, info = tfds.load('celeb_a:2.*.*', with_info=True)\n",
        "\n",
        "# therefore for the time being instead enjoy some pet photos \n",
        "\n",
        "dataset, info = tfds.load('oxford_iiit_pet:3.*.*', with_info=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mAJ9lj3SUSzE"
      },
      "source": [
        "AUTOTUNE = tf.data.experimental.AUTOTUNE"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mhlHg9AUUSzF"
      },
      "source": [
        "def normalise(img):\n",
        "    img = img / 255.0\n",
        "    \n",
        "    return img"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sE_7fIi5USzF"
      },
      "source": [
        "def load_img(datapoint):\n",
        "    img = tf.image.resize(datapoint['image'], (DIMENSION, DIMENSION))\n",
        "    img = tf.image.rgb_to_grayscale(img)\n",
        "    img = tf.image.convert_image_dtype(img, tf.float32)\n",
        "    img = normalise(img)\n",
        "    \n",
        "    return img"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LUTrssklUSzF"
      },
      "source": [
        "train = dataset['train'].map(load_img, num_parallel_calls=tf.data.AUTOTUNE)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "o0zg5dKdUSzG"
      },
      "source": [
        "SHUFFLE_BUFFER_SIZE = 100\n",
        "train_dataset = train.cache().shuffle(SHUFFLE_BUFFER_SIZE).batch(BATCH_SIZE)\n",
        "train_dataset = train_dataset.prefetch(buffer_size=tf.data.AUTOTUNE)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "47TcIAEoUSzG"
      },
      "source": [
        "NUM_IMAGES = info.splits['train'].num_examples"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kXOv7_xBUSzG"
      },
      "source": [
        "for image in train_dataset.take(1):\n",
        "    print(\"Image shape: \", image.numpy().shape)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZgVNZlhSUSzH"
      },
      "source": [
        "sample_training_images = next(iter(train_dataset))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "b932veUuUSzH"
      },
      "source": [
        "plt.imshow(np.array(sample_training_images[0].numpy().reshape(DIMENSION,DIMENSION)), cmap='gray')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aMGv7HyVUSzH"
      },
      "source": [
        "def plotImages(images_arr):\n",
        "    fig, axes = plt.subplots(1, 5, figsize=(DIMENSION,DIMENSION))\n",
        "    axes = axes.flatten()\n",
        "    for img, ax in zip( images_arr, axes):\n",
        "        ax.imshow(img.numpy().reshape(DIMENSION, DIMENSION), cmap='gray')           \n",
        "        ax.axis('off')    \n",
        "    plt.tight_layout()\n",
        "    plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Yn84V0sAUSzI"
      },
      "source": [
        "plotImages(sample_training_images[:5])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hmGl2X7PUSzI"
      },
      "source": [
        "def make_generator_model():\n",
        "    model = tf.keras.Sequential()\n",
        "    model.add(layers.Dense((DIMENSION//8)*(DIMENSION//8)*128, use_bias=False, input_shape=(16*16,)))\n",
        "    model.add(layers.BatchNormalization())\n",
        "    model.add(layers.LeakyReLU())\n",
        "\n",
        "    model.add(layers.Reshape((DIMENSION//8,DIMENSION//8,128)))\n",
        "    assert model.output_shape == (None, DIMENSION//8,DIMENSION//8,128)\n",
        "\n",
        "    model.add(layers.Conv2DTranspose(128, (5, 5), strides=(2, 2), padding='same', use_bias=False))\n",
        "    assert model.output_shape == (None, DIMENSION//4, DIMENSION//4, 128)\n",
        "    model.add(layers.BatchNormalization())\n",
        "    model.add(layers.LeakyReLU())\n",
        "\n",
        "    model.add(layers.Conv2DTranspose(128, (5, 5), strides=(2, 2), padding='same', use_bias=False))\n",
        "    assert model.output_shape == (None, DIMENSION//2, DIMENSION//2, 128)\n",
        "    model.add(layers.BatchNormalization())\n",
        "    model.add(layers.LeakyReLU())\n",
        "    \n",
        "    model.add(layers.Conv2DTranspose(128, (5, 5), strides=(2, 2), padding='same', use_bias=False))\n",
        "    assert model.output_shape == (None, DIMENSION, DIMENSION, 128)\n",
        "    model.add(layers.BatchNormalization())\n",
        "    model.add(layers.LeakyReLU())\n",
        "\n",
        "    model.add(layers.Conv2DTranspose(1, (5, 5), strides=(1, 1), padding='same', use_bias=False, activation='tanh'))\n",
        "    assert model.output_shape == (None, DIMENSION, DIMENSION, 1)\n",
        "\n",
        "    return model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7Iz6rIpKUSzI"
      },
      "source": [
        "generator = make_generator_model()\n",
        "\n",
        "noise = tf.random.normal([1, noise_dim])\n",
        "generated_image = generator(noise, training=False)\n",
        "\n",
        "plt.imshow(generated_image[0, :, :, 0], cmap='gray')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8SM9OqDwUSzJ"
      },
      "source": [
        "def make_discriminator_model():\n",
        "    model = tf.keras.Sequential()\n",
        "    model.add(layers.Conv2D(64, (5, 5), strides=(2, 2), padding='same',\n",
        "                                     input_shape=[DIMENSION, DIMENSION, 1]))\n",
        "    model.add(layers.LeakyReLU())\n",
        "    model.add(layers.Dropout(0.3))\n",
        "\n",
        "    model.add(layers.Conv2D(64, (5, 5), strides=(2, 2), padding='same'))\n",
        "    model.add(layers.LeakyReLU())\n",
        "    model.add(layers.Dropout(0.3))\n",
        "    \n",
        "    model.add(layers.Conv2D(64, (5, 5), strides=(2, 2), padding='same'))\n",
        "    model.add(layers.LeakyReLU())\n",
        "    model.add(layers.Dropout(0.3))\n",
        "\n",
        "    model.add(layers.Flatten())\n",
        "    model.add(layers.Dense(1))\n",
        "\n",
        "    return model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2a5gJ5UTUSzJ"
      },
      "source": [
        "discriminator = make_discriminator_model()\n",
        "decision = discriminator(generated_image)\n",
        "print (decision)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rL54-sPLUSzK"
      },
      "source": [
        "cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v1L9dG83USzK"
      },
      "source": [
        "def discriminator_loss(real_output, fake_output):\n",
        "    real_loss = cross_entropy(tf.ones_like(real_output), real_output)\n",
        "    fake_loss = cross_entropy(tf.zeros_like(fake_output), fake_output)\n",
        "    total_loss = real_loss + fake_loss\n",
        "    return total_loss"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "n6Ka0qsqUSzK"
      },
      "source": [
        "def generator_loss(fake_output):\n",
        "    return cross_entropy(tf.ones_like(fake_output), fake_output)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z53EMbQwUSzL"
      },
      "source": [
        "generator_optimizer = tf.keras.optimizers.Adam(1e-4,)\n",
        "discriminator_optimizer = tf.keras.optimizers.Adam(1e-4)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "q7P4kgVrUSzL"
      },
      "source": [
        "checkpoint_dir = './training_checkpoints'\n",
        "checkpoint_prefix = os.path.join(checkpoint_dir, \"ckpt\")\n",
        "checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer,\n",
        "                                 discriminator_optimizer=discriminator_optimizer,\n",
        "                                 generator=generator,\n",
        "                                 discriminator=discriminator)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IqgKSY_tUSzL"
      },
      "source": [
        "num_examples_to_generate = 16\n",
        "\n",
        "seed = tf.random.normal([num_examples_to_generate, noise_dim])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "a2pKaz4PUSzL"
      },
      "source": [
        "def train_step(images):\n",
        "    noise = tf.random.normal([BATCH_SIZE, noise_dim])\n",
        "\n",
        "    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:\n",
        "        generated_images = generator(noise, training=True)\n",
        "\n",
        "        real_output = discriminator(images, training=True)\n",
        "        fake_output = discriminator(generated_images, training=True)\n",
        "\n",
        "        gen_loss = generator_loss(fake_output)\n",
        "        disc_loss = discriminator_loss(real_output, fake_output)\n",
        "\n",
        "    gradients_of_generator = gen_tape.gradient(gen_loss, generator.trainable_variables)\n",
        "    gradients_of_discriminator = disc_tape.gradient(disc_loss, discriminator.trainable_variables)\n",
        "\n",
        "    generator_optimizer.apply_gradients(zip(gradients_of_generator, generator.trainable_variables))\n",
        "    discriminator_optimizer.apply_gradients(zip(gradients_of_discriminator, discriminator.trainable_variables))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "63WcEOetUSzM"
      },
      "source": [
        "def generate_and_save_images(model, epoch, test_input):\n",
        "    predictions = model(test_input, training=False)\n",
        "\n",
        "    fig = plt.figure(figsize=(8,8))\n",
        "\n",
        "    for i in range(predictions.shape[0]):\n",
        "        plt.subplot(4, 4, i+1)\n",
        "        plt.imshow(predictions[i, :, :, 0] * 127.5 + 127.5, cmap='gray')\n",
        "        plt.axis('off')\n",
        "\n",
        "    plt.savefig('image_at_epoch_{:04d}.png'.format(epoch))\n",
        "    plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "T38et9bYUSzM"
      },
      "source": [
        "STEP_SIZE = NUM_IMAGES // BATCH_SIZE\n",
        "def train(epochs):\n",
        "    for epoch in range(epochs):\n",
        "        start = time.time()\n",
        "\n",
        "        for _ in range(STEP_SIZE):\n",
        "             train_step(next(iter(train_dataset)))\n",
        "\n",
        "        display.clear_output(wait=True)\n",
        "        generate_and_save_images(generator,\n",
        "                                 epoch + 1,\n",
        "                                 seed)\n",
        "\n",
        "        if (epoch + 1) % 15 == 0:\n",
        "            checkpoint.save(file_prefix = checkpoint_prefix)\n",
        "\n",
        "        print ('Time for epoch {} is {} sec'.format(epoch + 1, time.time()-start))\n",
        "\n",
        "    display.clear_output(wait=True)\n",
        "    generate_and_save_images(generator,\n",
        "                           epochs,\n",
        "                           seed)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rcE1htjPUSzM"
      },
      "source": [
        "%%time\n",
        "train(EPOCHS)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wxk1TsReUSzN"
      },
      "source": [
        "checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tWKLW2wXUSzN"
      },
      "source": [
        "def display_image(epoch_no):\n",
        "      return PIL.Image.open('image_at_epoch_{:04d}.png'.format(epoch_no))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "t1_sbAX8USzN"
      },
      "source": [
        "display_image(EPOCHS)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fQEb6HquUSzN"
      },
      "source": [
        "anim_file = 'dcgan.gif'\n",
        "\n",
        "with imageio.get_writer(anim_file, mode='I') as writer:\n",
        "    filenames = glob.glob('image*.png')\n",
        "    filenames = sorted(filenames)\n",
        "    last = -1\n",
        "    for i,filename in enumerate(filenames):\n",
        "        frame = 2*(i**0.5)\n",
        "        if round(frame) > round(last):\n",
        "            last = frame\n",
        "        else:\n",
        "            continue\n",
        "            \n",
        "        image = imageio.imread(filename)\n",
        "        writer.append_data(image)\n",
        "        image = imageio.imread(filename)\n",
        "        writer.append_data(image)\n",
        "\n",
        "import IPython\n",
        "\n",
        "if IPython.version_info > (6,2,0,''):\n",
        "    display.Image(filename=anim_file)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}