cgan-map-generation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_kg_hide-output": true,
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# This Python 3 environment comes with many helpful analytics libraries installed\n",
    "# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n",
    "# For example, here's several helpful packages to load\n",
    "\n",
    "import numpy as np # linear algebra\n",
    "import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n",
    "\n",
    "# Input data files are available in the read-only \"../input/\" directory\n",
    "# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n",
    "\n",
    "import os\n",
    "for dirname, _, filenames in os.walk('/kaggle/input'):\n",
    "    for filename in filenames:\n",
    "        print(os.path.join(dirname, filename))\n",
    "\n",
    "# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n",
    "# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "from tqdm import tqdm\n",
    "from PIL import Image\n",
    "from IPython import display\n",
    "\n",
    "import tensorflow as tf\n",
    "from tensorflow.keras import models, layers, initializers, losses, optimizers\n",
    "from tensorflow.keras.utils import plot_model\n",
    "from tensorflow.keras.preprocessing.image import img_to_array, load_img"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "root_dir = \"/kaggle/input/pix2pix-dataset/maps/maps\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "folders = os.listdir(root_dir)\n",
    "folders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "folders = [\"train\", \"val\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_image_paths = []\n",
    "\n",
    "for file in tqdm(os.listdir(os.path.join(root_dir, folders[0]))):\n",
    "    train_image_paths.append(os.path.join(root_dir, folders[0], file))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "valid_image_paths = []\n",
    "\n",
    "for file in tqdm(os.listdir(os.path.join(root_dir, folders[1]))):\n",
    "    valid_image_paths.append(os.path.join(root_dir, folders[1], file))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Train Images:\", len(train_image_paths))\n",
    "print(\"Valid Images:\", len(valid_image_paths))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_image(path):\n",
    "    img = load_img(path)\n",
    "    img = img_to_array(img)\n",
    "    img = tf.cast(img, tf.float32)\n",
    "    \n",
    "    width = tf.shape(img)[1] // 2\n",
    "    \n",
    "    image = img[:, :width, :]\n",
    "    mask = img[:, width:, :]\n",
    "    \n",
    "    image = (tf.image.resize(image, (256, 256))) / 255\n",
    "    mask = (tf.image.resize(mask, (256, 256))) / 255\n",
    "\n",
    "    return image, mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_images = []\n",
    "train_masks = []\n",
    "\n",
    "for i, path in tqdm(enumerate(train_image_paths)):\n",
    "    image, mask = load_image(path)\n",
    "    \n",
    "    train_images.append(image)\n",
    "    train_masks.append(mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_images = np.array(train_images)\n",
    "train_masks = np.array(train_masks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "valid_images = []\n",
    "valid_masks = []\n",
    "\n",
    "for i, path in tqdm(enumerate(valid_image_paths)):\n",
    "    image, mask = load_image(path)\n",
    "    \n",
    "    valid_images.append(image)\n",
    "    valid_masks.append(mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "valid_images = np.array(valid_images)\n",
    "valid_masks = np.array(valid_masks)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_image_and_mask(images, masks):\n",
    "    plt.figure(figsize=(10, 6))\n",
    "\n",
    "    for i in range(8):\n",
    "        idx = np.random.randint(len(images))\n",
    "        plt.subplot(3, 8, i + 1)\n",
    "        plt.imshow(images[idx], cmap=\"gray\")\n",
    "        plt.axis(\"off\")\n",
    "        plt.title(train_image_paths[idx].split(\"/\")[-1])\n",
    "\n",
    "        plt.subplot(3, 8, i + 9)\n",
    "        plt.imshow(masks[idx], cmap=\"gray\")\n",
    "        plt.axis(\"off\")\n",
    "        plt.title(\"Mask\")\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_image_and_mask(train_images, train_masks)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def encoder_block(input_layer, num_filters):\n",
    "    conv = layers.Conv2D(num_filters, kernel_size=(3, 3), strides=(2, 2), padding=\"same\", kernel_initializer=\"he_normal\", use_bias=False)(input_layer)\n",
    "    conv = layers.LeakyReLU(0.2)(conv)\n",
    "    conv = layers.BatchNormalization()(conv)\n",
    "    return conv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def decoder_block(input_layer, skip_features, num_filters):\n",
    "    uconv = layers.Conv2DTranspose(num_filters, kernel_size=(4, 4), strides=(2, 2), padding=\"same\", kernel_initializer=\"he_normal\", use_bias=False)(input_layer)\n",
    "    uconv = layers.LeakyReLU(0.2)(uconv)\n",
    "    uconv = layers.BatchNormalization()(uconv)\n",
    "    uconv = layers.ReLU()(uconv)\n",
    "    cnt = layers.concatenate([uconv, skip_features])\n",
    "    return cnt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g_input_layer = layers.Input(shape=(256, 256, 3))\n",
    "\n",
    "e1 = encoder_block(g_input_layer, 64)\n",
    "e2 = encoder_block(e1, 128)\n",
    "e3 = encoder_block(e2, 256)\n",
    "e4 = encoder_block(e3, 512)\n",
    "e5 = encoder_block(e4, 512)\n",
    "e6 = encoder_block(e5, 512)\n",
    "e7 = encoder_block(e6, 512)\n",
    "\n",
    "g_bottle = encoder_block(e7, 512)\n",
    "\n",
    "d1 = decoder_block(g_bottle, e7, 512)\n",
    "d2 = decoder_block(d1, e6, 512)\n",
    "d3 = decoder_block(d2, e5, 512)\n",
    "d4 = decoder_block(d3, e4, 512)\n",
    "d5 = decoder_block(d4, e3, 256)\n",
    "d6 = decoder_block(d5, e2, 128)\n",
    "d7 = decoder_block(d6, e1, 64)\n",
    "\n",
    "g_initializer = initializers.RandomNormal(stddev=0.02, seed=42)\n",
    "\n",
    "g_output_layer = layers.Conv2DTranspose(3, kernel_size=(4, 4), strides=(2, 2), padding=\"same\", kernel_initializer=g_initializer, activation=\"tanh\")(d7)\n",
    "\n",
    "generator = models.Model(inputs=g_input_layer, outputs=g_output_layer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_model(generator)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Discriminator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "d_input_layer1 = layers.Input(shape=(256, 256, 3))\n",
    "d_input_layer2 = layers.Input(shape=(256, 256, 3))\n",
    "\n",
    "d_cnt = layers.concatenate([d_input_layer1, d_input_layer2])\n",
    "\n",
    "de1 = encoder_block(d_cnt, 64)\n",
    "de2 = encoder_block(de1, 128)\n",
    "de3 = encoder_block(de2, 512)\n",
    "\n",
    "d_initializer = initializers.RandomNormal(stddev=0.02, seed=42)\n",
    "\n",
    "d_conv1 = layers.Conv2D(512, kernel_size=(4, 4), kernel_initializer=d_initializer, use_bias=False)(de3)\n",
    "d_batchnorm1 = layers.BatchNormalization()(d_conv1)\n",
    "d_leakyrelu1 = layers.LeakyReLU()(d_batchnorm1)\n",
    "\n",
    "d_output_layer = layers.Conv2D(1, kernel_size=4, kernel_initializer=d_initializer)(d_leakyrelu1)\n",
    "\n",
    "discriminator = models.Model(inputs=[d_input_layer1, d_input_layer2], outputs=d_output_layer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_model(discriminator)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "EPOCHS = 100\n",
    "NUM_TESTS = 16\n",
    "NOISE_DIM = 100\n",
    "BUFFER_SIZE = len(train_images)\n",
    "BATCH_SIZE = 32\n",
    "LR_G = 0.0004\n",
    "LR_D = 0.0004\n",
    "BETA_G = 0.5\n",
    "BETA_D = 0.5\n",
    "LAMBDA = 100\n",
    "\n",
    "generator_optimizer = optimizers.Adam(learning_rate=LR_G, beta_1=BETA_G)\n",
    "discriminator_optimizer = optimizers.Adam(learning_rate=LR_D, beta_1=BETA_D)\n",
    "cross_entropy = losses.BinaryCrossentropy(from_logits=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = (tf.cast(train_images, tf.float32), tf.cast(train_masks, tf.float32))\n",
    "data = tf.data.Dataset.from_tensor_slices(dataset).batch(BATCH_SIZE, drop_remainder=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generator_loss(real_output, fake_output, image):\n",
    "    gen_loss = cross_entropy(tf.ones_like(real_output), real_output)\n",
    "    l1_loss = tf.reduce_mean(tf.abs(image - fake_output))\n",
    "    total = gen_loss + (LAMBDA * l1_loss)\n",
    "    return total"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_generated_images(epoch, test_input, real_input, output):\n",
    "    generated_images = generator(test_input, training=False)\n",
    "    \n",
    "    fig = plt.figure(figsize=(12, 8))\n",
    "    \n",
    "    plt.subplot(1, 3, 1)\n",
    "    plt.tight_layout()\n",
    "    plt.imshow(test_input[0] * 0.5 + 0.5)\n",
    "    plt.axis(\"off\")\n",
    "    \n",
    "    plt.subplot(1, 3, 2)\n",
    "    plt.tight_layout()\n",
    "    plt.imshow(real_input[0] * 0.5 + 0.5)\n",
    "    plt.axis(\"off\")\n",
    "    \n",
    "    plt.subplot(1, 3, 3)\n",
    "    plt.tight_layout()\n",
    "    plt.imshow(generated_images[0] * 0.5 + 0.5)\n",
    "    plt.axis(\"off\")\n",
    "        \n",
    "    if not os.path.exists(output):\n",
    "        os.makedirs(output)\n",
    "\n",
    "    fig.savefig(f\"{output}/cgan_{epoch}.png\")\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_step(inputs, target):\n",
    "    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:\n",
    "        fake_images = generator(inputs, training=True)\n",
    "        real_preds = discriminator([inputs, target], training=True)\n",
    "        fake_preds = discriminator([inputs, fake_images], training=True)\n",
    "        gen_loss = generator_loss(fake_preds, fake_images, target)\n",
    "        disc_loss = discriminator_loss(real_preds, fake_preds)\n",
    "        \n",
    "    generator_gradients = gen_tape.gradient(gen_loss, generator.trainable_variables)\n",
    "    generator_optimizer.apply_gradients(zip(generator_gradients, generator.trainable_variables))\n",
    "    \n",
    "    discriminator_gradients = disc_tape.gradient(disc_loss, discriminator.trainable_variables)\n",
    "    discriminator_optimizer.apply_gradients(zip(discriminator_gradients, discriminator.trainable_variables))\n",
    "    return gen_loss, disc_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_input, real_input = next(iter(data.take(1)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "g_loss = []\n",
    "d_loss = []\n",
    "\n",
    "for epoch in range(EPOCHS):\n",
    "    for image, mask in data:\n",
    "        gen_loss, disc_loss = train_step(image, mask)\n",
    "        \n",
    "    g_loss.append(gen_loss)\n",
    "    d_loss.append(disc_loss)\n",
    "    print(f\"[+] Epoch {epoch + 1}, Generator Loss: {gen_loss}, Discriminator Loss: {disc_loss}\")\n",
    "    show_generated_images(epoch, test_input, real_input, \"output\")\n",
    "    display.clear_output(wait=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 8))\n",
    "\n",
    "plt.plot(g_loss)\n",
    "plt.plot(d_loss)\n",
    "plt.legend([\"Generator Loss\", \"Discriminator Loss\"])\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_test = (tf.cast(valid_images, tf.float32), tf.cast(valid_masks, tf.float32))\n",
    "data_test = tf.data.Dataset.from_tensor_slices(dataset_test).batch(BATCH_SIZE, drop_remainder=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_predictions(test_images):\n",
    "    fig = plt.figure(figsize=(12, 8))\n",
    "    test_input, real_input = next(iter(test_images.take(3)))\n",
    "\n",
    "    generated_images = generator(test_input, training=False)\n",
    "\n",
    "    plt.subplot(1, 3, 1)\n",
    "    plt.tight_layout()\n",
    "    plt.imshow(test_input[0] * 0.5 + 0.5)\n",
    "    plt.axis(\"off\")\n",
    "    \n",
    "    plt.subplot(1, 3, 2)\n",
    "    plt.tight_layout()\n",
    "    plt.imshow(real_input[0] * 0.5 + 0.5)\n",
    "    plt.axis(\"off\")\n",
    "    \n",
    "    plt.subplot(1, 3, 3)\n",
    "    plt.tight_layout()\n",
    "    plt.imshow(generated_images[0] * 0.5 + 0.5)\n",
    "    plt.axis(\"off\")\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_predictions(data_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}