Pix2Pix-DepthEstimation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_kg_hide-output": true,
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import imageio\n",
    "import numpy as np\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore',category=FutureWarning)\n",
    "import tensorflow as tf\n",
    "import matplotlib.pyplot as plt\n",
    "from glob import glob\n",
    "import cv2\n",
    "import shutil\n",
    "tf.logging.set_verbosity(tf.logging.ERROR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Helpers():\n",
    "    \n",
    "    @staticmethod\n",
    "    def normalize(images):\n",
    "        return np.array(images)/127.5-1.0\n",
    "    \n",
    "    @staticmethod\n",
    "    def unnormalize(images):\n",
    "        return (0.5*np.array(images)+0.5)*255\n",
    "    \n",
    "    @staticmethod\n",
    "    def resize(image, size):\n",
    "        return np.array(cv2.resize(image, size))\n",
    "    \n",
    "    @staticmethod\n",
    "    def split_images(image, is_testing):\n",
    "        image = imageio.imread(image).astype(np.float)\n",
    "        _, width, _ = image.shape\n",
    "        half_width = int(width/2)\n",
    "        source_image = image[:, half_width:, :]\n",
    "        destination_image = image[:, :half_width, :]\n",
    "        source_image = Helpers.resize(source_image, (IMAGE_SIZE, IMAGE_SIZE))\n",
    "        destination_image = Helpers.resize(destination_image, (IMAGE_SIZE, IMAGE_SIZE))\n",
    "        if not is_testing and np.random.random() > 0.5:\n",
    "            source_image = np.fliplr(source_image)\n",
    "            destination_image = np.fliplr(destination_image)\n",
    "        return source_image, destination_image\n",
    "    \n",
    "    @staticmethod\n",
    "    def new_dir(path):\n",
    "        shutil.rmtree(path, ignore_errors=True)\n",
    "        os.makedirs(path, exist_ok=True)\n",
    "        \n",
    "    @staticmethod\n",
    "    def archive_output():\n",
    "        shutil.make_archive(\"output\", \"zip\", \"./output\")\n",
    "        \n",
    "    @staticmethod\n",
    "    def image_pairs(batch, is_testing):\n",
    "        source_images, destination_images = [], []\n",
    "        for image_path in batch:\n",
    "            source_image, destination_image = Helpers.split_images(image_path, is_testing)\n",
    "            source_images.append(source_image)\n",
    "            destination_images.append(destination_image)\n",
    "        return source_images, destination_images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_cell_guid": "79c7e3d0-c299-4dcb-8224-4455121ee9b0",
    "_kg_hide-output": false,
    "_uuid": "d629ff2d2480ee46fbb7e2d37f6b5fab8052498a"
   },
   "outputs": [],
   "source": [
    "# Requires following dataset structure:\n",
    "# dataset_name\n",
    "# └── dataset_name\n",
    "#     ├── testing\n",
    "#     │   └── ... (image files)\n",
    "#     ├── testing_raw\n",
    "#     │   ├── ... (image files)\n",
    "#     ├── training\n",
    "#     │   └── ... (image files)\n",
    "#     └── validation (optional)\n",
    "#         └── ... (image files)\n",
    "class DataLoader():\n",
    "    \n",
    "    def __init__(self, dataset_name=\"pix2pix-depth\"):\n",
    "        self.dataset_name = dataset_name\n",
    "        base_path = BASE_INPUT_PATH + self.dataset_name + \"/\" + self.dataset_name + \"/\"\n",
    "        self.training_path = base_path + \"training/\"\n",
    "        self.validation_path = base_path + \"validation/\"\n",
    "        self.testing_path = base_path + \"testing/\"\n",
    "        self.testing_raw_path = base_path + \"testing_raw/\"\n",
    "\n",
    "    def load_random_data(self, data_size, is_testing=False):\n",
    "        paths = glob(self.training_path+\"*\") if is_testing else glob(self.testing_path+\"*\")\n",
    "        source_images, destination_images = Helpers.image_pairs(np.random.choice(paths, size=data_size), is_testing)\n",
    "        return Helpers.normalize(source_images), Helpers.normalize(destination_images)\n",
    "\n",
    "    def yield_batch(self, batch_size, is_testing=False):\n",
    "        paths = glob(self.training_path+\"*\") if is_testing else glob(self.validation_path+\"*\")\n",
    "        for i in range(int(len(paths)/batch_size)-1):\n",
    "            batch = paths[i*batch_size:(i+1)*batch_size]\n",
    "            source_images, destination_images = Helpers.image_pairs(batch, is_testing)\n",
    "            yield Helpers.normalize(source_images), Helpers.normalize(destination_images)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "8ff0cb940babed054508cf32b3ef5d383302fa12"
   },
   "outputs": [],
   "source": [
    "# Model architecture from: https://phillipi.github.io/pix2pix/\n",
    "class Pix2Pix(): \n",
    "    \n",
    "    def __init__(self):\n",
    "        Helpers.new_dir(BASE_OUTPUT_PATH + \"training/\")\n",
    "        Helpers.new_dir(BASE_OUTPUT_PATH + \"training/losses/\")\n",
    "\n",
    "        self.image_shape = (IMAGE_SIZE, IMAGE_SIZE, IMAGE_CHANNELS)\n",
    "        self.data_loader = DataLoader()\n",
    "\n",
    "        patch = int(IMAGE_SIZE / 2**4)\n",
    "        self.disc_patch = (patch, patch, 1)\n",
    "\n",
    "        self.generator_filters = 64\n",
    "        self.discriminator_filters = 64\n",
    "        \n",
    "        optimizer = tf.keras.optimizers.Adam(LEARNING_RATE, BETA_1)\n",
    "\n",
    "        self.discriminator = self.discriminator()\n",
    "        self.discriminator.compile(loss=\"mse\", optimizer=optimizer, metrics=[\"accuracy\"])\n",
    "        self.generator = self.generator()\n",
    "\n",
    "        source_image = tf.keras.layers.Input(shape=self.image_shape)\n",
    "        destination_image = tf.keras.layers.Input(shape=self.image_shape)\n",
    "        generated_image = self.generator(destination_image)\n",
    "\n",
    "        self.discriminator.trainable = False\n",
    "        valid = self.discriminator([generated_image, destination_image])\n",
    "        self.combined = tf.keras.models.Model(inputs=[source_image, destination_image], outputs=[valid, generated_image])\n",
    "        self.combined.compile(loss=[\"mse\", \"mae\"], loss_weights=[1, 100], optimizer=optimizer)\n",
    "\n",
    "    def generator(self):\n",
    "        def conv2d(layer_input, filters, bn=True):\n",
    "            downsample = tf.keras.layers.Conv2D(filters, kernel_size=4, strides=2, padding=\"same\")(layer_input)\n",
    "            downsample = tf.keras.layers.LeakyReLU(alpha=LEAKY_RELU_ALPHA)(downsample)\n",
    "            if bn:\n",
    "                downsample = tf.keras.layers.BatchNormalization(momentum=BN_MOMENTUM)(downsample)\n",
    "            return downsample\n",
    "\n",
    "        def deconv2d(layer_input, skip_input, filters, dropout_rate=0):\n",
    "            upsample = tf.keras.layers.UpSampling2D(size=2)(layer_input)\n",
    "            upsample = tf.keras.layers.Conv2D(filters, kernel_size=4, strides=1, padding=\"same\", activation=\"relu\")(upsample)\n",
    "            if dropout_rate:\n",
    "                upsample = tf.keras.layers.Dropout(dropout_rate)(upsample)\n",
    "            upsample = tf.keras.layers.BatchNormalization(momentum=BN_MOMENTUM)(upsample)\n",
    "            upsample = tf.keras.layers.Concatenate()([upsample, skip_input])\n",
    "            return upsample\n",
    "\n",
    "        downsample_0 = tf.keras.layers.Input(shape=self.image_shape)\n",
    "        downsample_1 = conv2d(downsample_0, self.generator_filters, bn=False)\n",
    "        downsample_2 = conv2d(downsample_1, self.generator_filters*2)\n",
    "        downsample_3 = conv2d(downsample_2, self.generator_filters*4)\n",
    "        downsample_4 = conv2d(downsample_3, self.generator_filters*8)\n",
    "        downsample_5 = conv2d(downsample_4, self.generator_filters*8)\n",
    "        downsample_6 = conv2d(downsample_5, self.generator_filters*8)\n",
    "        downsample_7 = conv2d(downsample_6, self.generator_filters*8)\n",
    "\n",
    "        upsample_1 = deconv2d(downsample_7, downsample_6, self.generator_filters*8)\n",
    "        upsample_2 = deconv2d(upsample_1, downsample_5, self.generator_filters*8)\n",
    "        upsample_3 = deconv2d(upsample_2, downsample_4, self.generator_filters*8)\n",
    "        upsample_4 = deconv2d(upsample_3, downsample_3, self.generator_filters*4)\n",
    "        upsample_5 = deconv2d(upsample_4, downsample_2, self.generator_filters*2)\n",
    "        upsample_6 = deconv2d(upsample_5, downsample_1, self.generator_filters)\n",
    "        upsample_7 = tf.keras.layers.UpSampling2D(size=2)(upsample_6)\n",
    "        \n",
    "        output_image = tf.keras.layers.Conv2D(IMAGE_CHANNELS, kernel_size=4, strides=1, padding=\"same\", activation=\"tanh\")(upsample_7)\n",
    "        return tf.keras.models.Model(downsample_0, output_image)\n",
    "\n",
    "    def discriminator(self):\n",
    "        def discriminator_layer(layer_input, filters, bn=True):\n",
    "            discriminator_layer = tf.keras.layers.Conv2D(filters, kernel_size=4, strides=2, padding=\"same\")(layer_input)\n",
    "            discriminator_layer = tf.keras.layers.LeakyReLU(alpha=LEAKY_RELU_ALPHA)(discriminator_layer)\n",
    "            if bn:\n",
    "                discriminator_layer = tf.keras.layers.BatchNormalization(momentum=BN_MOMENTUM)(discriminator_layer)\n",
    "            return discriminator_layer\n",
    "\n",
    "        source_image = tf.keras.layers.Input(shape=self.image_shape)\n",
    "        destination_image = tf.keras.layers.Input(shape=self.image_shape)\n",
    "        combined_images = tf.keras.layers.Concatenate(axis=-1)([source_image, destination_image])\n",
    "        discriminator_layer_1 = discriminator_layer(combined_images, self.discriminator_filters, bn=False)\n",
    "        discriminator_layer_2 = discriminator_layer(discriminator_layer_1, self.discriminator_filters*2)\n",
    "        discriminator_layer_3 = discriminator_layer(discriminator_layer_2, self.discriminator_filters*4)\n",
    "        discriminator_layer_4 = discriminator_layer(discriminator_layer_3, self.discriminator_filters*8)\n",
    "        validity = tf.keras.layers.Conv2D(1, kernel_size=4, strides=1, padding=\"same\")(discriminator_layer_4)\n",
    "        return tf.keras.models.Model([source_image, destination_image], validity)\n",
    "        \n",
    "    def preview_training_progress(self, epoch, size=3):\n",
    "        def preview_outputs(epoch, size):\n",
    "            source_images, destination_images = self.data_loader.load_random_data(size, is_testing=True)\n",
    "            generated_images = self.generator.predict(destination_images)\n",
    "            grid_image = None\n",
    "            for i in range(size):\n",
    "                row = Helpers.unnormalize(np.concatenate([destination_images[i], generated_images[i], source_images[i]], axis=1))\n",
    "                if grid_image is None:\n",
    "                    grid_image = row\n",
    "                else:\n",
    "                    grid_image = np.concatenate([grid_image, row], axis=0)\n",
    "            plt.imshow(grid_image/255.0)\n",
    "            plt.show()\n",
    "            plt.close()\n",
    "            grid_image = cv2.cvtColor(np.float32(grid_image), cv2.COLOR_RGB2BGR)\n",
    "            cv2.imwrite(BASE_OUTPUT_PATH + \"training/ \" + str(epoch) + \".png\", grid_image)\n",
    "            \n",
    "        def preview_losses():\n",
    "            def plot(title, data):\n",
    "                plt.plot(data, alpha=0.6)\n",
    "                plt.title(title + \"_\" + str(i))\n",
    "                plt.savefig(BASE_OUTPUT_PATH + \"training/losses/\" + title + \"_\" + str(i) + \".png\")\n",
    "                plt.close()\n",
    "            for i, d in enumerate(self.d_losses):\n",
    "                plot(\"discriminator\", d)\n",
    "            for i, g in enumerate(self.g_losses):\n",
    "                plot(\"generator\", g)\n",
    "                \n",
    "        preview_outputs(epoch, size)\n",
    "        #preview_losses()\n",
    "\n",
    "    def train(self):\n",
    "        valid = np.ones((BATCH_SIZE,) + self.disc_patch)\n",
    "        fake = np.zeros((BATCH_SIZE,) + self.disc_patch)\n",
    "        self.d_losses = []\n",
    "        self.g_losses = []\n",
    "        self.preview_training_progress(0)\n",
    "        for epoch in range(EPOCHS):\n",
    "            epoch_d_losses = []\n",
    "            epoch_g_losses = []\n",
    "            for iteration, (source_images, destination_images) in enumerate(self.data_loader.yield_batch(BATCH_SIZE)):\n",
    "                generated_images = self.generator.predict(destination_images)\n",
    "                d_loss_real = self.discriminator.train_on_batch([source_images, destination_images], valid)\n",
    "                d_loss_fake = self.discriminator.train_on_batch([generated_images, destination_images], fake)\n",
    "                d_losses = 0.5 * np.add(d_loss_real, d_loss_fake)\n",
    "                g_losses = self.combined.train_on_batch([source_images, destination_images], [valid, source_images])\n",
    "                epoch_d_losses.append(d_losses)\n",
    "                epoch_g_losses.append(g_losses)\n",
    "                print(\"\\repoch: \" + str(epoch) \n",
    "                      +\", iteration: \"+ str(iteration) \n",
    "                      + \", d_losses: \" + str(d_losses) \n",
    "                      + \", g_losses: \" + str(g_losses)\n",
    "                      , sep=\" \", end=\" \", flush=True)\n",
    "            self.d_losses.append(np.average(epoch_d_losses, axis=0))\n",
    "            self.g_losses.append(np.average(epoch_g_losses, axis=0))\n",
    "            self.preview_training_progress(epoch)\n",
    "        \n",
    "    def test(self):\n",
    "        image_paths = glob(self.data_loader.testing_raw_path+\"*\")\n",
    "        for image_path in image_paths:\n",
    "            image = np.array(imageio.imread(image_path))\n",
    "            image_normalized = Helpers.normalize(image)\n",
    "            generated_batch = self.generator.predict(np.array([image_normalized]))\n",
    "            concat = Helpers.unnormalize(np.concatenate([image_normalized, generated_batch[0]], axis=1))\n",
    "            cv2.imwrite(BASE_OUTPUT_PATH+os.path.basename(image_path), cv2.cvtColor(np.float32(concat), cv2.COLOR_RGB2BGR))\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_uuid": "abbde5b9129e01502039e3a4638b678385bfbf77"
   },
   "outputs": [],
   "source": [
    "BASE_INPUT_PATH = \"\" # Kaggle: \"../input/pix2pix-depth/\" \n",
    "BASE_OUTPUT_PATH = \"./output/\"\n",
    "\n",
    "IMAGE_SIZE = 256\n",
    "IMAGE_CHANNELS = 3\n",
    "LEARNING_RATE = 0.00015\n",
    "BETA_1 = 0.5\n",
    "LEAKY_RELU_ALPHA = 0.2\n",
    "BN_MOMENTUM = 0.8\n",
    "EPOCHS = 50\n",
    "BATCH_SIZE = 32\n",
    "\n",
    "gan = Pix2Pix()\n",
    "gan.train()\n",
    "gan.test()"
   ]
  }
 ],
 "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}