train_google_colab.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "horizon-unet.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T-2blStPzGfx"
      },
      "source": [
        "#Horizon UAV\n",
        "Using U-Net"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FeR-3e4LzPFe"
      },
      "source": [
        "Load Dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4geyEGySzSKx"
      },
      "source": [
        "# Get dataset\n",
        "!wget -q -O dataset_horizon.zip https://www.dropbox.com/s/3mc9u1dtrogq9xc/dataset_horizon.zip?dl=0\n",
        "!unzip -q -o dataset_horizon.zip\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qYaeZsV5zviP"
      },
      "source": [
        "Library Import"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7BpiyiS9dIsP"
      },
      "source": [
        "import tensorflow as tf\n",
        "import matplotlib.pyplot as plt\n",
        "import os\n",
        "import datetime\n",
        "import random\n",
        "import numpy as np\n",
        "import io\n",
        "import math\n",
        "import cv2\n",
        "import pandas as pd\n",
        "import albumentations as albu\n",
        "from sklearn.model_selection import train_test_split\n",
        "from pathlib import Path\n",
        "from time import time as timer\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N9fP0te8dK8N"
      },
      "source": [
        "Hyperparameter"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lF9OoXtTdN5Q"
      },
      "source": [
        "# Training parameter\n",
        "test_size = 0.2\n",
        "random_seed = 192\n",
        "\n",
        "# Hyperparameter\n",
        "epoch = 50\n",
        "batch_size = 32\n",
        "learning_rate = 0.001\n",
        "n_encoder_decoder = 3\n",
        "initial_filter = 8\n",
        "image_size = (128, 128)\n",
        "\n",
        "# Augmentation\n",
        "transformations = [albu.HorizontalFlip(p=0.5),\n",
        "                   albu.VerticalFlip(p=0.5),\n",
        "                   albu.ShiftScaleRotate(p=0.5, border_mode=1),\n",
        "                   albu.RandomBrightnessContrast(p=0.25),\n",
        "                   albu.RandomGamma(p=0.25)\n",
        "                   ]\n",
        "\n",
        "aug = albu.Compose(transformations)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-O22eAg1dRqj"
      },
      "source": [
        "Data Generator"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7hqZsokudTbS"
      },
      "source": [
        "# Data Generator\n",
        "class data_generator(tf.keras.utils.Sequence):\n",
        "    def __init__(self, file_list, batch_size, image_size, \\\n",
        "        shuffle=True, augmentation=None):\n",
        "\n",
        "        self.file_list = file_list\n",
        "        self.batch_size = batch_size\n",
        "        self.image_size = image_size\n",
        "        self.shuffle = shuffle\n",
        "        self.aug = augmentation\n",
        "        self.on_epoch_end()\n",
        "    \n",
        "    def __len__(self):\n",
        "        return math.ceil(len(self.file_list) / self.batch_size)\n",
        "\n",
        "    def on_epoch_end(self):\n",
        "        self.indexes = np.arange(len(self.file_list))\n",
        "        if self.shuffle == True:\n",
        "            np.random.shuffle(self.indexes)\n",
        "    \n",
        "    def __getitem__(self, index):\n",
        "        indexes = self.indexes[index*self.batch_size:(index+1)*\\\n",
        "            self.batch_size]\n",
        "\n",
        "        batch = [self.file_list[k] for k in indexes]\n",
        "\n",
        "        # Create batch list\n",
        "        batch_x = []\n",
        "        batch_y = []\n",
        "\n",
        "        for filename in batch:\n",
        "            # Load Image\n",
        "            image = cv2.imread(os.path.join(\"dataset\", \"images\", \\\n",
        "                filename + \".jpg\"))\n",
        "\n",
        "            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
        "\n",
        "            # Load mask\n",
        "            mask_land = cv2.imread(os.path.join(\"dataset\", \\\n",
        "                \"masks\", \"land\", filename + \".png\"), \\\n",
        "                cv2.IMREAD_GRAYSCALE)\n",
        "\n",
        "            mask_sky = cv2.imread(os.path.join(\"dataset\", \\\n",
        "                \"masks\", \"sky\", filename + \".png\"), \\\n",
        "                cv2.IMREAD_GRAYSCALE)\n",
        "\n",
        "            mask = np.dstack((mask_land, mask_sky))\n",
        "\n",
        "            # Resize\n",
        "            image = cv2.resize(image, self.image_size)\n",
        "            mask = cv2.resize(mask, self.image_size, \\\n",
        "                interpolation = cv2.INTER_NEAREST)\n",
        "            \n",
        "            # Augmentation\n",
        "            if self.aug is not None:\n",
        "                augmented = self.aug(image=image, mask=mask)\n",
        "                image = augmented[\"image\"]\n",
        "                mask = augmented[\"mask\"]\n",
        "            \n",
        "            # Normalize\n",
        "            image = cv2.normalize(image, None, 0, 1, \\\n",
        "                cv2.NORM_MINMAX, cv2.CV_32F)\n",
        "            \n",
        "            mask = cv2.normalize(mask, None, 0, 1, \\\n",
        "                cv2.NORM_MINMAX, cv2.CV_32F)\n",
        "            \n",
        "            # Load to batch\n",
        "            batch_x.append(image)\n",
        "            batch_y.append(mask)\n",
        "\n",
        "        # Convert batch as array\n",
        "        batch_x = np.array(batch_x)\n",
        "        batch_y = np.array(batch_y)\n",
        "\n",
        "        return batch_x, batch_y\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NFuAJEAGdVBP"
      },
      "source": [
        "Custom Loss Function"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fqMltujhdXpq"
      },
      "source": [
        "# Loss Function\n",
        "def dice_loss(y_true, y_pred, num_classes=2):\n",
        "    smooth = tf.keras.backend.epsilon()\n",
        "    dice = 0\n",
        "    for index in range(num_classes):\n",
        "        y_true_f = tf.keras.backend.flatten(y_true[:,:,:,index])\n",
        "        y_pred_f = tf.keras.backend.flatten(y_pred[:,:,:,index])\n",
        "        intersection = tf.keras.backend.sum(y_true_f * y_pred_f)\n",
        "        union = tf.keras.backend.sum(y_true_f) + \\\n",
        "            tf.keras.backend.sum(y_pred_f)\n",
        "        dice -= (2. * intersection + smooth) / (union + smooth)\n",
        "    return dice/num_classes\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QARKXnqQ9BGY"
      },
      "source": [
        "Custom MeanIoU"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DO1w62M39Ecn"
      },
      "source": [
        "# Metric Function\n",
        "class MaxMeanIoU(tf.keras.metrics.MeanIoU):\n",
        "    def update_state(self, y_true, y_pred, sample_weight=None):\n",
        "        return super().update_state(tf.argmax(y_true, axis=-1), tf.argmax(y_pred, axis=-1), sample_weight)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xx1g-SlRdZnQ"
      },
      "source": [
        "Create Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5jM_aC7rdakE"
      },
      "source": [
        "# Create model\n",
        "def create_model():\n",
        "    # Input\n",
        "    input_shape = (image_size[0], image_size[1], 3)\n",
        "    inputs = tf.keras.layers.Input(shape=input_shape)\n",
        "    x = inputs\n",
        "\n",
        "    encoder_layers = []\n",
        "    # Encoder\n",
        "    for i in range(n_encoder_decoder):\n",
        "        filter_number = int(2**(math.log2(initial_filter)+i))\n",
        "        x = tf.keras.layers.Conv2D(filter_number, 3, \\\n",
        "            activation='relu', padding='same')(x)\n",
        "\n",
        "        x = tf.keras.layers.Conv2D(filter_number, 3, \\\n",
        "            activation='relu', padding='same')(x)\n",
        "\n",
        "        encoder_layers.append(x)\n",
        "        x = tf.keras.layers.MaxPool2D()(x)\n",
        "\n",
        "    # Bridge\n",
        "    filter_number = int(2**(math.log2(initial_filter)+\\\n",
        "        n_encoder_decoder))\n",
        "\n",
        "    x = tf.keras.layers.Conv2D(filter_number, 3, activation='relu', \\\n",
        "        padding='same')(x)\n",
        "\n",
        "    x = tf.keras.layers.Conv2D(filter_number, 3, activation='relu', \\\n",
        "        padding='same')(x)\n",
        "\n",
        "    # Decoder\n",
        "    for i in reversed(range(n_encoder_decoder)):\n",
        "        filter_number = int(2**(math.log2(initial_filter)+i))\n",
        "        x = tf.keras.layers.Conv2DTranspose(filter_number, 2, \\\n",
        "            strides=(2, 2), activation='relu', padding='same')(x)\n",
        "\n",
        "        x = tf.keras.layers.Concatenate(axis=3)([x, encoder_layers[i]])\n",
        "        x = tf.keras.layers.Conv2D(filter_number, 3, \\\n",
        "            activation='relu', padding='same')(x)\n",
        "\n",
        "        x = tf.keras.layers.Conv2D(filter_number, 3, \\\n",
        "            activation='relu', padding='same')(x)\n",
        "    \n",
        "    # Output\n",
        "    outputs = tf.keras.layers.Conv2D(2, 1)(x)\n",
        "    outputs = tf.keras.layers.Activation('softmax')(outputs)\n",
        "\n",
        "    # Create Optimizer\n",
        "    opt = tf.keras.optimizers.Adam(learning_rate=learning_rate)\n",
        "\n",
        "    # Create Loss Function\n",
        "    loss = dice_loss\n",
        "\n",
        "    # Create Model\n",
        "    model = tf.keras.models.Model(inputs=[inputs], outputs=[outputs])\n",
        "    model.compile(optimizer = opt, loss = loss, metrics=[\"accuracy\", MaxMeanIoU(num_classes=2)])\n",
        "    \n",
        "    return model\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PjmYn3HNdbgT"
      },
      "source": [
        "Create Custom Training Callback"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zIQ71R9odf-_"
      },
      "source": [
        "# Create Callback\n",
        "def create_callback():\n",
        "    # Tensorboard Callbacks\n",
        "    logdir = os.path.join(\"logs\", datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\"))\n",
        "    tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=logdir)\n",
        "\n",
        "    # Checkpoint Callbacks\n",
        "    Path(\"checkpoint\").mkdir(parents=True, exist_ok=True)\n",
        "    \n",
        "    best_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(filepath=os.path.join(\"checkpoint\", \"best\" + datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\") + \".h5\"), \n",
        "                                                             monitor='max_mean_io_u', verbose=1, save_best_only=True, mode='max')\n",
        "\n",
        "    # Predict Image Callbacks\n",
        "    file_writer_cm = tf.summary.create_file_writer(os.path.join(logdir, \"predict_output\"))\n",
        "    def predict_epoch(epoch, logs):\n",
        "        # Load image\n",
        "        filename = np.random.choice(test_list)\n",
        "        image = cv2.imread(os.path.join(\"dataset\", \"images\", filename + \".jpg\"))\n",
        "        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
        "        image = cv2.resize(image, image_size)\n",
        "        image = cv2.normalize(image, None, 0, 1, cv2.NORM_MINMAX, cv2.CV_32F)\n",
        "        \n",
        "        # Predict mask\n",
        "        pred = model.predict(np.expand_dims(image, 0))\n",
        "\n",
        "        # Process mask\n",
        "        mask = pred.squeeze()\n",
        "        mask = np.stack((mask,)*3, axis=-1)\n",
        "        mask[mask >= 0.5] = 1\n",
        "        mask[mask < 0.5] = 0\n",
        "\n",
        "        class_land = np.concatenate([image, mask[:, :, 0], image * mask[:, :, 0]], axis = 1)\n",
        "        class_sky = np.concatenate([image, mask[:, :, 1], image * mask[:, :, 1]], axis = 1)\n",
        "        \n",
        "        # Log the image as an image summary.\n",
        "        with file_writer_cm.as_default():\n",
        "            tf.summary.image(\"class_land\", np.reshape(class_land, (1, image_size[0], image_size[1]*3, 3)), step=epoch)\n",
        "            tf.summary.image(\"class_sky\", np.reshape(class_sky, (1, image_size[0], image_size[1]*3, 3)), step=epoch)\n",
        "\n",
        "    # Define per-epoch callback.\n",
        "    predict_callback = tf.keras.callbacks.LambdaCallback(on_epoch_end=predict_epoch)\n",
        "\n",
        "    return [tensorboard_callback, best_checkpoint_callback, predict_callback]\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zfyM_-hmdhZi"
      },
      "source": [
        "Training"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "m8iYkwPOditf"
      },
      "source": [
        "# Training\n",
        "\n",
        "# Load Data\n",
        "dataset_path = \"dataset/images\"\n",
        "file_list = [os.path.splitext(filename)[0] for filename in os.listdir(dataset_path)]\n",
        "\n",
        "# Data Split\n",
        "train_list, test_list = train_test_split(file_list, shuffle=True, \\\n",
        "    test_size=test_size, random_state=random_seed)\n",
        "\n",
        "start = timer()\n",
        "\n",
        "loss = []\n",
        "accuracy = []\n",
        "mean_io_u = []\n",
        "\n",
        "# Load data\n",
        "train_generator = data_generator(train_list, \\\n",
        "    batch_size=batch_size, image_size=image_size, \\\n",
        "    augmentation=aug)\n",
        "\n",
        "val_generator = data_generator(test_list, \\\n",
        "    batch_size=batch_size, image_size=image_size)\n",
        "\n",
        "# Create model\n",
        "model = create_model()\n",
        "\n",
        "# Train model\n",
        "model.fit(train_generator, epochs=epoch, \\\n",
        "    validation_data=val_generator,\\\n",
        "    callbacks=create_callback(), max_queue_size=5)\n",
        "\n",
        "# Evaluate Model\n",
        "result = model.evaluate(val_generator)\n",
        "\n",
        "loss = result[0]\n",
        "accuracy = result[1]\n",
        "mean_io_u = result[2]\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SeOTpJbsdlUM"
      },
      "source": [
        "Print Result"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xNeHkVSOdnEl"
      },
      "source": [
        "# Get evaluation metric\n",
        "print(\"Run Date:\", datetime.datetime.now().strftime(\"%Y%m%d-%H%M%S\"))\n",
        "print(\"Elapsed Time:\", timer() - start, \"Seconds\")\n",
        "print(\"Training parameter\")\n",
        "print(\"test_size:\", test_size)\n",
        "print(\"random_seed:\", random_seed)\n",
        "print()\n",
        "print(\"Hyperparameter:\")\n",
        "print(\"epoch:\", epoch)\n",
        "print(\"batch_size:\", batch_size)\n",
        "print(\"learning_rate:\", learning_rate)\n",
        "print(\"n_encoder_decoder:\", n_encoder_decoder)\n",
        "print(\"initial_filter:\", initial_filter)\n",
        "print()\n",
        "print(\"Result:\")\n",
        "print(\"loss:\", loss)\n",
        "print(\"accuracy:\", accuracy)\n",
        "print(\"mean_io_u:\", mean_io_u)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "522qE15UvNUl"
      },
      "source": [
        "Save Result to Google Drive"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Xm1l4-GDvTbd"
      },
      "source": [
        "# Mount Google Drive\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "We1FlLEjvOmd"
      },
      "source": [
        "# Zip and save result to google drive\n",
        "!zip -r -q template_new.zip checkpoint logs\n",
        "!cp /content/template_new.zip /content/drive/My\\ Drive/horizon-uav/\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ha4Hahi7ioSP"
      },
      "source": [
        "Tensorboard"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ng-Mct_Qipa4"
      },
      "source": [
        "# Load the TensorBoard notebook extension\n",
        "%load_ext tensorboard\n",
        "%tensorboard --logdir logs\n"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}