MNIST_Object_Localization.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "name": "C3_W1_Lab_3_Object_Localization.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/Sujay8070/Advanced-Computer-Vision-myTasks/blob/main/MNIST_Object_Localization.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KsjDCIat6_UK"
      },
      "source": [
        "# Image Classification and Object Localization\n",
        "\n",
        "In this lab, you'll build a CNN from scratch to:\n",
        "- classify the main subject in an image\n",
        "- localize it by drawing bounding boxes around it.\n",
        "\n",
        "You'll use the [MNIST](http://yann.lecun.com/exdb/mnist/) dataset to synthesize a custom dataset for the task:\n",
        "- Place each \"digit\" image on a black canvas of width 75 x 75 at random locations.\n",
        "- Calculate the corresponding bounding boxes for those \"digits\".\n",
        "\n",
        "The bounding box prediction can be modelled as a \"regression\" task, which means that the model will predict a numeric value (as opposed to a category)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qpiJj8ym0v0-"
      },
      "source": [
        "## Imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AoilhmYe1b5t"
      },
      "source": [
        "import os, re, time, json\n",
        "import PIL.Image, PIL.ImageFont, PIL.ImageDraw\n",
        "import numpy as np\n",
        "try:\n",
        "  # %tensorflow_version only exists in Colab.\n",
        "  %tensorflow_version 2.x\n",
        "except Exception:\n",
        "  pass\n",
        "import tensorflow as tf\n",
        "from matplotlib import pyplot as plt\n",
        "import tensorflow_datasets as tfds\n",
        "\n",
        "print(\"Tensorflow version \" + tf.__version__)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xmoFKEd98MP3"
      },
      "source": [
        "# Visualization Utilities\n",
        "\n",
        "These functions are used to draw bounding boxes around the digits."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tBjj1Fg-i_lc"
      },
      "source": [
        "#@title Plot Utilities for Bounding Boxes [RUN ME]\n",
        "\n",
        "im_width = 75\n",
        "im_height = 75\n",
        "use_normalized_coordinates = True\n",
        "\n",
        "def draw_bounding_boxes_on_image_array(image,\n",
        "                                       boxes,\n",
        "                                       color=[],\n",
        "                                       thickness=1,\n",
        "                                       display_str_list=()):\n",
        "  \"\"\"Draws bounding boxes on image (numpy array).\n",
        "  Args:\n",
        "    image: a numpy array object.\n",
        "    boxes: a 2 dimensional numpy array of [N, 4]: (ymin, xmin, ymax, xmax).\n",
        "           The coordinates are in normalized format between [0, 1].\n",
        "    color: color to draw bounding box. Default is red.\n",
        "    thickness: line thickness. Default value is 4.\n",
        "    display_str_list_list: a list of strings for each bounding box.\n",
        "  Raises:\n",
        "    ValueError: if boxes is not a [N, 4] array\n",
        "  \"\"\"\n",
        "  image_pil = PIL.Image.fromarray(image)\n",
        "  rgbimg = PIL.Image.new(\"RGBA\", image_pil.size)\n",
        "  rgbimg.paste(image_pil)\n",
        "  draw_bounding_boxes_on_image(rgbimg, boxes, color, thickness,\n",
        "                               display_str_list)\n",
        "  return np.array(rgbimg)\n",
        "  \n",
        "\n",
        "def draw_bounding_boxes_on_image(image,\n",
        "                                 boxes,\n",
        "                                 color=[],\n",
        "                                 thickness=1,\n",
        "                                 display_str_list=()):\n",
        "  \"\"\"Draws bounding boxes on image.\n",
        "  Args:\n",
        "    image: a PIL.Image object.\n",
        "    boxes: a 2 dimensional numpy array of [N, 4]: (ymin, xmin, ymax, xmax).\n",
        "           The coordinates are in normalized format between [0, 1].\n",
        "    color: color to draw bounding box. Default is red.\n",
        "    thickness: line thickness. Default value is 4.\n",
        "    display_str_list: a list of strings for each bounding box.\n",
        "                           \n",
        "  Raises:\n",
        "    ValueError: if boxes is not a [N, 4] array\n",
        "  \"\"\"\n",
        "  boxes_shape = boxes.shape\n",
        "  if not boxes_shape:\n",
        "    return\n",
        "  if len(boxes_shape) != 2 or boxes_shape[1] != 4:\n",
        "    raise ValueError('Input must be of size [N, 4]')\n",
        "  for i in range(boxes_shape[0]):\n",
        "    draw_bounding_box_on_image(image, boxes[i, 1], boxes[i, 0], boxes[i, 3],\n",
        "                               boxes[i, 2], color[i], thickness, display_str_list[i])\n",
        "        \n",
        "def draw_bounding_box_on_image(image,\n",
        "                               ymin,\n",
        "                               xmin,\n",
        "                               ymax,\n",
        "                               xmax,\n",
        "                               color='red',\n",
        "                               thickness=1,\n",
        "                               display_str=None,\n",
        "                               use_normalized_coordinates=True):\n",
        "  \"\"\"Adds a bounding box to an image.\n",
        "  Bounding box coordinates can be specified in either absolute (pixel) or\n",
        "  normalized coordinates by setting the use_normalized_coordinates argument.\n",
        "  Args:\n",
        "    image: a PIL.Image object.\n",
        "    ymin: ymin of bounding box.\n",
        "    xmin: xmin of bounding box.\n",
        "    ymax: ymax of bounding box.\n",
        "    xmax: xmax of bounding box.\n",
        "    color: color to draw bounding box. Default is red.\n",
        "    thickness: line thickness. Default value is 4.\n",
        "    display_str_list: string to display in box\n",
        "    use_normalized_coordinates: If True (default), treat coordinates\n",
        "      ymin, xmin, ymax, xmax as relative to the image.  Otherwise treat\n",
        "      coordinates as absolute.\n",
        "  \"\"\"\n",
        "  draw = PIL.ImageDraw.Draw(image)\n",
        "  im_width, im_height = image.size\n",
        "  if use_normalized_coordinates:\n",
        "    (left, right, top, bottom) = (xmin * im_width, xmax * im_width,\n",
        "                                  ymin * im_height, ymax * im_height)\n",
        "  else:\n",
        "    (left, right, top, bottom) = (xmin, xmax, ymin, ymax)\n",
        "  draw.line([(left, top), (left, bottom), (right, bottom),\n",
        "             (right, top), (left, top)], width=thickness, fill=color)\n",
        "  \n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "USx9tRBF8hWy"
      },
      "source": [
        "These utilities are used to visualize the data and predictions."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qhdz68Xm3Z4Z"
      },
      "source": [
        "#@title Visualization Utilities [RUN ME]\n",
        "\"\"\"\n",
        "This cell contains helper functions used for visualization\n",
        "and downloads only. \n",
        "\n",
        "You can skip reading it, as there is very\n",
        "little Keras or Tensorflow related code here.\n",
        "\"\"\"\n",
        "\n",
        "# Matplotlib config\n",
        "plt.rc('image', cmap='gray')\n",
        "plt.rc('grid', linewidth=0)\n",
        "plt.rc('xtick', top=False, bottom=False, labelsize='large')\n",
        "plt.rc('ytick', left=False, right=False, labelsize='large')\n",
        "plt.rc('axes', facecolor='F8F8F8', titlesize=\"large\", edgecolor='white')\n",
        "plt.rc('text', color='a8151a')\n",
        "plt.rc('figure', facecolor='F0F0F0')# Matplotlib fonts\n",
        "MATPLOTLIB_FONT_DIR = os.path.join(os.path.dirname(plt.__file__), \"mpl-data/fonts/ttf\")\n",
        "\n",
        "# pull a batch from the datasets. This code is not very nice, it gets much better in eager mode (TODO)\n",
        "def dataset_to_numpy_util(training_dataset, validation_dataset, N):\n",
        "  \n",
        "  # get one batch from each: 10000 validation digits, N training digits\n",
        "  batch_train_ds = training_dataset.unbatch().batch(N)\n",
        "  \n",
        "  # eager execution: loop through datasets normally\n",
        "  if tf.executing_eagerly():\n",
        "    for validation_digits, (validation_labels, validation_bboxes) in validation_dataset:\n",
        "      validation_digits = validation_digits.numpy()\n",
        "      validation_labels = validation_labels.numpy()\n",
        "      validation_bboxes = validation_bboxes.numpy()\n",
        "      break\n",
        "    for training_digits, (training_labels, training_bboxes) in batch_train_ds:\n",
        "      training_digits = training_digits.numpy()\n",
        "      training_labels = training_labels.numpy()\n",
        "      training_bboxes = training_bboxes.numpy()\n",
        "      break\n",
        "  \n",
        "  # these were one-hot encoded in the dataset\n",
        "  validation_labels = np.argmax(validation_labels, axis=1)\n",
        "  training_labels = np.argmax(training_labels, axis=1)\n",
        "  \n",
        "  return (training_digits, training_labels, training_bboxes,\n",
        "          validation_digits, validation_labels, validation_bboxes)\n",
        "\n",
        "# create digits from local fonts for testing\n",
        "def create_digits_from_local_fonts(n):\n",
        "  font_labels = []\n",
        "  img = PIL.Image.new('LA', (75*n, 75), color = (0,255)) # format 'LA': black in channel 0, alpha in channel 1\n",
        "  font1 = PIL.ImageFont.truetype(os.path.join(MATPLOTLIB_FONT_DIR, 'DejaVuSansMono-Oblique.ttf'), 25)\n",
        "  font2 = PIL.ImageFont.truetype(os.path.join(MATPLOTLIB_FONT_DIR, 'STIXGeneral.ttf'), 25)\n",
        "  d = PIL.ImageDraw.Draw(img)\n",
        "  for i in range(n):\n",
        "    font_labels.append(i%10)\n",
        "    d.text((7+i*75,0 if i<10 else -4), str(i%10), fill=(255,255), font=font1 if i<10 else font2)\n",
        "  font_digits = np.array(img.getdata(), np.float32)[:,0] / 255.0 # black in channel 0, alpha in channel 1 (discarded)\n",
        "  font_digits = np.reshape(np.stack(np.split(np.reshape(font_digits, [75, 75*n]), n, axis=1), axis=0), [n, 75*75])\n",
        "  return font_digits, font_labels\n",
        "\n",
        "\n",
        "# utility to display a row of digits with their predictions\n",
        "def display_digits_with_boxes(digits, predictions, labels, pred_bboxes, bboxes, iou, title):\n",
        "\n",
        "  n = 10\n",
        "\n",
        "  indexes = np.random.choice(len(predictions), size=n)\n",
        "  n_digits = digits[indexes]\n",
        "  n_predictions = predictions[indexes]\n",
        "  n_labels = labels[indexes]\n",
        "\n",
        "  n_iou = []\n",
        "  if len(iou) > 0:\n",
        "    n_iou = iou[indexes]\n",
        "\n",
        "  if (len(pred_bboxes) > 0):\n",
        "    n_pred_bboxes = pred_bboxes[indexes,:]\n",
        "\n",
        "  if (len(bboxes) > 0):\n",
        "    n_bboxes = bboxes[indexes,:]\n",
        "\n",
        "\n",
        "  n_digits = n_digits * 255.0\n",
        "  n_digits = n_digits.reshape(n, 75, 75)\n",
        "  fig = plt.figure(figsize=(20, 4))\n",
        "  plt.title(title)\n",
        "  plt.yticks([])\n",
        "  plt.xticks([])\n",
        "  \n",
        "  for i in range(10):\n",
        "    ax = fig.add_subplot(1, 10, i+1)\n",
        "    bboxes_to_plot = []\n",
        "    if (len(pred_bboxes) > i):\n",
        "      bboxes_to_plot.append(n_pred_bboxes[i])\n",
        "    \n",
        "    if (len(bboxes) > i):\n",
        "      bboxes_to_plot.append(n_bboxes[i])\n",
        "\n",
        "    img_to_draw = draw_bounding_boxes_on_image_array(image=n_digits[i], boxes=np.asarray(bboxes_to_plot), color=['red', 'green'], display_str_list=[\"true\", \"pred\"])\n",
        "    plt.xlabel(n_predictions[i])\n",
        "    plt.xticks([])\n",
        "    plt.yticks([])\n",
        "    \n",
        "    if n_predictions[i] != n_labels[i]:\n",
        "      ax.xaxis.label.set_color('red')\n",
        "\n",
        "    \n",
        "    \n",
        "    plt.imshow(img_to_draw)\n",
        "\n",
        "    if len(iou) > i :\n",
        "      color = \"black\"\n",
        "      if (n_iou[i][0] < iou_threshold):\n",
        "        color = \"red\"\n",
        "      ax.text(0.2, -0.3, \"iou: %s\" %(n_iou[i][0]), color=color, transform=ax.transAxes)\n",
        "\n",
        "\n",
        "# utility to display training and validation curves\n",
        "def plot_metrics(metric_name, title, ylim=5):\n",
        "  plt.title(title)\n",
        "  plt.ylim(0,ylim)\n",
        "  plt.plot(history.history[metric_name],color='blue',label=metric_name)\n",
        "  plt.plot(history.history['val_' + metric_name],color='green',label='val_' + metric_name)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_Ok__0RB-M8S"
      },
      "source": [
        "## Selecting Between Strategies"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R4jujVYWY9-6"
      },
      "source": [
        "### TPU or GPU detection\n",
        "\n",
        "Depending on the hardware available, you'll use different distribution strategies.  For a review on distribution strategies, please check out the second course in this specialization [\"Custom and Distributed Training with TensorFlow\"](https://www.coursera.org/learn/custom-distributed-training-with-tensorflow), week 4, \"Distributed Training\".\n",
        "\n",
        "- If the TPU is available, then you'll be using the TPU Strategy.\n",
        "Otherwise:\n",
        "- If more than one GPU is available, then you'll use the Mirrored Strategy\n",
        "- If one GPU is available or if just the CPU is available, you'll use the default strategy."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Hd5zB1G7Y9-7"
      },
      "source": [
        "# Detect hardware\n",
        "try:\n",
        "  tpu = tf.distribute.cluster_resolver.TPUClusterResolver() # TPU detection\n",
        "except ValueError:\n",
        "  tpu = None\n",
        "  gpus = tf.config.experimental.list_logical_devices(\"GPU\")\n",
        "    \n",
        "# Select appropriate distribution strategy\n",
        "if tpu:\n",
        "  tf.config.experimental_connect_to_cluster(tpu)\n",
        "  tf.tpu.experimental.initialize_tpu_system(tpu)\n",
        "  strategy = tf.distribute.experimental.TPUStrategy(tpu) # Going back and forth between TPU and host is expensive. Better to run 128 batches on the TPU before reporting back.\n",
        "  print('Running on TPU ', tpu.cluster_spec().as_dict()['worker'])  \n",
        "elif len(gpus) > 1:\n",
        "  strategy = tf.distribute.MirroredStrategy([gpu.name for gpu in gpus])\n",
        "  print('Running on multiple GPUs ', [gpu.name for gpu in gpus])\n",
        "elif len(gpus) == 1:\n",
        "  strategy = tf.distribute.get_strategy() # default strategy that works on CPU and single GPU\n",
        "  print('Running on single GPU ', gpus[0].name)\n",
        "else:\n",
        "  strategy = tf.distribute.get_strategy() # default strategy that works on CPU and single GPU\n",
        "  print('Running on CPU')\n",
        "print(\"Number of accelerators: \", strategy.num_replicas_in_sync)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Lvo0t7XVIkWZ"
      },
      "source": [
        "### Parameters\n",
        "\n",
        "The global batch size is the batch size per replica (64 in this case) times the number of replicas in the distribution strategy."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cCpkS9C_H7Tl"
      },
      "source": [
        "BATCH_SIZE = 64 * strategy.num_replicas_in_sync # Gobal batch size.\n",
        "# The global batch size will be automatically sharded across all\n",
        "# replicas by the tf.data.Dataset API. A single TPU has 8 cores.\n",
        "# The best practice is to scale the batch size by the number of\n",
        "# replicas (cores). The learning rate should be increased as well.\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JVkc7nzg-WUy"
      },
      "source": [
        "## Loading and Preprocessing the Dataset\n",
        "\n",
        "Define some helper functions that will pre-process your data:\n",
        "- `read_image_tfds`: randomly overlays the \"digit\" image on top of a larger canvas.\n",
        "- `get_training_dataset`: loads data and splits it to get the training set.\n",
        "- `get_validation_dataset`: loads and splits the data to get the validation set."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZE8dgyPC1_6m"
      },
      "source": [
        "\n",
        "'''\n",
        "Transforms each image in dataset by pasting it on a 75x75 canvas at random locations.\n",
        "'''\n",
        "def read_image_tfds(image, label):\n",
        "    xmin = tf.random.uniform((), 0 , 48, dtype=tf.int32)\n",
        "    ymin = tf.random.uniform((), 0 , 48, dtype=tf.int32)\n",
        "    image = tf.reshape(image, (28,28,1,))\n",
        "    image = tf.image.pad_to_bounding_box(image, ymin, xmin, 75, 75)\n",
        "    image = tf.cast(image, tf.float32)/255.0\n",
        "    xmin = tf.cast(xmin, tf.float32)\n",
        "    ymin = tf.cast(ymin, tf.float32)\n",
        "   \n",
        "    xmax = (xmin + 28) / 75\n",
        "    ymax = (ymin + 28) / 75\n",
        "    xmin = xmin / 75\n",
        "    ymin = ymin / 75\n",
        "    return image, (tf.one_hot(label, 10), [xmin, ymin, xmax, ymax])\n",
        "  \n",
        "'''\n",
        "Loads and maps the training split of the dataset using the map function. Note that we try to load the gcs version since TPU can only work with datasets on Google Cloud Storage.\n",
        "'''\n",
        "def get_training_dataset():\n",
        "      \n",
        "      with  strategy.scope():\n",
        "        dataset = tfds.load(\"mnist\", split=\"train\", as_supervised=True, try_gcs=True)\n",
        "        dataset = dataset.map(read_image_tfds, num_parallel_calls=16)\n",
        "        dataset = dataset.shuffle(5000, reshuffle_each_iteration=True)\n",
        "        dataset = dataset.repeat() # Mandatory for Keras for now\n",
        "        dataset = dataset.batch(BATCH_SIZE, drop_remainder=True) # drop_remainder is important on TPU, batch size must be fixed\n",
        "        dataset = dataset.prefetch(-1)  # fetch next batches while training on the current one (-1: autotune prefetch buffer size)\n",
        "      return dataset\n",
        "\n",
        "'''\n",
        "Loads and maps the validation split of the dataset using the map function. Note that we try to load the gcs version since TPU can only work with datasets on Google Cloud Storage.\n",
        "'''  \n",
        "def get_validation_dataset():\n",
        "    dataset = tfds.load(\"mnist\", split=\"test\", as_supervised=True, try_gcs=True)\n",
        "    dataset = dataset.map(read_image_tfds, num_parallel_calls=16)\n",
        "\n",
        "    #dataset = dataset.cache() # this small dataset can be entirely cached in RAM\n",
        "    dataset = dataset.batch(10000, drop_remainder=True) # 10000 items in eval dataset, all in one batch\n",
        "    dataset = dataset.repeat() # Mandatory for Keras for now\n",
        "    return dataset\n",
        "\n",
        "# instantiate the datasets\n",
        "with strategy.scope():\n",
        "  training_dataset = get_training_dataset()\n",
        "  validation_dataset = get_validation_dataset()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_fXo6GuvL3EB"
      },
      "source": [
        "### Visualize Data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yZ4tjPKvL2eh"
      },
      "source": [
        "(training_digits, training_labels, training_bboxes,\n",
        " validation_digits, validation_labels, validation_bboxes) = dataset_to_numpy_util(training_dataset, validation_dataset, 10)\n",
        "\n",
        "display_digits_with_boxes(training_digits, training_labels, training_labels, np.array([]), training_bboxes, np.array([]), \"training digits and their labels\")\n",
        "display_digits_with_boxes(validation_digits, validation_labels, validation_labels, np.array([]), validation_bboxes, np.array([]), \"validation digits and their labels\")\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f8nHWWkS_eeZ"
      },
      "source": [
        "## Define the Network\n",
        "\n",
        "Here, you'll define your custom CNN. \n",
        "- `feature_extractor`: these convolutional layers extract the features of the image.\n",
        "- `classifier`:  This define the output layer that predicts among 10 categories (digits 0 through 9)\n",
        "- `bounding_box_regression`: This defines the output layer that predicts 4 numeric values, which define the coordinates of the bounding box (xmin, ymin, xmax, ymax)\n",
        "- `final_model`: This combines the layers for feature extraction, classification and bounding box prediction.  \n",
        "  - Notice that this is another example of a branching model, because the model splits to produce two kinds of output (a category and set of numbers).  \n",
        "  - Since you've learned to use the Functional API earlier in the specialization (course 1), you have the flexibility to define this kind of branching model!\n",
        "- `define_and_compile_model`: choose the optimizer and metrics, then compile the model."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "56y8UNFQIVwj"
      },
      "source": [
        "'''\n",
        "Feature extractor is the CNN that is made up of convolution and pooling layers.\n",
        "'''\n",
        "def feature_extractor(inputs):\n",
        "    x = tf.keras.layers.Conv2D(16, activation='relu', kernel_size=3, input_shape=(75, 75, 1))(inputs)\n",
        "    x = tf.keras.layers.AveragePooling2D((2, 2))(x)\n",
        "\n",
        "    x = tf.keras.layers.Conv2D(32,kernel_size=3,activation='relu')(x)\n",
        "    x = tf.keras.layers.AveragePooling2D((2, 2))(x)\n",
        "\n",
        "    x = tf.keras.layers.Conv2D(64,kernel_size=3,activation='relu')(x)\n",
        "    x = tf.keras.layers.AveragePooling2D((2, 2))(x)\n",
        "\n",
        "    return x\n",
        "\n",
        "'''\n",
        "dense_layers adds a flatten and dense layer.\n",
        "This will follow the feature extraction layers\n",
        "'''\n",
        "def dense_layers(inputs):\n",
        "  x = tf.keras.layers.Flatten()(inputs)\n",
        "  x = tf.keras.layers.Dense(128, activation='relu')(x)\n",
        "  return x\n",
        "\n",
        "\n",
        "'''\n",
        "Classifier defines the classification output.\n",
        "This has a set of fully connected layers and a softmax layer.\n",
        "'''\n",
        "def classifier(inputs):\n",
        "\n",
        "  classification_output = tf.keras.layers.Dense(10, activation='softmax', name = 'classification')(inputs)\n",
        "  return classification_output\n",
        "\n",
        "\n",
        "'''\n",
        "This function defines the regression output for bounding box prediction. \n",
        "Note that we have four outputs corresponding to (xmin, ymin, xmax, ymax)\n",
        "'''\n",
        "def bounding_box_regression(inputs):\n",
        "    bounding_box_regression_output = tf.keras.layers.Dense(units = '4', name = 'bounding_box')(inputs)\n",
        "    return bounding_box_regression_output\n",
        "\n",
        "\n",
        "def final_model(inputs):\n",
        "    feature_cnn = feature_extractor(inputs)\n",
        "    dense_output = dense_layers(feature_cnn)\n",
        "\n",
        "    '''\n",
        "    The model branches here.  \n",
        "    The dense layer's output gets fed into two branches:\n",
        "    classification_output and bounding_box_output\n",
        "    '''\n",
        "    classification_output = classifier(dense_output)\n",
        "    bounding_box_output = bounding_box_regression(dense_output)\n",
        "\n",
        "    model = tf.keras.Model(inputs = inputs, outputs = [classification_output, bounding_box_output])\n",
        "\n",
        "    return model\n",
        "  \n",
        "\n",
        "def define_and_compile_model(inputs):\n",
        "  model = final_model(inputs)\n",
        "  \n",
        "  model.compile(optimizer='adam', \n",
        "              loss = {'classification' : 'categorical_crossentropy',\n",
        "                      'bounding_box' : 'mse'\n",
        "                     },\n",
        "              metrics = {'classification' : 'accuracy',\n",
        "                         'bounding_box' : 'mse'\n",
        "                        })\n",
        "  return model\n",
        "\n",
        "    \n",
        "with strategy.scope():\n",
        "  inputs = tf.keras.layers.Input(shape=(75, 75, 1,))\n",
        "  model = define_and_compile_model(inputs)\n",
        "\n",
        "# print model layers\n",
        "model.summary()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CuhDh8ao8VyB"
      },
      "source": [
        "### Train and validate the model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kv0BQTPsKrkt"
      },
      "source": [
        "Train the model.  \n",
        "- You can choose the number of epochs depending on the level of performance that you want and the time that you have.\n",
        "- Each epoch will take just a few seconds if you're using the TPU."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TTwH_P-ZJ_xx"
      },
      "source": [
        "EPOCHS = 10 # 45\n",
        "steps_per_epoch = 60000//BATCH_SIZE  # 60,000 items in this dataset\n",
        "validation_steps = 1\n",
        "\n",
        "history = model.fit(training_dataset,\n",
        "                    steps_per_epoch=steps_per_epoch, validation_data=validation_dataset, validation_steps=validation_steps, epochs=EPOCHS)\n",
        "\n",
        "loss, classification_loss, bounding_box_loss, classification_accuracy, bounding_box_mse = model.evaluate(validation_dataset, steps=1)\n",
        "print(\"Validation accuracy: \", classification_accuracy)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Cz-b8TxU6EDj"
      },
      "source": [
        "plot_metrics(\"classification_loss\", \"Classification Loss\")\n",
        "plot_metrics(\"bounding_box_loss\", \"Bounding Box Loss\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3FBn4V5-Krkt"
      },
      "source": [
        "## Intersection over union\n",
        "\n",
        "Calculate the I-O-U metric to evaluate the model's performance."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YFqJxt3_VrCm"
      },
      "source": [
        "def intersection_over_union(pred_box, true_box):\n",
        "    xmin_pred, ymin_pred, xmax_pred, ymax_pred =  np.split(pred_box, 4, axis = 1)\n",
        "    xmin_true, ymin_true, xmax_true, ymax_true = np.split(true_box, 4, axis = 1)\n",
        "\n",
        "    smoothing_factor = 1e-10\n",
        "\n",
        "    xmin_overlap = np.maximum(xmin_pred, xmin_true)\n",
        "    xmax_overlap = np.minimum(xmax_pred, xmax_true)\n",
        "    ymin_overlap = np.maximum(ymin_pred, ymin_true)\n",
        "    ymax_overlap = np.minimum(ymax_pred, ymax_true)\n",
        "\n",
        "    pred_box_area = (xmax_pred - xmin_pred) * (ymax_pred - ymin_pred)\n",
        "    true_box_area = (xmax_true - xmin_true) * (ymax_true - ymin_true)\n",
        "\n",
        "    overlap_area = np.maximum((xmax_overlap - xmin_overlap), 0)  * np.maximum((ymax_overlap - ymin_overlap), 0)\n",
        "    union_area = (pred_box_area + true_box_area) - overlap_area\n",
        "    \n",
        "    iou = (overlap_area + smoothing_factor) / (union_area + smoothing_factor)\n",
        "\n",
        "    return iou"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9jFVovcUUVs1"
      },
      "source": [
        "### Visualize predictions\n",
        "The following code will make predictions and visualize both the classification and the predicted bounding boxes.\n",
        "- The true bounding box labels will be in green, and the model's predicted bounding boxes are in red.\n",
        "- The predicted number is shown below the image."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "w12OId8Mz7dF"
      },
      "source": [
        "# recognize validation digits\n",
        "predictions = model.predict(validation_digits, batch_size=64)\n",
        "predicted_labels = np.argmax(predictions[0], axis=1)\n",
        "\n",
        "predicted_bboxes = predictions[1]\n",
        "\n",
        "iou = intersection_over_union(predicted_bboxes, validation_bboxes)\n",
        "\n",
        "iou_threshold = 0.6\n",
        "\n",
        "print(\"Number of predictions where iou > threshold(%s): %s\" % (iou_threshold, (iou >= iou_threshold).sum()))\n",
        "print(\"Number of predictions where iou < threshold(%s): %s\" % (iou_threshold, (iou < iou_threshold).sum()))\n",
        "\n",
        "\n",
        "display_digits_with_boxes(validation_digits, predicted_labels, validation_labels, predicted_bboxes, validation_bboxes, iou, \"True and Predicted values\")"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}