Fully_Convolutional_Neural_Networks_for_Image_Segmentation.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "Fully_Convolutional_Neural_Networks_for_Image_Segmentation.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [],
      "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/souvikg544/AdvancedComputerVision/blob/main/Fully_Convolutional_Neural_Networks_for_Image_Segmentation.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NSFfcSDxhIS7"
      },
      "source": [
        "# Ungraded Lab: Fully Convolutional Neural Networks for Image Segmentation\n",
        "\n",
        "This notebook illustrates how to build a Fully Convolutional Neural Network for semantic image segmentation.\n",
        "\n",
        "You will train the model on a [custom dataset](https://drive.google.com/file/d/0B0d9ZiqAgFkiOHR1NTJhWVJMNEU/view?usp=sharing) prepared by [divamgupta](https://github.com/divamgupta/image-segmentation-keras). This contains video frames from a moving vehicle and is a subsample of the [CamVid](http://mi.eng.cam.ac.uk/research/projects/VideoRec/CamVid/) dataset. \n",
        "\n",
        "You will be using a pretrained VGG-16 network for the feature extraction path, then followed by an FCN-8 network for upsampling and generating the predictions. The output will be a label map (i.e. segmentation mask) with predictions for 12 classes. Let's begin!"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Install this package to use Colab's GPU for training\n",
        "!apt install --allow-change-held-packages libcudnn8=8.4.1.50-1+cuda11.6"
      ],
      "metadata": {
        "id": "1Npv-7kDziWl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MZ3g9dJxSxmN"
      },
      "source": [
        "## Imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aifz2907kxYN"
      },
      "source": [
        "import os\n",
        "import zipfile\n",
        "import PIL.Image, PIL.ImageFont, PIL.ImageDraw\n",
        "import numpy as np\n",
        "\n",
        "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",
        "from matplotlib import pyplot as plt\n",
        "import tensorflow_datasets as tfds\n",
        "import seaborn as sns\n",
        "\n",
        "print(\"Tensorflow version \" + tf.__version__)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3sLz8mI2S62W"
      },
      "source": [
        "## Download the Dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fwC8gfP6jTZC"
      },
      "source": [
        "We hosted the dataset in a Google bucket so you will need to download it first and unzip to a local directory."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1Z5V1XMBNJso"
      },
      "source": [
        "# download the dataset (zipped file)\n",
        "!gdown --id 0B0d9ZiqAgFkiOHR1NTJhWVJMNEU -O /tmp/fcnn-dataset.zip "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "*Troubleshooting: If you get a download error saying \"Cannot retrieve the public link of the file.\", please run the next two cells below to download the dataset. Otherwise, please skip them.*"
      ],
      "metadata": {
        "id": "Ls9hGczYWBwS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%writefile download.sh\n",
        "\n",
        "#!/bin/bash\n",
        "fileid=\"0B0d9ZiqAgFkiOHR1NTJhWVJMNEU\"\n",
        "filename=\"/tmp/fcnn-dataset.zip\"\n",
        "html=`curl -c ./cookie -s -L \"https://drive.google.com/uc?export=download&id=${fileid}\"`\n",
        "curl -Lb ./cookie \"https://drive.google.com/uc?export=download&`echo ${html}|grep -Po '(confirm=[a-zA-Z0-9\\-_]+)'`&id=${fileid}\" -o ${filename}"
      ],
      "metadata": {
        "id": "fjQKhE_3WFLF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# NOTE: Please only run this if downloading with gdown did not work.\n",
        "# This will run the script created above.\n",
        "!bash download.sh"
      ],
      "metadata": {
        "id": "Nvl1E05nWMoc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "You can extract the downloaded zip files with this code:"
      ],
      "metadata": {
        "id": "sIFZoDLIWO-v"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# extract the downloaded dataset to a local directory: /tmp/fcnn\n",
        "local_zip = '/tmp/fcnn-dataset.zip'\n",
        "zip_ref = zipfile.ZipFile(local_zip, 'r')\n",
        "zip_ref.extractall('/tmp/fcnn')\n",
        "zip_ref.close()"
      ],
      "metadata": {
        "id": "-vPfaklRWC_L"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "10uOBZzfMrqZ"
      },
      "source": [
        "The dataset you just downloaded contains folders for images and annotations. The *images* contain the video frames while the *annotations* contain the pixel-wise label maps. Each label map has the shape `(height, width , 1)` with each point in this space denoting the corresponding pixel's class. Classes are in the range `[0, 11]` (i.e. 12 classes) and the pixel labels correspond to these classes:\n",
        "\n",
        "| Value  | Class Name    |\n",
        "| -------| -------------| \n",
        "| 0      | sky |\n",
        "| 1      | building      |\n",
        "| 2      | column/pole      |\n",
        "| 3      | road |\n",
        "| 4      | side walk     |\n",
        "| 5      | vegetation      |\n",
        "| 6      | traffic light |\n",
        "| 7      | fence      |\n",
        "| 8      | vehicle     |\n",
        "| 9      | pedestrian |\n",
        "| 10      | byciclist      |\n",
        "| 11      | void      |\n",
        "\n",
        "For example, if a pixel is part of a road, then that point will be labeled `3` in the label map. Run the cell below to create a list containing the class names:\n",
        "- Note: bicyclist is mispelled as 'byciclist' in the dataset.  We won't handle data cleaning in this example, but you can inspect and clean the data if you want to use this as a starting point for a personal project."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OlK0lVK_-Z46"
      },
      "source": [
        "# pixel labels in the video frames\n",
        "class_names = ['sky', 'building','column/pole', 'road', 'side walk', 'vegetation', 'traffic light', 'fence', 'vehicle', 'pedestrian', 'byciclist', 'void']"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mi0NcGqESgj2"
      },
      "source": [
        "## Load and Prepare the Dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OsRr8WWIjwTg"
      },
      "source": [
        "\n",
        "Next, you will load and prepare the train and validation sets for training. There are some preprocessing steps needed before the data is fed to the model. These include:\n",
        "\n",
        "* resizing the height and width of the input images and label maps (224 x 224px by default)\n",
        "* normalizing the input images' pixel values to fall in the range `[-1, 1]`\n",
        "* reshaping the label maps from `(height, width, 1)` to `(height, width, 12)` with each slice along the third axis having `1` if it belongs to the class corresponding to that slice's index else `0`. For example, if a pixel is part of a road, then using the table above, that point at slice #3 will be labeled `1` and it will be `0` in all other slices. To illustrate using simple arrays:\n",
        "```\n",
        "# if we have a label map with 3 classes...\n",
        "n_classes = 3\n",
        "# and this is the original annotation...\n",
        "orig_anno = [0 1 2]\n",
        "# then the reshaped annotation will have 3 slices and its contents will look like this:\n",
        "reshaped_anno = [1 0 0][0 1 0][0 0 1]\n",
        "```\n",
        "\n",
        "The following function will do the preprocessing steps mentioned above.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Lsc-_7Xu_twj"
      },
      "source": [
        "def map_filename_to_image_and_mask(t_filename, a_filename, height=224, width=224):\n",
        "  '''\n",
        "  Preprocesses the dataset by:\n",
        "    * resizing the input image and label maps\n",
        "    * normalizing the input image pixels\n",
        "    * reshaping the label maps from (height, width, 1) to (height, width, 12)\n",
        "\n",
        "  Args:\n",
        "    t_filename (string) -- path to the raw input image\n",
        "    a_filename (string) -- path to the raw annotation (label map) file\n",
        "    height (int) -- height in pixels to resize to\n",
        "    width (int) -- width in pixels to resize to\n",
        "\n",
        "  Returns:\n",
        "    image (tensor) -- preprocessed image\n",
        "    annotation (tensor) -- preprocessed annotation\n",
        "  '''\n",
        "\n",
        "  # Convert image and mask files to tensors \n",
        "  img_raw = tf.io.read_file(t_filename)\n",
        "  anno_raw = tf.io.read_file(a_filename)\n",
        "  image = tf.image.decode_jpeg(img_raw)\n",
        "  annotation = tf.image.decode_jpeg(anno_raw)\n",
        " \n",
        "  # Resize image and segmentation mask\n",
        "  image = tf.image.resize(image, (height, width,))\n",
        "  annotation = tf.image.resize(annotation, (height, width,))\n",
        "  image = tf.reshape(image, (height, width, 3,))\n",
        "  annotation = tf.cast(annotation, dtype=tf.int32)\n",
        "  annotation = tf.reshape(annotation, (height, width, 1,))\n",
        "  stack_list = []\n",
        "\n",
        "  # Reshape segmentation masks\n",
        "  for c in range(len(class_names)):\n",
        "    mask = tf.equal(annotation[:,:,0], tf.constant(c))\n",
        "    stack_list.append(tf.cast(mask, dtype=tf.int32))\n",
        "  \n",
        "  annotation = tf.stack(stack_list, axis=2)\n",
        "\n",
        "  # Normalize pixels in the input image\n",
        "  image = image/127.5\n",
        "  image -= 1\n",
        "\n",
        "  return image, annotation"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fK9tYK28lUhx"
      },
      "source": [
        "The dataset also already has separate folders for train and test sets. As described earlier, these sets will have two folders: one corresponding to the images, and the other containing the annotations. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "k5Z8UbTQGRcC"
      },
      "source": [
        "# show folders inside the dataset you downloaded\n",
        "!ls /tmp/fcnn/dataset1"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eal6745BGkV7"
      },
      "source": [
        "You will use the following functions to create the tensorflow datasets from the images in these folders. Notice that before creating the batches in the `get_training_dataset()` and `get_validation_set()`, the images are first preprocessed using the `map_filename_to_image_and_mask()` function you defined earlier."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "S8YE6w9g-ZEF"
      },
      "source": [
        "# Utilities for preparing the datasets\n",
        "\n",
        "BATCH_SIZE = 64\n",
        "\n",
        "def get_dataset_slice_paths(image_dir, label_map_dir):\n",
        "  '''\n",
        "  generates the lists of image and label map paths\n",
        "  \n",
        "  Args:\n",
        "    image_dir (string) -- path to the input images directory\n",
        "    label_map_dir (string) -- path to the label map directory\n",
        "\n",
        "  Returns:\n",
        "    image_paths (list of strings) -- paths to each image file\n",
        "    label_map_paths (list of strings) -- paths to each label map\n",
        "  '''\n",
        "  image_file_list = os.listdir(image_dir)\n",
        "  label_map_file_list = os.listdir(label_map_dir)\n",
        "  image_paths = [os.path.join(image_dir, fname) for fname in image_file_list]\n",
        "  label_map_paths = [os.path.join(label_map_dir, fname) for fname in label_map_file_list]\n",
        "\n",
        "  return image_paths, label_map_paths\n",
        "\n",
        "\n",
        "def get_training_dataset(image_paths, label_map_paths):\n",
        "  '''\n",
        "  Prepares shuffled batches of the training set.\n",
        "  \n",
        "  Args:\n",
        "    image_paths (list of strings) -- paths to each image file in the train set\n",
        "    label_map_paths (list of strings) -- paths to each label map in the train set\n",
        "\n",
        "  Returns:\n",
        "    tf Dataset containing the preprocessed train set\n",
        "  '''\n",
        "  training_dataset = tf.data.Dataset.from_tensor_slices((image_paths, label_map_paths))\n",
        "  training_dataset = training_dataset.map(map_filename_to_image_and_mask)\n",
        "  training_dataset = training_dataset.shuffle(100, reshuffle_each_iteration=True)\n",
        "  training_dataset = training_dataset.batch(BATCH_SIZE)\n",
        "  training_dataset = training_dataset.repeat()\n",
        "  training_dataset = training_dataset.prefetch(-1)\n",
        "\n",
        "  return training_dataset\n",
        "\n",
        "\n",
        "def get_validation_dataset(image_paths, label_map_paths):\n",
        "  '''\n",
        "  Prepares batches of the validation set.\n",
        "  \n",
        "  Args:\n",
        "    image_paths (list of strings) -- paths to each image file in the val set\n",
        "    label_map_paths (list of strings) -- paths to each label map in the val set\n",
        "\n",
        "  Returns:\n",
        "    tf Dataset containing the preprocessed validation set\n",
        "  '''\n",
        "  validation_dataset = tf.data.Dataset.from_tensor_slices((image_paths, label_map_paths))\n",
        "  validation_dataset = validation_dataset.map(map_filename_to_image_and_mask)\n",
        "  validation_dataset = validation_dataset.batch(BATCH_SIZE)\n",
        "  validation_dataset = validation_dataset.repeat()  \n",
        "\n",
        "  return validation_dataset\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FWnstwNEiXa0"
      },
      "source": [
        "You can now generate the training and validation sets by running the cell below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "skVGwEPmeiwz"
      },
      "source": [
        "# get the paths to the images\n",
        "training_image_paths, training_label_map_paths = get_dataset_slice_paths('/tmp/fcnn/dataset1/images_prepped_train/','/tmp/fcnn/dataset1/annotations_prepped_train/')\n",
        "validation_image_paths, validation_label_map_paths = get_dataset_slice_paths('/tmp/fcnn/dataset1/images_prepped_test/','/tmp/fcnn/dataset1/annotations_prepped_test/')\n",
        "\n",
        "# generate the train and val sets\n",
        "training_dataset = get_training_dataset(training_image_paths, training_label_map_paths)\n",
        "validation_dataset = get_validation_dataset(validation_image_paths, validation_label_map_paths)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gu9v343Pa4Eg"
      },
      "source": [
        "## Let's Take a Look at the Dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ad3a1E7JjNXc"
      },
      "source": [
        "You will also need utilities to help visualize the dataset and the model predictions later. First, you need to assign a color mapping to the classes in the label maps. Since our dataset has 12 classes, you need to have a list of 12 colors. We can use the [color_palette()](https://seaborn.pydata.org/generated/seaborn.color_palette.html) from Seaborn to generate this."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KuciRFRpajTe"
      },
      "source": [
        "# generate a list that contains one color for each class\n",
        "colors = sns.color_palette(None, len(class_names))\n",
        "\n",
        "# print class name - normalized RGB tuple pairs\n",
        "# the tuple values will be multiplied by 255 in the helper functions later\n",
        "# to convert to the (0,0,0) to (255,255,255) RGB values you might be familiar with\n",
        "for class_name, color in zip(class_names, colors):\n",
        "  print(f'{class_name} -- {color}')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "d46YCbvPafbp"
      },
      "source": [
        "# Visualization Utilities\n",
        "\n",
        "def fuse_with_pil(images):\n",
        "  '''\n",
        "  Creates a blank image and pastes input images\n",
        "\n",
        "  Args:\n",
        "    images (list of numpy arrays) - numpy array representations of the images to paste\n",
        "  \n",
        "  Returns:\n",
        "    PIL Image object containing the images\n",
        "  '''\n",
        "\n",
        "  widths = (image.shape[1] for image in images)\n",
        "  heights = (image.shape[0] for image in images)\n",
        "  total_width = sum(widths)\n",
        "  max_height = max(heights)\n",
        "\n",
        "  new_im = PIL.Image.new('RGB', (total_width, max_height))\n",
        "\n",
        "  x_offset = 0\n",
        "  for im in images:\n",
        "    pil_image = PIL.Image.fromarray(np.uint8(im))\n",
        "    new_im.paste(pil_image, (x_offset,0))\n",
        "    x_offset += im.shape[1]\n",
        "  \n",
        "  return new_im\n",
        "\n",
        "\n",
        "def give_color_to_annotation(annotation):\n",
        "  '''\n",
        "  Converts a 2-D annotation to a numpy array with shape (height, width, 3) where\n",
        "  the third axis represents the color channel. The label values are multiplied by\n",
        "  255 and placed in this axis to give color to the annotation\n",
        "\n",
        "  Args:\n",
        "    annotation (numpy array) - label map array\n",
        "  \n",
        "  Returns:\n",
        "    the annotation array with an additional color channel/axis\n",
        "  '''\n",
        "  seg_img = np.zeros( (annotation.shape[0],annotation.shape[1], 3) ).astype('float')\n",
        "  \n",
        "  for c in range(12):\n",
        "    segc = (annotation == c)\n",
        "    seg_img[:,:,0] += segc*( colors[c][0] * 255.0)\n",
        "    seg_img[:,:,1] += segc*( colors[c][1] * 255.0)\n",
        "    seg_img[:,:,2] += segc*( colors[c][2] * 255.0)\n",
        "  \n",
        "  return seg_img\n",
        "\n",
        "\n",
        "def show_predictions(image, labelmaps, titles, iou_list, dice_score_list):\n",
        "  '''\n",
        "  Displays the images with the ground truth and predicted label maps\n",
        "\n",
        "  Args:\n",
        "    image (numpy array) -- the input image\n",
        "    labelmaps (list of arrays) -- contains the predicted and ground truth label maps\n",
        "    titles (list of strings) -- display headings for the images to be displayed\n",
        "    iou_list (list of floats) -- the IOU values for each class\n",
        "    dice_score_list (list of floats) -- the Dice Score for each vlass\n",
        "  '''\n",
        "\n",
        "  true_img = give_color_to_annotation(labelmaps[1])\n",
        "  pred_img = give_color_to_annotation(labelmaps[0])\n",
        "\n",
        "  image = image + 1\n",
        "  image = image * 127.5\n",
        "  images = np.uint8([image, pred_img, true_img])\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",
        "  plt.figure(figsize=(15, 4))\n",
        "\n",
        "  for idx, im in enumerate(images):\n",
        "    plt.subplot(1, 3, idx+1)\n",
        "    if idx == 1:\n",
        "      plt.xlabel(display_string)\n",
        "    plt.xticks([])\n",
        "    plt.yticks([])\n",
        "    plt.title(titles[idx], fontsize=12)\n",
        "    plt.imshow(im)\n",
        "\n",
        "\n",
        "def show_annotation_and_image(image, annotation):\n",
        "  '''\n",
        "  Displays the image and its annotation side by side\n",
        "\n",
        "  Args:\n",
        "    image (numpy array) -- the input image\n",
        "    annotation (numpy array) -- the label map\n",
        "  '''\n",
        "  new_ann = np.argmax(annotation, axis=2)\n",
        "  seg_img = give_color_to_annotation(new_ann)\n",
        "  \n",
        "  image = image + 1\n",
        "  image = image * 127.5\n",
        "  image = np.uint8(image)\n",
        "  images = [image, seg_img]\n",
        "  \n",
        "  images = [image, seg_img]\n",
        "  fused_img = fuse_with_pil(images)\n",
        "  plt.imshow(fused_img)\n",
        "\n",
        "\n",
        "def list_show_annotation(dataset):\n",
        "  '''\n",
        "  Displays images and its annotations side by side\n",
        "\n",
        "  Args:\n",
        "    dataset (tf Dataset) - batch of images and annotations\n",
        "  '''\n",
        "\n",
        "  ds = dataset.unbatch()\n",
        "  ds = ds.shuffle(buffer_size=100)\n",
        "\n",
        "  plt.figure(figsize=(25, 15))\n",
        "  plt.title(\"Images And Annotations\")\n",
        "  plt.subplots_adjust(bottom=0.1, top=0.9, hspace=0.05)\n",
        "\n",
        "  # we set the number of image-annotation pairs to 9\n",
        "  # feel free to make this a function parameter if you want\n",
        "  for idx, (image, annotation) in enumerate(ds.take(9)):\n",
        "    plt.subplot(3, 3, idx + 1)\n",
        "    plt.yticks([])\n",
        "    plt.xticks([])\n",
        "    show_annotation_and_image(image.numpy(), annotation.numpy())\n",
        "\n",
        " "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cK621Jm8bJyj"
      },
      "source": [
        "\n",
        "Please run the cells below to see sample images from the train and validation sets. You will see the image and the label maps side side by side."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AFO_hIhLWYT4"
      },
      "source": [
        "list_show_annotation(training_dataset)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bdgVkp8wZua0"
      },
      "source": [
        "list_show_annotation(validation_dataset)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QFv2k8xabRb8"
      },
      "source": [
        "## Define the Model\n",
        "\n",
        "You will now build the model and prepare it for training. AS mentioned earlier, this will use a VGG-16 network for the encoder and FCN-8 for the decoder. This is the diagram as shown in class:\n",
        "\n",
        "<img src='https://drive.google.com/uc?export=view&id=1lrqB4YegV8jXWNfyYAaeuFlwXIc54aRP' alt='fcn-8'>\n",
        "\n",
        "For this exercise, you will notice a slight difference from the lecture because the dataset images are 224x224 instead of 32x32. You'll see how this is handled in the next cells as you build the encoder."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HHB1BGmF965d"
      },
      "source": [
        "### Define Pooling Block of VGG"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TWnI0IcNxNRh"
      },
      "source": [
        "As you saw in Course 1 of this specialization, VGG networks have repeating blocks so to make the code neat, it's best to create a function to encapsulate this process. Each block has convolutional layers followed by a max pooling layer which downsamples the image."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pL578pjdmXXf"
      },
      "source": [
        "def block(x, n_convs, filters, kernel_size, activation, pool_size, pool_stride, block_name):\n",
        "  '''\n",
        "  Defines a block in the VGG network.\n",
        "\n",
        "  Args:\n",
        "    x (tensor) -- input image\n",
        "    n_convs (int) -- number of convolution layers to append\n",
        "    filters (int) -- number of filters for the convolution layers\n",
        "    activation (string or object) -- activation to use in the convolution\n",
        "    pool_size (int) -- size of the pooling layer\n",
        "    pool_stride (int) -- stride of the pooling layer\n",
        "    block_name (string) -- name of the block\n",
        "\n",
        "  Returns:\n",
        "    tensor containing the max-pooled output of the convolutions\n",
        "  '''\n",
        "\n",
        "  for i in range(n_convs):\n",
        "      x = tf.keras.layers.Conv2D(filters=filters, kernel_size=kernel_size, activation=activation, padding='same', name=\"{}_conv{}\".format(block_name, i + 1))(x)\n",
        "    \n",
        "  x = tf.keras.layers.MaxPooling2D(pool_size=pool_size, strides=pool_stride, name=\"{}_pool{}\".format(block_name, i+1 ))(x)\n",
        "\n",
        "  return x\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Lm_8Jp4PbVV5"
      },
      "source": [
        "### Download VGG weights\n",
        "\n",
        "First, please run the cell below to get pre-trained weights for VGG-16. You will load this in the next section when you build the encoder network."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vKPpXapoYxAc"
      },
      "source": [
        "# download the weights\n",
        "!wget https://github.com/fchollet/deep-learning-models/releases/download/v0.1/vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5\n",
        "\n",
        "# assign to a variable\n",
        "vgg_weights_path = \"vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZLeQCxf99_tn"
      },
      "source": [
        "### Define VGG-16\n",
        "\n",
        "You can build the encoder as shown below. \n",
        "\n",
        "* You will create 5 blocks with increasing number of filters at each stage. \n",
        "* The number of convolutions, filters, kernel size, activation, pool size and pool stride will remain constant.\n",
        "* You will load the pretrained weights after creating the VGG 16 network.\n",
        "* Additional convolution layers will be appended to extract more features.\n",
        "* The output will contain the output of the last layer and the previous four convolution blocks."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Z4_WZnAmoOnZ"
      },
      "source": [
        "def VGG_16(image_input):\n",
        "  '''\n",
        "  This function defines the VGG encoder.\n",
        "\n",
        "  Args:\n",
        "    image_input (tensor) - batch of images\n",
        "\n",
        "  Returns:\n",
        "    tuple of tensors - output of all encoder blocks plus the final convolution layer\n",
        "  '''\n",
        "\n",
        "  # create 5 blocks with increasing filters at each stage. \n",
        "  # you will save the output of each block (i.e. p1, p2, p3, p4, p5). \"p\" stands for the pooling layer.\n",
        "  x = block(image_input,n_convs=2, filters=64, kernel_size=(3,3), activation='relu',pool_size=(2,2), pool_stride=(2,2), block_name='block1')\n",
        "  p1= x\n",
        "\n",
        "  x = block(x,n_convs=2, filters=128, kernel_size=(3,3), activation='relu',pool_size=(2,2), pool_stride=(2,2), block_name='block2')\n",
        "  p2 = x\n",
        "\n",
        "  x = block(x,n_convs=3, filters=256, kernel_size=(3,3), activation='relu',pool_size=(2,2), pool_stride=(2,2), block_name='block3')\n",
        "  p3 = x\n",
        "\n",
        "  x = block(x,n_convs=3, filters=512, kernel_size=(3,3), activation='relu',pool_size=(2,2), pool_stride=(2,2), block_name='block4')\n",
        "  p4 = x\n",
        "\n",
        "  x = block(x,n_convs=3, filters=512, kernel_size=(3,3), activation='relu',pool_size=(2,2), pool_stride=(2,2), block_name='block5')\n",
        "  p5 = x\n",
        "\n",
        "  # create the vgg model\n",
        "  vgg  = tf.keras.Model(image_input , p5)\n",
        "\n",
        "  # load the pretrained weights you downloaded earlier\n",
        "  vgg.load_weights(vgg_weights_path) \n",
        "\n",
        "  # number of filters for the output convolutional layers\n",
        "  n = 4096\n",
        "\n",
        "  # our input images are 224x224 pixels so they will be downsampled to 7x7 after the pooling layers above.\n",
        "  # we can extract more features by chaining two more convolution layers.\n",
        "  c6 = tf.keras.layers.Conv2D( n , ( 7 , 7 ) , activation='relu' , padding='same', name=\"conv6\")(p5)\n",
        "  c7 = tf.keras.layers.Conv2D( n , ( 1 , 1 ) , activation='relu' , padding='same', name=\"conv7\")(c6)\n",
        "\n",
        "  # return the outputs at each stage. you will only need two of these in this particular exercise \n",
        "  # but we included it all in case you want to experiment with other types of decoders.\n",
        "  return (p1, p2, p3, p4, c7)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "45pH17d__KUW"
      },
      "source": [
        "### Define FCN 8 Decoder\n",
        "\n",
        "Next, you will build the decoder using deconvolution layers. Please refer to the diagram for FCN-8 at the start of this section to visualize what the code below is doing. It will involve two summations before upsampling to the original image size and generating the predicted mask."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hX2V0E2gs-ZQ"
      },
      "source": [
        "def fcn8_decoder(convs, n_classes):\n",
        "  '''\n",
        "  Defines the FCN 8 decoder.\n",
        "\n",
        "  Args:\n",
        "    convs (tuple of tensors) - output of the encoder network\n",
        "    n_classes (int) - number of classes\n",
        "\n",
        "  Returns:\n",
        "    tensor with shape (height, width, n_classes) containing class probabilities\n",
        "  '''\n",
        "\n",
        "  # unpack the output of the encoder\n",
        "  f1, f2, f3, f4, f5 = convs\n",
        "  \n",
        "  # upsample the output of the encoder then crop extra pixels that were introduced\n",
        "  o = tf.keras.layers.Conv2DTranspose(n_classes , kernel_size=(4,4) ,  strides=(2,2) , use_bias=False )(f5)\n",
        "  o = tf.keras.layers.Cropping2D(cropping=(1,1))(o)\n",
        "\n",
        "  # load the pool 4 prediction and do a 1x1 convolution to reshape it to the same shape of `o` above\n",
        "  o2 = f4\n",
        "  o2 = ( tf.keras.layers.Conv2D(n_classes , ( 1 , 1 ) , activation='relu' , padding='same'))(o2)\n",
        "\n",
        "  # add the results of the upsampling and pool 4 prediction\n",
        "  o = tf.keras.layers.Add()([o, o2])\n",
        "\n",
        "  # upsample the resulting tensor of the operation you just did\n",
        "  o = (tf.keras.layers.Conv2DTranspose( n_classes , kernel_size=(4,4) ,  strides=(2,2) , use_bias=False ))(o)\n",
        "  o = tf.keras.layers.Cropping2D(cropping=(1, 1))(o)\n",
        "\n",
        "  # load the pool 3 prediction and do a 1x1 convolution to reshape it to the same shape of `o` above\n",
        "  o2 = f3\n",
        "  o2 = ( tf.keras.layers.Conv2D(n_classes , ( 1 , 1 ) , activation='relu' , padding='same'))(o2)\n",
        "\n",
        "  # add the results of the upsampling and pool 3 prediction\n",
        "  o = tf.keras.layers.Add()([o, o2])\n",
        "  \n",
        "  # upsample up to the size of the original image\n",
        "  o = tf.keras.layers.Conv2DTranspose(n_classes , kernel_size=(8,8) ,  strides=(8,8) , use_bias=False )(o)\n",
        "\n",
        "  # append a softmax to get the class probabilities\n",
        "  o = (tf.keras.layers.Activation('softmax'))(o)\n",
        "\n",
        "  return o"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yyn3xXSf_Ogl"
      },
      "source": [
        "### Define Final Model\n",
        "\n",
        "You can now build the final model by connecting the encoder and decoder blocks."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "T29n8_dbuZNm"
      },
      "source": [
        "def segmentation_model():\n",
        "  '''\n",
        "  Defines the final segmentation model by chaining together the encoder and decoder.\n",
        "\n",
        "  Returns:\n",
        "    keras Model that connects the encoder and decoder networks of the segmentation model\n",
        "  '''\n",
        "  \n",
        "  inputs = tf.keras.layers.Input(shape=(224,224,3,))\n",
        "  convs = VGG_16(image_input=inputs)\n",
        "  outputs = fcn8_decoder(convs, 12)\n",
        "  model = tf.keras.Model(inputs=inputs, outputs=outputs)\n",
        "  \n",
        "  return model\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_w8qNGG1vQHZ"
      },
      "source": [
        "# instantiate the model and see how it looks\n",
        "model = segmentation_model()\n",
        "model.summary()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dekOKLw0_Rgg"
      },
      "source": [
        "### Compile the Model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6cAcCHkyxqt4"
      },
      "source": [
        "Next, the model will be configured for training. You will need to specify the loss, optimizer and metrics. You will use `categorical_crossentropy` as the loss function since the label map is transformed to one hot encoded vectors for each pixel in the image (i.e. `1` in one slice and `0` for other slices as described earlier)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZpWpp8h4g_rE"
      },
      "source": [
        "sgd = tf.keras.optimizers.SGD(lr=1E-2, momentum=0.9, nesterov=True)\n",
        "\n",
        "model.compile(loss='categorical_crossentropy',\n",
        "              optimizer=sgd,\n",
        "              metrics=['accuracy'])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L9zxLlNZ_XbT"
      },
      "source": [
        "## Train the Model\n",
        "\n",
        "The model can now be trained. This will take around 30 minutes to run and you will reach around 85% accuracy for both train and val sets."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8HoZwpGWhMB-"
      },
      "source": [
        "# number of training images\n",
        "train_count = 367\n",
        "\n",
        "# number of validation images\n",
        "validation_count = 101\n",
        "\n",
        "EPOCHS = 50\n",
        "\n",
        "steps_per_epoch = train_count//BATCH_SIZE\n",
        "validation_steps = validation_count//BATCH_SIZE\n",
        "\n",
        "history = model.fit(training_dataset,\n",
        "                    steps_per_epoch=steps_per_epoch, validation_data=validation_dataset, validation_steps=validation_steps, epochs=EPOCHS)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y1luX1e7_bEd"
      },
      "source": [
        "## Evaluate the Model\n",
        "\n",
        "After training, you will want to see how your model is doing on a test set. For segmentation models, you can use the intersection-over-union and the dice score as metrics to evaluate your model. You'll see how it is implemented in this section."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zENjQuK0luH5"
      },
      "source": [
        "def get_images_and_segments_test_arrays():\n",
        "  '''\n",
        "  Gets a subsample of the val set as your test set\n",
        "\n",
        "  Returns:\n",
        "    Test set containing ground truth images and label maps\n",
        "  '''\n",
        "  y_true_segments = []\n",
        "  y_true_images = []\n",
        "  test_count = 64\n",
        "\n",
        "  ds = validation_dataset.unbatch()\n",
        "  ds = ds.batch(101)\n",
        "\n",
        "  for image, annotation in ds.take(1):\n",
        "    y_true_images = image\n",
        "    y_true_segments = annotation\n",
        "\n",
        "\n",
        "  y_true_segments = y_true_segments[:test_count, : ,: , :]\n",
        "  y_true_segments = np.argmax(y_true_segments, axis=3)  \n",
        "\n",
        "  return y_true_images, y_true_segments\n",
        "\n",
        "# load the ground truth images and segmentation masks\n",
        "y_true_images, y_true_segments = get_images_and_segments_test_arrays()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ly1ErxSA_kpb"
      },
      "source": [
        "### Make Predictions\n",
        "\n",
        "You can get output segmentation masks by using the `predict()` method. As you may recall, the output of our segmentation model has the shape `(height, width, 12)` where `12` is the number of classes. Each pixel value in those 12 slices indicates the probability of that pixel belonging to that particular class. If you want to create the predicted label map, then you can get the `argmax()` of that axis. This is shown in the following cell."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1CpEeUiN7ey9"
      },
      "source": [
        "# get the model prediction\n",
        "results = model.predict(validation_dataset, steps=validation_steps)\n",
        "\n",
        "# for each pixel, get the slice number which has the highest probability\n",
        "results = np.argmax(results, axis=3)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c-nmUQkp_dc6"
      },
      "source": [
        "Compute Metrics\n",
        "\n",
        "The function below generates the IOU and dice score of the prediction and ground truth masks. From the lectures, it is given that:\n",
        "\n",
        "$$IOU = \\frac{area\\_of\\_overlap}{area\\_of\\_union}$$\n",
        "<br>\n",
        "$$Dice Score = 2 * \\frac{area\\_of\\_overlap}{combined\\_area}$$\n",
        "\n",
        "The code below does that for you. A small smoothening factor is introduced in the denominators to prevent possible division by zero."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EobztGe_66sA"
      },
      "source": [
        "def compute_metrics(y_true, y_pred):\n",
        "  '''\n",
        "  Computes IOU and Dice Score.\n",
        "\n",
        "  Args:\n",
        "    y_true (tensor) - ground truth label map\n",
        "    y_pred (tensor) - predicted label map\n",
        "  '''\n",
        "  \n",
        "  class_wise_iou = []\n",
        "  class_wise_dice_score = []\n",
        "\n",
        "  smoothening_factor = 0.00001\n",
        "\n",
        "  for i in range(12):\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": "duS-cSFMy1VH"
      },
      "source": [
        "### Show Predictions and Metrics\n",
        "\n",
        "You can now see the predicted segmentation masks side by side with the ground truth. The metrics are also overlayed so you can evaluate how your model is doing."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Hkbsk_P1fpRM"
      },
      "source": [
        "# input a number from 0 to 63 to pick an image from the test set\n",
        "integer_slider = 0\n",
        "\n",
        "# compute metrics\n",
        "iou, dice_score = compute_metrics(y_true_segments[integer_slider], results[integer_slider])  \n",
        "\n",
        "# visualize the output and metrics\n",
        "show_predictions(y_true_images[integer_slider], [results[integer_slider], y_true_segments[integer_slider]], [\"Image\", \"Predicted Mask\", \"True Mask\"], iou, dice_score)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "psu5U4jRy5px"
      },
      "source": [
        "### Display Class Wise Metrics\n",
        "\n",
        "You can also compute the class-wise metrics so you can see how your model performs across all images in the test set."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XqzDRh0e6_8G"
      },
      "source": [
        "# compute class-wise metrics\n",
        "cls_wise_iou, cls_wise_dice_score = compute_metrics(y_true_segments, results)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7mnS0UPtsMeB"
      },
      "source": [
        "# print IOU for each class\n",
        "for idx, iou in enumerate(cls_wise_iou):\n",
        "  spaces = ' ' * (13-len(class_names[idx]) + 2)\n",
        "  print(\"{}{}{} \".format(class_names[idx], spaces, iou)) \n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MVxYk02pJm8O"
      },
      "source": [
        "# print the dice score for each class\n",
        "for idx, dice_score in enumerate(cls_wise_dice_score):\n",
        "  spaces = ' ' * (13-len(class_names[idx]) + 2)\n",
        "  print(\"{}{}{} \".format(class_names[idx], spaces, dice_score)) \n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nm6399aII-m-"
      },
      "source": [
        "**That's all for this lab! In the next section, you will work on another architecture for building a segmentation model: the UNET.**"
      ]
    }
  ]
}