C3_W3_Lab_2_OxfordPets_UNet.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "C3_W3_Lab_2_OxfordPets-UNet.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "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"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/vsnupoudel/AdvancedCvCourses/blob/main/C3_W3_Lab_2_OxfordPets_UNet.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VxdtCvC2mpR6"
      },
      "source": [
        "# Ungraded Lab: U-Net for Image Segmentation\n",
        "\n",
        "This notebook illustrates how to build a [UNet](https://arxiv.org/abs/1505.04597) for semantic image segmentation. This architecture is also a fully convolutional network and is similar to the model you just built in the previous lesson. A key difference is the use of skip connections from the encoder to the decoder. You will see how this is implemented later as you build each part of the network.\n",
        "\n",
        "At the end of this lab, you will be able to use the UNet to output segmentation masks that shows which pixels of an input image are part of the background, foreground, and outline. \n",
        "\n",
        "<img src='https://drive.google.com/uc?export=view&id=1w8up90xVOYRT8vs5lNo2kEiVhm-7hWxZ' alt='sample_output'>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NbSvE6h4mZyO"
      },
      "source": [
        "# Imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YQX7R4bhZy5h"
      },
      "source": [
        "try:\n",
        "  # %tensorflow_version only exists in Colab.\n",
        "  %tensorflow_version 2.x\n",
        "except Exception:\n",
        "  pass\n",
        "\n",
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oWe0_rQM4JbC"
      },
      "source": [
        "## Download the Oxford-IIIT Pets dataset\n",
        "\n",
        "You will be training the model on the [Oxford Pets - IIT dataset](https://www.robots.ox.ac.uk/~vgg/data/pets/) dataset. This contains pet images, their classes, segmentation masks and head region-of-interest. You will only use the images and segmentation masks in this lab.\n",
        "\n",
        "This dataset is already included in TensorFlow Datasets and you can simply download it. The segmentation masks are included in versions 3 and above. The cell below will download the dataset and place the results in a dictionary named `dataset`. It will also collect information about the dataset and we'll assign it to a variable named `info`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YDEza-tcR4SQ"
      },
      "source": [
        "# If you hit a problem with checksums, you can execute the following line first\n",
        "!python -m tensorflow_datasets.scripts.download_and_prepare --register_checksums --datasets=oxford_iiit_pet:3.2.0\n",
        "\n",
        "# download the dataset and get info\n",
        "dataset, info = tfds.load('oxford_iiit_pet:3.*.*', with_info=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CKIvihxy4q0K"
      },
      "source": [
        "Let's briefly examine the contents of the dataset you just downloaded."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RLqGeTL39uLz"
      },
      "source": [
        "# see the possible keys we can access in the dataset dict.\n",
        "# this contains the test and train splits.\n",
        "print(dataset.keys())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hpedAXHz-Qwl"
      },
      "source": [
        "# see information about the dataset\n",
        "print(info)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GYMHMvk5na7_"
      },
      "source": [
        "## Prepare the Dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rJcVdj_U4vzf"
      },
      "source": [
        "You will now prepare the train and test sets. The following utility functions preprocess the data. These include:\n",
        "\n",
        "* simple augmentation by flipping the image\n",
        "* normalizing the pixel values  \n",
        "* resizing the images\n",
        "\n",
        "Another preprocessing step is to adjust the segmentation mask's pixel values. The `README` in the [annotations](https://www.robots.ox.ac.uk/~vgg/data/pets/data/annotations.tar.gz) folder of the dataset mentions that the pixels in the segmentation mask are labeled as such:\n",
        "\n",
        "| Label            | Class Name     |\n",
        "| -------------    | -------------  |\n",
        "| 1                | foreground     |\n",
        "| 2                | background     |\n",
        "| 3                | Not Classified |\n",
        "\n",
        "<br>\n",
        "<br>\n",
        "\n",
        "For convenience, let's subtract `1` from these values and we will interpret these as `{'pet', 'background', 'outline'}`:\n",
        "\n",
        "| Label            | Class Name     |\n",
        "| -------------    | -------------  |\n",
        "| 0                | pet            |\n",
        "| 1                | background     |\n",
        "| 2                | outline        |"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FD60EbcAQqov"
      },
      "source": [
        "# Preprocessing Utilities\n",
        "\n",
        "def random_flip(input_image, input_mask):\n",
        "  '''does a random flip of the image and mask'''\n",
        "  if tf.random.uniform(()) > 0.5:\n",
        "    input_image = tf.image.flip_left_right(input_image)\n",
        "    input_mask = tf.image.flip_left_right(input_mask)\n",
        "\n",
        "  return input_image, input_mask\n",
        "\n",
        "\n",
        "def normalize(input_image, input_mask):\n",
        "  '''\n",
        "  normalizes the input image pixel values to be from [0,1].\n",
        "  subtracts 1 from the mask labels to have a range from [0,2]\n",
        "  '''\n",
        "  input_image = tf.cast(input_image, tf.float32) / 255.0\n",
        "  input_mask -= 1\n",
        "  return input_image, input_mask\n",
        "\n",
        "\n",
        "@tf.function\n",
        "def load_image_train(datapoint):\n",
        "  '''resizes, normalizes, and flips the training data'''\n",
        "  input_image = tf.image.resize(datapoint['image'], (128, 128), method='nearest')\n",
        "  input_mask = tf.image.resize(datapoint['segmentation_mask'], (128, 128), method='nearest')\n",
        "  input_image, input_mask = random_flip(input_image, input_mask)\n",
        "  input_image, input_mask = normalize(input_image, input_mask)\n",
        "  \n",
        "  return input_image, input_mask\n",
        "\n",
        "\n",
        "def load_image_test(datapoint):\n",
        "  '''resizes and normalizes the test data'''\n",
        "  input_image = tf.image.resize(datapoint['image'], (128, 128), method='nearest')\n",
        "  input_mask = tf.image.resize(datapoint['segmentation_mask'], (128, 128), method='nearest')\n",
        "  input_image, input_mask = normalize(input_image, input_mask)\n",
        "\n",
        "  return input_image, input_mask"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "65-qHTjX5VZh"
      },
      "source": [
        "You can now call the utility functions above to prepare the train and test sets. The dataset you downloaded from TFDS already contains these splits and you will use those by simpling accessing the `train` and `test` keys of the `dataset` dictionary.\n",
        "\n",
        "*Note*: The `tf.data.experimental.AUTOTUNE` you see in this notebook is simply a constant equal to `-1`. This value is passed to allow certain methods to automatically set parameters based on available resources. For instance, `num_parallel_calls` parameter below will be set dynamically based on the available CPUs. The docstrings will show if a parameter can be autotuned. [Here](https://github.com/tensorflow/tensorflow/blob/v2.3.1/tensorflow/python/data/ops/dataset_ops.py#L1557-L1702) is the entry describing what it does to `num_parallel_calls`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "39fYScNz9lmo"
      },
      "source": [
        "# preprocess the train and test sets\n",
        "train = dataset['train'].map(load_image_train, num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "test = dataset['test'].map(load_image_test)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dWL5sxz7_L-J"
      },
      "source": [
        "Now that the splits are loaded, you can then prepare batches for training and testing."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DeFwFDN6EVoI"
      },
      "source": [
        "BATCH_SIZE = 64\n",
        "BUFFER_SIZE = 1000\n",
        "\n",
        "# shuffle and group the train set into batches\n",
        "train_dataset = train.cache().shuffle(BUFFER_SIZE).batch(BATCH_SIZE).repeat()\n",
        "\n",
        "# do a prefetch to optimize processing\n",
        "train_dataset = train_dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)\n",
        "\n",
        "# group the test set into batches\n",
        "test_dataset = test.batch(BATCH_SIZE)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ck3QY8r0AUfK"
      },
      "source": [
        "Let's define a few more utilities to help us visualize our data and metrics."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "n34OGwJXzFEu"
      },
      "source": [
        "# class list of the mask pixels\n",
        "class_names = ['pet', 'background', 'outline']\n",
        "\n",
        "\n",
        "def display_with_metrics(display_list, iou_list, dice_score_list):\n",
        "  '''displays a list of images/masks and overlays a list of IOU and Dice Scores'''\n",
        "  \n",
        "  metrics_by_id = [(idx, iou, dice_score) for idx, (iou, dice_score) in enumerate(zip(iou_list, dice_score_list)) if iou > 0.0]\n",
        "  metrics_by_id.sort(key=lambda tup: tup[1], reverse=True)  # sorts in place\n",
        "  \n",
        "  display_string_list = [\"{}: IOU: {} Dice Score: {}\".format(class_names[idx], iou, dice_score) for idx, iou, dice_score in metrics_by_id]\n",
        "  display_string = \"\\n\\n\".join(display_string_list)\n",
        "\n",
        "  display(display_list, [\"Image\", \"Predicted Mask\", \"True Mask\"], display_string=display_string) \n",
        "\n",
        "\n",
        "def display(display_list,titles=[], display_string=None):\n",
        "  '''displays a list of images/masks'''\n",
        "\n",
        "  plt.figure(figsize=(15, 15))\n",
        "\n",
        "  for i in range(len(display_list)):\n",
        "    plt.subplot(1, len(display_list), i+1)\n",
        "    plt.title(titles[i])\n",
        "    plt.xticks([])\n",
        "    plt.yticks([])\n",
        "    if display_string and i == 1:\n",
        "      plt.xlabel(display_string, fontsize=12)\n",
        "    img_arr = tf.keras.preprocessing.image.array_to_img(display_list[i])\n",
        "    plt.imshow(img_arr)\n",
        "  \n",
        "  plt.show()\n",
        "\n",
        "\n",
        "def show_image_from_dataset(dataset):\n",
        "  '''displays the first image and its mask from a dataset'''\n",
        "\n",
        "  for image, mask in dataset.take(1):\n",
        "    sample_image, sample_mask = image, mask\n",
        "  display([sample_image, sample_mask], titles=[\"Image\", \"True Mask\"])\n",
        "\n",
        "\n",
        "def plot_metrics(metric_name, title, ylim=5):\n",
        "  '''plots a given metric from the model history'''\n",
        "  plt.title(title)\n",
        "  plt.ylim(0,ylim)\n",
        "  plt.plot(model_history.history[metric_name],color='blue',label=metric_name)\n",
        "  plt.plot(model_history.history['val_' + metric_name],color='green',label='val_' + metric_name)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xa3gMAE_9qNa"
      },
      "source": [
        "Finally, you can take a look at an image example and it's correponding mask from the dataset."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "a6u_Rblkteqb"
      },
      "source": [
        "# display an image from the train set\n",
        "show_image_from_dataset(train)\n",
        "\n",
        "# display an image from the test set\n",
        "show_image_from_dataset(test)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FAOe93FRMk3w"
      },
      "source": [
        "## Define the model\n",
        "\n",
        "With the dataset prepared, you can now build the UNet. Here is the overall architecture as shown in class:\n",
        "\n",
        "<img src='https://drive.google.com/uc?export=view&id=1BeQSKL2Eq6Fw9iRXsN1hgunY-CS2nH7V' alt='unet'>\n",
        "\n",
        "A UNet consists of an encoder (downsampler) and decoder (upsampler) with a bottleneck in between. The gray arrows correspond to the skip connections that concatenate encoder block outputs to each stage of the decoder. Let's see how to implement these starting with the encoder."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_HtSYSzxFMBh"
      },
      "source": [
        "### Encoder\n",
        "\n",
        "Like the FCN model you built in the previous lesson, the encoder here will have repeating blocks (red boxes in the figure below) so it's best to create functions for it to make the code modular. These encoder blocks will contain two Conv2D layers activated by ReLU, followed by a MaxPooling and Dropout layer. As discussed in class, each stage will have increasing number of filters and the dimensionality of the features will reduce because of the pooling layer."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ul8_qU1scs5C"
      },
      "source": [
        "<img src='https://drive.google.com/uc?export=view&id=1Gs9K3_8ZBn2_ntOtJL_-_ww4ZOgfyhrS' alt='unet'>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SJqbVK3FDUjK"
      },
      "source": [
        "The encoder utilities will have three functions:\n",
        "\n",
        "* `conv2d_block()` - to add two convolution layers and ReLU activations\n",
        "* `encoder_block()` - to add pooling and dropout to the conv2d blocks. Recall that in UNet, you need to save the output of the convolution layers at each block so this function will return two values to take that into account (i.e. output of the conv block and the dropout)\n",
        "* `encoder()` - to build the entire encoder. This will return the output of the last encoder block as well as the output of the previous conv blocks. These will be concatenated to the decoder blocks as you'll see later."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VoGZBIzs8Ln-"
      },
      "source": [
        "# Encoder Utilities\n",
        "\n",
        "def conv2d_block(input_tensor, n_filters, kernel_size = 3):\n",
        "  '''\n",
        "  Adds 2 convolutional layers with the parameters passed to it\n",
        "\n",
        "  Args:\n",
        "    input_tensor (tensor) -- the input tensor\n",
        "    n_filters (int) -- number of filters\n",
        "    kernel_size (int) -- kernel size for the convolution\n",
        "\n",
        "  Returns:\n",
        "    tensor of output features\n",
        "  '''\n",
        "  # first layer\n",
        "  x = input_tensor\n",
        "  for i in range(2):\n",
        "    x = tf.keras.layers.Conv2D(filters = n_filters, kernel_size = (kernel_size, kernel_size),\\\n",
        "            kernel_initializer = 'he_normal', padding = 'same')(x)\n",
        "    x = tf.keras.layers.Activation('relu')(x)\n",
        "  \n",
        "  return x\n",
        "\n",
        "\n",
        "def encoder_block(inputs, n_filters=64, pool_size=(2,2), dropout=0.3):\n",
        "  '''\n",
        "  Adds two convolutional blocks and then perform down sampling on output of convolutions.\n",
        "\n",
        "  Args:\n",
        "    input_tensor (tensor) -- the input tensor\n",
        "    n_filters (int) -- number of filters\n",
        "    kernel_size (int) -- kernel size for the convolution\n",
        "\n",
        "  Returns:\n",
        "    f - the output features of the convolution block \n",
        "    p - the maxpooled features with dropout\n",
        "  '''\n",
        "\n",
        "  f = conv2d_block(inputs, n_filters=n_filters)\n",
        "  p = tf.keras.layers.MaxPooling2D(pool_size=(2,2))(f)\n",
        "  p = tf.keras.layers.Dropout(0.3)(p)\n",
        "\n",
        "  return f, p\n",
        "\n",
        "\n",
        "def encoder(inputs):\n",
        "  '''\n",
        "  This function defines the encoder or downsampling path.\n",
        "\n",
        "  Args:\n",
        "    inputs (tensor) -- batch of input images\n",
        "\n",
        "  Returns:\n",
        "    p4 - the output maxpooled features of the last encoder block\n",
        "    (f1, f2, f3, f4) - the output features of all the encoder blocks\n",
        "  '''\n",
        "  f1, p1 = encoder_block(inputs, n_filters=64, pool_size=(2,2), dropout=0.3)\n",
        "  f2, p2 = encoder_block(p1, n_filters=128, pool_size=(2,2), dropout=0.3)\n",
        "  f3, p3 = encoder_block(p2, n_filters=256, pool_size=(2,2), dropout=0.3)\n",
        "  f4, p4 = encoder_block(p3, n_filters=512, pool_size=(2,2), dropout=0.3)\n",
        "\n",
        "  return p4, (f1, f2, f3, f4)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i6lSYsoOc6j6"
      },
      "source": [
        "### Bottleneck\n",
        "\n",
        "\n",
        "A bottleneck follows the encoder block and is used to extract more features. This does not have a pooling layer so the dimensionality remains the same. You can use the `conv2d_block()` function defined earlier to implement this."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YLzUf31Cuh-f"
      },
      "source": [
        "def bottleneck(inputs):\n",
        "  '''\n",
        "  This function defines the bottleneck convolutions to extract more features before the upsampling layers.\n",
        "  '''\n",
        "  \n",
        "  bottle_neck = conv2d_block(inputs, n_filters=1024)\n",
        "\n",
        "  return bottle_neck"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-__-6WcUa1s2"
      },
      "source": [
        "### Decoder\n",
        "\n",
        "Finally, we have the decoder which upsamples the features back to the original image size. At each upsampling level, you will take the output of the corresponding encoder block and concatenate it before feeding to the next decoder block. This is summarized in the figure below.\n",
        "\n",
        "<img src='https://drive.google.com/uc?export=view&id=1Ql5vdw6l88vxaHgk7VjcMc4vfyoWYx2w' alt='unet_decoder'>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XACX8TJh1oKd"
      },
      "source": [
        "# Decoder Utilities\n",
        "\n",
        "def decoder_block(inputs, conv_output, n_filters=64, kernel_size=3, strides=3, dropout=0.3):\n",
        "  '''\n",
        "  defines the one decoder block of the UNet\n",
        "\n",
        "  Args:\n",
        "    inputs (tensor) -- batch of input features\n",
        "    conv_output (tensor) -- features from an encoder block\n",
        "    n_filters (int) -- number of filters\n",
        "    kernel_size (int) -- kernel size\n",
        "    strides (int) -- strides for the deconvolution/upsampling\n",
        "    padding (string) -- \"same\" or \"valid\", tells if shape will be preserved by zero padding\n",
        "\n",
        "  Returns:\n",
        "    c (tensor) -- output features of the decoder block\n",
        "  '''\n",
        "  u = tf.keras.layers.Conv2DTranspose(n_filters, kernel_size, strides = strides, padding = 'same')(inputs)\n",
        "  c = tf.keras.layers.concatenate([u, conv_output])\n",
        "  c = tf.keras.layers.Dropout(dropout)(c)\n",
        "  c = conv2d_block(c, n_filters, kernel_size=3)\n",
        "\n",
        "  return c\n",
        "\n",
        "\n",
        "def decoder(inputs, convs, output_channels):\n",
        "  '''\n",
        "  Defines the decoder of the UNet chaining together 4 decoder blocks. \n",
        "  \n",
        "  Args:\n",
        "    inputs (tensor) -- batch of input features\n",
        "    convs (tuple) -- features from the encoder blocks\n",
        "    output_channels (int) -- number of classes in the label map\n",
        "\n",
        "  Returns:\n",
        "    outputs (tensor) -- the pixel wise label map of the image\n",
        "  '''\n",
        "  \n",
        "  f1, f2, f3, f4 = convs\n",
        "\n",
        "  c6 = decoder_block(inputs, f4, n_filters=512, kernel_size=(3,3), strides=(2,2), dropout=0.3)\n",
        "  c7 = decoder_block(c6, f3, n_filters=256, kernel_size=(3,3), strides=(2,2), dropout=0.3)\n",
        "  c8 = decoder_block(c7, f2, n_filters=128, kernel_size=(3,3), strides=(2,2), dropout=0.3)\n",
        "  c9 = decoder_block(c8, f1, n_filters=64, kernel_size=(3,3), strides=(2,2), dropout=0.3)\n",
        "\n",
        "  outputs = tf.keras.layers.Conv2D(output_channels, (1, 1), activation='softmax')(c9)\n",
        "\n",
        "  return outputs"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NAtWsYwGExtB"
      },
      "source": [
        "### Putting it all together\n",
        "\n",
        "You can finally build the UNet by chaining the encoder, bottleneck, and decoder. You will specify the number of output channels and in this particular set, that would be `3`. That is because there are three possible labels for each pixel: 'pet', 'background', and 'outline'."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-gE1jiz5u6Zg"
      },
      "source": [
        "OUTPUT_CHANNELS = 3\n",
        "\n",
        "def unet():\n",
        "  '''\n",
        "  Defines the UNet by connecting the encoder, bottleneck and decoder.\n",
        "  '''\n",
        "\n",
        "  # specify the input shape\n",
        "  inputs = tf.keras.layers.Input(shape=(128, 128,3,))\n",
        "\n",
        "  # feed the inputs to the encoder\n",
        "  encoder_output, convs = encoder(inputs)\n",
        "\n",
        "  # feed the encoder output to the bottleneck\n",
        "  bottle_neck = bottleneck(encoder_output)\n",
        "\n",
        "  # feed the bottleneck and encoder block outputs to the decoder\n",
        "  # specify the number of classes via the `output_channels` argument\n",
        "  outputs = decoder(bottle_neck, convs, output_channels=OUTPUT_CHANNELS)\n",
        "  \n",
        "  # create the model\n",
        "  model = tf.keras.Model(inputs=inputs, outputs=outputs)\n",
        "\n",
        "  return model\n",
        "\n",
        "# instantiate the model\n",
        "model = unet()\n",
        "\n",
        "# see the resulting model architecture\n",
        "model.summary()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j0DGH_4T0VYn"
      },
      "source": [
        "## Compile and Train the model\n",
        "\n",
        "Now, all that is left to do is to compile and train the model. The loss you will use is `sparse_categorical_crossentropy`. The reason is because the network is trying to assign each pixel a label, just like multi-class prediction. In the true segmentation mask, each pixel has either a {0,1,2}. The network here is outputting three channels. Essentially, each channel is trying to learn to predict a class and `sparse_categorical_crossentropy` is the recommended loss for such a scenario. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WEyXtFjCzZv5"
      },
      "source": [
        "# configure the optimizer, loss and metrics for training\n",
        "model.compile(optimizer=tf.keras.optimizers.Adam(), loss='sparse_categorical_crossentropy',\n",
        "              metrics=['accuracy'])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "StKDH_B9t4SD"
      },
      "source": [
        "# configure the training parameters and train the model\n",
        "\n",
        "TRAIN_LENGTH = info.splits['train'].num_examples\n",
        "EPOCHS = 10\n",
        "VAL_SUBSPLITS = 5\n",
        "STEPS_PER_EPOCH = TRAIN_LENGTH // BATCH_SIZE\n",
        "VALIDATION_STEPS = info.splits['test'].num_examples//BATCH_SIZE//VAL_SUBSPLITS\n",
        "\n",
        "# this will take around 20 minutes to run\n",
        "model_history = model.fit(train_dataset, epochs=EPOCHS,\n",
        "                          steps_per_epoch=STEPS_PER_EPOCH,\n",
        "                          validation_steps=VALIDATION_STEPS,\n",
        "                          validation_data=test_dataset)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DqMfaelrwTES"
      },
      "source": [
        "You can plot the train and validation loss to see how the training went. This should show generally decreasing values per epoch."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P_mu0SAbt40Q"
      },
      "source": [
        "# Plot the training and validation loss\n",
        "plot_metrics(\"loss\", title=\"Training vs Validation Loss\", ylim=1)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "unP3cnxo_N72"
      },
      "source": [
        "## Make predictions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7BVXldSo-0mW"
      },
      "source": [
        "The model is now ready to make some predictions. You will use the test dataset you prepared earlier to feed input images that the model has not seen before. The utilities below will help in processing the test dataset and model predictions."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MuHPDMC1yYGB"
      },
      "source": [
        "# Prediction Utilities\n",
        "\n",
        "def get_test_image_and_annotation_arrays():\n",
        "  '''\n",
        "  Unpacks the test dataset and returns the input images and segmentation masks\n",
        "  '''\n",
        "\n",
        "  ds = test_dataset.unbatch()\n",
        "  ds = ds.batch(info.splits['test'].num_examples)\n",
        "  \n",
        "  images = []\n",
        "  y_true_segments = []\n",
        "\n",
        "  for image, annotation in ds.take(1):\n",
        "    y_true_segments = annotation.numpy()\n",
        "    images = image.numpy()\n",
        "  \n",
        "  y_true_segments = y_true_segments[:(info.splits['test'].num_examples - (info.splits['test'].num_examples % BATCH_SIZE))]\n",
        "  \n",
        "  return images[:(info.splits['test'].num_examples - (info.splits['test'].num_examples % BATCH_SIZE))], y_true_segments\n",
        "\n",
        "\n",
        "def create_mask(pred_mask):\n",
        "  '''\n",
        "  Creates the segmentation mask by getting the channel with the highest probability. Remember that we\n",
        "  have 3 channels in the output of the UNet. For each pixel, the predicition will be the channel with the\n",
        "  highest probability.\n",
        "  '''\n",
        "  pred_mask = tf.argmax(pred_mask, axis=-1)\n",
        "  pred_mask = pred_mask[..., tf.newaxis]\n",
        "  return pred_mask[0].numpy()\n",
        "\n",
        "\n",
        "def make_predictions(image, mask, num=1):\n",
        "  '''\n",
        "  Feeds an image to a model and returns the predicted mask.\n",
        "  '''\n",
        "\n",
        "  image = np.reshape(image,(1, image.shape[0], image.shape[1], image.shape[2]))\n",
        "  pred_mask = model.predict(image)\n",
        "  pred_mask = create_mask(pred_mask)\n",
        "\n",
        "  return pred_mask "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Nc8XlJ19zjpz"
      },
      "source": [
        "### Compute class wise metrics\n",
        "\n",
        "Like the previous lab, you will also want to compute the IOU and Dice Score. This is the same function you used previously."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Z3hfZSSIwi1y"
      },
      "source": [
        "def class_wise_metrics(y_true, y_pred):\n",
        "  class_wise_iou = []\n",
        "  class_wise_dice_score = []\n",
        "\n",
        "  smoothening_factor = 0.00001\n",
        "  for i in range(3):\n",
        "    \n",
        "    intersection = np.sum((y_pred == i) * (y_true == i))\n",
        "    y_true_area = np.sum((y_true == i))\n",
        "    y_pred_area = np.sum((y_pred == i))\n",
        "    combined_area = y_true_area + y_pred_area\n",
        "    \n",
        "    iou = (intersection + smoothening_factor) / (combined_area - intersection + smoothening_factor)\n",
        "    class_wise_iou.append(iou)\n",
        "    \n",
        "    dice_score =  2 * ((intersection + smoothening_factor) / (combined_area + smoothening_factor))\n",
        "    class_wise_dice_score.append(dice_score)\n",
        "\n",
        "  return class_wise_iou, class_wise_dice_score"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wHze5dwDzorO"
      },
      "source": [
        "With all the utilities defined, you can now proceed to showing the metrics and feeding test images.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yEV6XHFoCDTa"
      },
      "source": [
        "# Setup the ground truth and predictions.\n",
        "\n",
        "# get the ground truth from the test set\n",
        "y_true_images, y_true_segments = get_test_image_and_annotation_arrays()\n",
        "\n",
        "# feed the test set to th emodel to get the predicted masks\n",
        "results = model.predict(test_dataset, steps=info.splits['test'].num_examples//BATCH_SIZE)\n",
        "results = np.argmax(results, axis=3)\n",
        "results = results[..., tf.newaxis]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xNcC3lBtwnsY"
      },
      "source": [
        "# compute the class wise metrics\n",
        "cls_wise_iou, cls_wise_dice_score = class_wise_metrics(y_true_segments, results)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Fgd1hmRK3mdp"
      },
      "source": [
        "# show the IOU for each class\n",
        "for idx, iou in enumerate(cls_wise_iou):\n",
        "  spaces = ' ' * (10-len(class_names[idx]) + 2)\n",
        "  print(\"{}{}{} \".format(class_names[idx], spaces, iou)) "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ECFMjlw63nHb"
      },
      "source": [
        "# show the Dice Score for each class\n",
        "for idx, dice_score in enumerate(cls_wise_dice_score):\n",
        "  spaces = ' ' * (10-len(class_names[idx]) + 2)\n",
        "  print(\"{}{}{} \".format(class_names[idx], spaces, dice_score)) \n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sAN56XW9zueE"
      },
      "source": [
        "### Show Predictions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_xzI28AfxFQi"
      },
      "source": [
        "# Please input a number between 0 to 3647 to pick an image from the dataset\n",
        "integer_slider = 3646\n",
        "\n",
        "# Get the prediction mask\n",
        "y_pred_mask = make_predictions(y_true_images[integer_slider], y_true_segments[integer_slider])\n",
        "\n",
        "# Compute the class wise metrics\n",
        "iou, dice_score = class_wise_metrics(y_true_segments[integer_slider], y_pred_mask)  \n",
        "\n",
        "# Overlay the metrics with the images\n",
        "display_with_metrics([y_true_images[integer_slider], y_pred_mask, y_true_segments[integer_slider]], iou, dice_score)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LkB4P0w24BJE"
      },
      "source": [
        "**That's all for this lab! In the next section, you will learn about another type of image segmentation model: Mask R-CNN for instance segmentation!**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1Gz4ybP4EwTr"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}