abd_10.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ilthBvnZCQto"
   },
   "source": [
    "# Algorithms for Big Data - Exercise 10\n",
    "This lecture is focused on using CNN for object localization tasks.\n",
    "\n",
    "You can download the dataset from this course on [Github](https://github.com/rasvob/2020-21-ARD/tree/master/datasets)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Fi2Jwhs35Itq"
   },
   "source": [
    "[Open in Google colab](https://colab.research.google.com/github/rasvob/2020-21-ARD/blob/master/abd_10.ipynb)\n",
    "[Download from Github](https://github.com/rasvob/2020-21-ARD/blob/master/abd_10.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import absolute_import\n",
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "\n",
    "import matplotlib.pyplot as plt # plotting\n",
    "import matplotlib.image as mpimg # images\n",
    "import numpy as np #numpy\n",
    "import seaborn as sns\n",
    "import tensorflow as tf\n",
    "# import tensorflow.compat.v2 as tf #use tensorflow v2 as a main \n",
    "import tensorflow.keras as keras # required for high level applications\n",
    "from sklearn.model_selection import train_test_split # split for validation sets\n",
    "from sklearn.metrics import accuracy_score, f1_score, confusion_matrix, classification_report\n",
    "from sklearn.preprocessing import normalize # normalization of the matrix\n",
    "import scipy\n",
    "import pandas as pd\n",
    "\n",
    "tf.version.VERSION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "from typing import List, Tuple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_history(history):\n",
    "    plt.figure()\n",
    "    for key in history.history.keys():\n",
    "        plt.plot(history.epoch, history.history[key], label=key)\n",
    "    plt.legend()\n",
    "    plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# What is Object Localization?\n",
    "Object localization is the name of the task of “classification with localization”. Namely, given an image, classify the object that appears in it, and find its location in the image, usually by using a bounding-box. \n",
    "\n",
    "In Object Localization, only a single object can appear in the image. If more than one object can appear, the task is called “Object Detection”.\n",
    "\n",
    "![model](https://github.com/rasvob/2020-21-ARD/raw/master/images/class_vs_loc.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Object Localization can be treated as a regression problem - predicting a continuous value, such as a weight or a salary. For instance, we can represent our output (a bounding-box) as a tuple of size 4, as follows:\n",
    "\n",
    "- (x,y, height, width)\n",
    "    - (x,y): the coordination of the left-top corner of the bounding box\n",
    "    - height: the height of the bounding box\n",
    "    - width: the width of the bounding box\n",
    "    \n",
    "![model](https://raw.githubusercontent.com/rasvob/2020-21-ARD/master/images/cat_bound.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Network architecture in general\n",
    "\n",
    "The coordination of the left-top corner of the bounding box must be inside the image and so do x+width and y+height. \n",
    "\n",
    "We will scale the image width and height to be 1.0. To make sure that the CNN outputs will be in the range [0,1], we will use the sigmoid activation layer - it will enforce that (x,y) will be inside the image, but not necessarily x+width and y+height. \n",
    "\n",
    "This property will be learned by the network during the training process.\n",
    "\n",
    "![model](https://raw.githubusercontent.com/rasvob/2020-21-ARD/master/images/loc_arch.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### What about the loss?\n",
    "\n",
    "The output of a sigmoid can be treated as probabilistic values, and therefore we can use **Binary crossentropy** loss. \n",
    "You can see [this](https://peltarion.com/knowledge-center/documentation/modeling-view/build-an-ai-model/loss-functions/binary-crossentropy) for more informations."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# We will start with purely synthetic use-case for educational purposes before we start to implement more complex one\n",
    "Our task will be the detection of white circles on pure black background. We will assume that the white blobs will be located in square bounding boxes for simplicity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_size = 512\n",
    "X = np.zeros((dataset_size, 128, 128, 1))\n",
    "labels = np.zeros((dataset_size, 3))\n",
    "# fill each image\n",
    "for i in range(dataset_size):\n",
    "    x = np.random.randint(8,120)\n",
    "    y = np.random.randint(8,120)\n",
    "    a = min(128 - max(x,y), min(x,y))\n",
    "    r = np.random.randint(4,a)\n",
    "    for x_i in range(128):\n",
    "      for y_i in range(128):\n",
    "        if ((x_i - x)**2) + ((y_i - y)**2) < r**2:\n",
    "          X[i, x_i, y_i,:] = 1\n",
    "    labels[i,0] = (x-r)/128.0\n",
    "    labels[i,1] = (y-r)/128.0\n",
    "    labels[i,2] = 2*r / 128.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Single example of generated image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(X[0].reshape(128, 128))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.patches import Rectangle\n",
    "def plot_pred(img,p):\n",
    "  fig, ax = plt.subplots(1)\n",
    "  ax.imshow(img.reshape(128, 128))\n",
    "  rect = Rectangle(xy=(p[1]*128,p[0]*128),width=p[2]*128, height=p[2]*128, linewidth=1,edgecolor='g',facecolor='none')\n",
    "  ax.add_patch(rect)\n",
    "  plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_pred(X[0], labels[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# More examples of our input data with bounding boxes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(8, 8, figsize=(20, 14))\n",
    "for i in range(8):\n",
    "    for j in range(8):\n",
    "        img = X[i*8 + j]\n",
    "        p = labels[i*8 + j]\n",
    "        ax[i, j].imshow(img.reshape(128, 128))\n",
    "        rect = Rectangle(xy=(p[1]*128,p[0]*128),width=p[2]*128, height=p[2]*128, linewidth=2,edgecolor='g',facecolor='none')\n",
    "        ax[i, j].add_patch(rect)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let's define our first object localization model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 637
    },
    "colab_type": "code",
    "id": "GqGF8pxYCQt2",
    "outputId": "854e6ae7-1750-4e71-bb94-5f32ec9446d4"
   },
   "outputs": [],
   "source": [
    "model = keras.Sequential([\n",
    "    keras.layers.Conv2D(64, (3,3), activation='relu', padding='same', input_shape=(128,128,1)),\n",
    "#     keras.layers.BatchNormalization(),\n",
    "    keras.layers.MaxPooling2D((2, 2), padding='same'),\n",
    "    keras.layers.Conv2D(64, (3,3), padding='same', activation='relu'),\n",
    "#     keras.layers.BatchNormalization(),\n",
    "    keras.layers.MaxPooling2D((2, 2), padding='same'),\n",
    "    keras.layers.Conv2D(16, (3,3), padding='same', activation='relu'),\n",
    "#     keras.layers.BatchNormalization(),\n",
    "    keras.layers.MaxPooling2D((2, 2), padding='same'),\n",
    "    \n",
    "    keras.layers.Flatten(),\n",
    "    keras.layers.Dense(64, activation='relu'),\n",
    "    keras.layers.Dropout(0.25),\n",
    "    keras.layers.Dense(32, activation='relu'),\n",
    "    keras.layers.Dropout(0.25),\n",
    "    keras.layers.Dense(3, activation='sigmoid'),\n",
    "])\n",
    "\n",
    "model.summary()\n",
    "\n",
    "model.compile(loss='binary_crossentropy',\n",
    "              optimizer='adam',\n",
    "              metrics=['accuracy'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "uz_cLvzw5Iub"
   },
   "source": [
    "#### Fit the model on the train data.\n",
    "Lets train the model on the training data and find the best model using the ModelCheckpoint callback to find the best model avaialble and achievable.\n",
    "\n",
    "# Beware! \n",
    "## Early stopping callback probably has bug (or perhaps it is a feature, one never knows nowadays... :)) in the of weights restoration!\n",
    "Early stopping restores best weights only if the training was stopped by the callback, if the training was finished it leaves the last trained weights in place.\n",
    "You can see [https://github.com/keras-team/keras/issues/12511](https://github.com/keras-team/keras/issues/12511) for more detailed discussion.\n",
    "\n",
    "Thanks go to Ondrej Chovanec who found about the issue on the internet.\n",
    "\n",
    "## There is a possible fix involving the ModelCheckpoint so we will try this one in this lecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_x, test_x, train_y, test_y = train_test_split(X, labels, test_size=0.2, random_state=42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(\n",
    "    filepath='weights.best.hdf5',\n",
    "    save_weights_only=True,\n",
    "    monitor='val_loss',\n",
    "    mode='auto',\n",
    "    save_best_only=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "PcNubiSECQt5"
   },
   "outputs": [],
   "source": [
    "batch_size = 32\n",
    "epochs = 100\n",
    "history = model.fit(train_x, train_y, validation_split=0.2, callbacks=[model_checkpoint_callback], epochs=epochs, batch_size=batch_size)\n",
    "\n",
    "show_history(history)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.load_weights(\"weights.best.hdf5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_loss, test_acc = model.evaluate(test_x, test_y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.evaluate(train_x, train_y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Now we can take a look at our predictions using the model\n",
    "We can see that sometimes the prediction is slightly off but usually not by much."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred = model.predict(test_x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nrows = 10\n",
    "fig, ax = plt.subplots(nrows, nrows, figsize=(20, 14))\n",
    "for i in range(nrows):\n",
    "    for j in range(nrows):\n",
    "        img = test_x[i*nrows + j]\n",
    "        p = test_y[i*nrows + j]\n",
    "        predicted = y_pred[i*nrows + j]\n",
    "        ax[i, j].imshow(img.reshape(128, 128))\n",
    "        rect = Rectangle(xy=(p[1]*128,p[0]*128),width=p[2]*128, height=p[2]*128, linewidth=2,edgecolor='g',facecolor='none')\n",
    "        ax[i, j].add_patch(rect)\n",
    "        rect = Rectangle(xy=(predicted[1]*128,predicted[0]*128),width=predicted[2]*128, height=predicted[2]*128, linewidth=2,edgecolor='r',facecolor='none')\n",
    "        ax[i, j].add_patch(rect)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Now we know the basics so we can focus on more realistic stuff\n",
    "Usually you don't want to train your own CNN model for the task, but you try to leverage transfer learning approach. Object localization is no exception. \n",
    "\n",
    "Object localization/detection is very common task and there is already wide variety of the models focused on the task. \n",
    "\n",
    "Tensorflow 2 has very nice high-level API available for these tasks. You can use already pre-trained model which can be used directly for the inference or fine-tuned.\n",
    "\n",
    "We will try to fine-tune the model for the object localization for which the model was not trained.\n",
    "\n",
    "You can read the [blog post](https://blog.tensorflow.org/2020/07/tensorflow-2-meets-object-detection-api.html) about the API or you can take a look at the [Github](https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/tf2.md)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pathlib\n",
    "\n",
    "# Clone the tensorflow models repository if it doesn't already exist\n",
    "if \"models\" in pathlib.Path.cwd().parts:\n",
    "  while \"models\" in pathlib.Path.cwd().parts:\n",
    "    os.chdir('..')\n",
    "elif not pathlib.Path('models').exists():\n",
    "  !git clone --depth 1 https://github.com/tensorflow/models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install the Object Detection API\n",
    "%%bash\n",
    "cd models/research/\n",
    "protoc object_detection/protos/*.proto --python_out=.\n",
    "cp object_detection/packages/tf2/setup.py .\n",
    "python -m pip install ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib\n",
    "\n",
    "import os\n",
    "import random\n",
    "import io\n",
    "import imageio\n",
    "import glob\n",
    "import scipy.misc\n",
    "from six import BytesIO\n",
    "from PIL import Image, ImageDraw, ImageFont\n",
    "from IPython.display import display, Javascript\n",
    "from IPython.display import Image as IPyImage\n",
    "\n",
    "from object_detection.utils import label_map_util\n",
    "from object_detection.utils import config_util\n",
    "from object_detection.utils import visualization_utils as viz_utils\n",
    "from object_detection.utils import colab_utils\n",
    "from object_detection.builders import model_builder\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_image_into_numpy_array(path):\n",
    "  \"\"\"Load an image from file into a numpy array.\n",
    "\n",
    "  Puts image into numpy array to feed into tensorflow graph.\n",
    "  Note that by convention we put it into a numpy array with shape\n",
    "  (height, width, channels), where channels=3 for RGB.\n",
    "\n",
    "  Args:\n",
    "    path: a file path.\n",
    "\n",
    "  Returns:\n",
    "    uint8 numpy array with shape (img_height, img_width, 3)\n",
    "  \"\"\"\n",
    "  img_data = tf.io.gfile.GFile(path, 'rb').read()\n",
    "  image = Image.open(BytesIO(img_data))\n",
    "  (im_width, im_height) = image.size\n",
    "  return np.array(image.getdata()).reshape(\n",
    "      (im_height, im_width, 3)).astype(np.uint8)\n",
    "\n",
    "def plot_detections(image_np,\n",
    "                    boxes,\n",
    "                    classes,\n",
    "                    scores,\n",
    "                    category_index,\n",
    "                    figsize=(12, 16),\n",
    "                    image_name=None):\n",
    "  \"\"\"Wrapper function to visualize detections.\n",
    "\n",
    "  Args:\n",
    "    image_np: uint8 numpy array with shape (img_height, img_width, 3)\n",
    "    boxes: a numpy array of shape [N, 4]\n",
    "    classes: a numpy array of shape [N]. Note that class indices are 1-based,\n",
    "      and match the keys in the label map.\n",
    "    scores: a numpy array of shape [N] or None.  If scores=None, then\n",
    "      this function assumes that the boxes to be plotted are groundtruth\n",
    "      boxes and plot all boxes as black with no classes or scores.\n",
    "    category_index: a dict containing category dictionaries (each holding\n",
    "      category index `id` and category name `name`) keyed by category indices.\n",
    "    figsize: size for the figure.\n",
    "    image_name: a name for the image file.\n",
    "  \"\"\"\n",
    "  image_np_with_annotations = image_np.copy()\n",
    "  viz_utils.visualize_boxes_and_labels_on_image_array(\n",
    "      image_np_with_annotations,\n",
    "      boxes,\n",
    "      classes,\n",
    "      scores,\n",
    "      category_index,\n",
    "      use_normalized_coordinates=True,\n",
    "      min_score_thresh=0.8)\n",
    "  if image_name:\n",
    "    plt.imsave(image_name, image_np_with_annotations)\n",
    "  else:\n",
    "    plt.imshow(image_np_with_annotations)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## We will start with some toy (literally) data consisting of 5 images of a rubber ducky.  \n",
    "\n",
    "Note that the [coco](https://cocodataset.org/#explore) dataset contains a number of animals, but notably, it does *not* contain rubber duckies (or even ducks for that matter), so this is a novel class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load images and visualize\n",
    "train_image_dir = 'models/research/object_detection/test_images/ducky/train/'\n",
    "train_images_np = []\n",
    "for i in range(1, 6):\n",
    "  image_path = os.path.join(train_image_dir, 'robertducky' + str(i) + '.jpg')\n",
    "  train_images_np.append(load_image_into_numpy_array(image_path))\n",
    "\n",
    "plt.rcParams['axes.grid'] = False\n",
    "plt.rcParams['xtick.labelsize'] = False\n",
    "plt.rcParams['ytick.labelsize'] = False\n",
    "plt.rcParams['xtick.top'] = False\n",
    "plt.rcParams['xtick.bottom'] = False\n",
    "plt.rcParams['ytick.left'] = False\n",
    "plt.rcParams['ytick.right'] = False\n",
    "plt.rcParams['figure.figsize'] = [14, 7]\n",
    "\n",
    "for idx, train_image_np in enumerate(train_images_np):\n",
    "  plt.subplot(2, 3, idx+1)\n",
    "  plt.imshow(train_image_np)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## We need to anotate the images first"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We can set bounding-boxes by hand but we will skip this step in the lecture\n",
    "# gt_boxes = []\n",
    "# colab_utils.annotate(train_images_np, box_storage_pointer=gt_boxes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gt_boxes = [\n",
    "    np.array([[0.436, 0.591, 0.629, 0.712]], dtype=np.float32),\n",
    "    np.array([[0.539, 0.583, 0.73, 0.71]], dtype=np.float32),\n",
    "    np.array([[0.464, 0.414, 0.626, 0.548]], dtype=np.float32),\n",
    "    np.array([[0.313, 0.308, 0.648, 0.526]], dtype=np.float32),\n",
    "    np.array([[0.256, 0.444, 0.484, 0.629]], dtype=np.float32)\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepare data for training\n",
    "\n",
    "Below we add the class annotations (for simplicity, we assume a single class in this colab; though it should be straightforward to extend this to handle multiple classes).  We also convert everything to the format that the training\n",
    "loop below expects (e.g., everything converted to tensors, classes converted to one-hot representations, etc.)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# By convention, our non-background classes start counting at 1.  Given\n",
    "# that we will be predicting just one class, we will therefore assign it a\n",
    "# `class id` of 1.\n",
    "duck_class_id = 1\n",
    "num_classes = 1\n",
    "\n",
    "category_index = {duck_class_id: {'id': duck_class_id, 'name': 'rubber_ducky'}}\n",
    "\n",
    "# Convert class labels to one-hot; convert everything to tensors.\n",
    "# The `label_id_offset` here shifts all classes by a certain number of indices;\n",
    "# we do this here so that the model receives one-hot labels where non-background\n",
    "# classes start counting at the zeroth index.  This is ordinarily just handled\n",
    "# automatically in our training binaries, but we need to reproduce it here.\n",
    "label_id_offset = 1\n",
    "train_image_tensors = []\n",
    "gt_classes_one_hot_tensors = []\n",
    "gt_box_tensors = []\n",
    "for (train_image_np, gt_box_np) in zip(\n",
    "    train_images_np, gt_boxes):\n",
    "  train_image_tensors.append(tf.expand_dims(tf.convert_to_tensor(\n",
    "      train_image_np, dtype=tf.float32), axis=0))\n",
    "  gt_box_tensors.append(tf.convert_to_tensor(gt_box_np, dtype=tf.float32))\n",
    "  zero_indexed_groundtruth_classes = tf.convert_to_tensor(\n",
    "      np.ones(shape=[gt_box_np.shape[0]], dtype=np.int32) - label_id_offset)\n",
    "  gt_classes_one_hot_tensors.append(tf.one_hot(\n",
    "      zero_indexed_groundtruth_classes, num_classes))\n",
    "print('Done prepping data.')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Let's just visualize the rubber duckies as a sanity check\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dummy_scores = np.array([1.0], dtype=np.float32)  # give boxes a score of 100%\n",
    "\n",
    "plt.figure(figsize=(30, 15))\n",
    "for idx in range(5):\n",
    "  plt.subplot(2, 3, idx+1)\n",
    "  plot_detections(\n",
    "      train_images_np[idx],\n",
    "      gt_boxes[idx],\n",
    "      np.ones(shape=[gt_boxes[idx].shape[0]], dtype=np.int32),\n",
    "      dummy_scores, category_index)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create model and restore weights for all but last layer\n",
    "\n",
    "In this cell we build a single stage detection architecture (RetinaNet) and restore all but the classification layer at the top (which will be automatically randomly initialized).\n",
    "\n",
    "For simplicity, we have hardcoded a number of things in this colab for the specific RetinaNet architecture at hand (including assuming that the image size will always be 640x640), however it is not difficult to generalize to other model configurations.\n",
    "\n",
    "If you are interested in the RetinaNet architecture, you can take a look at this [article](https://towardsdatascience.com/review-retinanet-focal-loss-object-detection-38fba6afabe4)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download the checkpoint and put it into models/research/object_detection/test_data/\n",
    "\n",
    "!wget http://download.tensorflow.org/models/object_detection/tf2/20200711/ssd_resnet50_v1_fpn_640x640_coco17_tpu-8.tar.gz\n",
    "!tar -xf ssd_resnet50_v1_fpn_640x640_coco17_tpu-8.tar.gz\n",
    "!mv ssd_resnet50_v1_fpn_640x640_coco17_tpu-8/checkpoint models/research/object_detection/test_data/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.keras.backend.clear_session()\n",
    "\n",
    "print('Building model and restoring weights for fine-tuning...', flush=True)\n",
    "num_classes = 1\n",
    "pipeline_config = 'models/research/object_detection/configs/tf2/ssd_resnet50_v1_fpn_640x640_coco17_tpu-8.config'\n",
    "checkpoint_path = 'models/research/object_detection/test_data/checkpoint/ckpt-0'\n",
    "\n",
    "# Load pipeline config and build a detection model.\n",
    "#\n",
    "# Since we are working off of a COCO architecture which predicts 90\n",
    "# class slots by default, we override the `num_classes` field here to be just\n",
    "# one (for our new rubber ducky class).\n",
    "configs = config_util.get_configs_from_pipeline_file(pipeline_config)\n",
    "model_config = configs['model']\n",
    "model_config.ssd.num_classes = num_classes\n",
    "model_config.ssd.freeze_batchnorm = True\n",
    "detection_model = model_builder.build(\n",
    "      model_config=model_config, is_training=True)\n",
    "\n",
    "# Set up object-based checkpoint restore --- RetinaNet has two prediction\n",
    "# `heads` --- one for classification, the other for box regression.  We will\n",
    "# restore the box regression head but initialize the classification head\n",
    "# from scratch (we show the omission below by commenting out the line that\n",
    "# we would add if we wanted to restore both heads)\n",
    "fake_box_predictor = tf.compat.v2.train.Checkpoint(\n",
    "    _base_tower_layers_for_heads=detection_model._box_predictor._base_tower_layers_for_heads,\n",
    "    # _prediction_heads=detection_model._box_predictor._prediction_heads,\n",
    "    #    (i.e., the classification head that we *will not* restore)\n",
    "    _box_prediction_head=detection_model._box_predictor._box_prediction_head,\n",
    "    )\n",
    "fake_model = tf.compat.v2.train.Checkpoint(\n",
    "          _feature_extractor=detection_model._feature_extractor,\n",
    "          _box_predictor=fake_box_predictor)\n",
    "ckpt = tf.compat.v2.train.Checkpoint(model=fake_model)\n",
    "ckpt.restore(checkpoint_path).expect_partial()\n",
    "\n",
    "# Run model through a dummy image so that variables are created\n",
    "image, shapes = detection_model.preprocess(tf.zeros([1, 640, 640, 3]))\n",
    "prediction_dict = detection_model.predict(image, shapes)\n",
    "_ = detection_model.postprocess(prediction_dict, shapes)\n",
    "print('Weights restored!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Now we can fine-tune the model!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.keras.backend.set_learning_phase(True)\n",
    "\n",
    "# These parameters can be tuned; since our training set has 5 images\n",
    "# it doesn't make sense to have a much larger batch size, though we could\n",
    "# fit more examples in memory if we wanted to.\n",
    "batch_size = 4\n",
    "learning_rate = 0.01\n",
    "num_batches = 100\n",
    "\n",
    "# Select variables in top layers to fine-tune.\n",
    "trainable_variables = detection_model.trainable_variables\n",
    "to_fine_tune = []\n",
    "prefixes_to_train = [\n",
    "  'WeightSharedConvolutionalBoxPredictor/WeightSharedConvolutionalBoxHead',\n",
    "  'WeightSharedConvolutionalBoxPredictor/WeightSharedConvolutionalClassHead']\n",
    "for var in trainable_variables:\n",
    "  if any([var.name.startswith(prefix) for prefix in prefixes_to_train]):\n",
    "    to_fine_tune.append(var)\n",
    "\n",
    "# Set up forward + backward pass for a single train step.\n",
    "def get_model_train_step_function(model, optimizer, vars_to_fine_tune):\n",
    "  \"\"\"Get a tf.function for training step.\"\"\"\n",
    "\n",
    "  # Use tf.function for a bit of speed.\n",
    "  # Comment out the tf.function decorator if you want the inside of the\n",
    "  # function to run eagerly.\n",
    "  @tf.function\n",
    "  def train_step_fn(image_tensors,\n",
    "                    groundtruth_boxes_list,\n",
    "                    groundtruth_classes_list):\n",
    "    \"\"\"A single training iteration.\n",
    "\n",
    "    Args:\n",
    "      image_tensors: A list of [1, height, width, 3] Tensor of type tf.float32.\n",
    "        Note that the height and width can vary across images, as they are\n",
    "        reshaped within this function to be 640x640.\n",
    "      groundtruth_boxes_list: A list of Tensors of shape [N_i, 4] with type\n",
    "        tf.float32 representing groundtruth boxes for each image in the batch.\n",
    "      groundtruth_classes_list: A list of Tensors of shape [N_i, num_classes]\n",
    "        with type tf.float32 representing groundtruth boxes for each image in\n",
    "        the batch.\n",
    "\n",
    "    Returns:\n",
    "      A scalar tensor representing the total loss for the input batch.\n",
    "    \"\"\"\n",
    "    shapes = tf.constant(batch_size * [[640, 640, 3]], dtype=tf.int32)\n",
    "    model.provide_groundtruth(\n",
    "        groundtruth_boxes_list=groundtruth_boxes_list,\n",
    "        groundtruth_classes_list=groundtruth_classes_list)\n",
    "    with tf.GradientTape() as tape:\n",
    "      preprocessed_images = tf.concat(\n",
    "          [detection_model.preprocess(image_tensor)[0]\n",
    "           for image_tensor in image_tensors], axis=0)\n",
    "      prediction_dict = model.predict(preprocessed_images, shapes)\n",
    "      losses_dict = model.loss(prediction_dict, shapes)\n",
    "      total_loss = losses_dict['Loss/localization_loss'] + losses_dict['Loss/classification_loss']\n",
    "      gradients = tape.gradient(total_loss, vars_to_fine_tune)\n",
    "      optimizer.apply_gradients(zip(gradients, vars_to_fine_tune))\n",
    "    return total_loss\n",
    "\n",
    "  return train_step_fn\n",
    "\n",
    "optimizer = tf.keras.optimizers.SGD(learning_rate=learning_rate, momentum=0.9)\n",
    "train_step_fn = get_model_train_step_function(\n",
    "    detection_model, optimizer, to_fine_tune)\n",
    "\n",
    "print('Start fine-tuning!', flush=True)\n",
    "for idx in range(num_batches):\n",
    "  # Grab keys for a random subset of examples\n",
    "  all_keys = list(range(len(train_images_np)))\n",
    "  random.shuffle(all_keys)\n",
    "  example_keys = all_keys[:batch_size]\n",
    "\n",
    "  # Note that we do not do data augmentation in this demo.  If you want a\n",
    "  # a fun exercise, we recommend experimenting with random horizontal flipping\n",
    "  # and random cropping :)\n",
    "  gt_boxes_list = [gt_box_tensors[key] for key in example_keys]\n",
    "  gt_classes_list = [gt_classes_one_hot_tensors[key] for key in example_keys]\n",
    "  image_tensors = [train_image_tensors[key] for key in example_keys]\n",
    "\n",
    "  # Training step (forward pass + backwards pass)\n",
    "  total_loss = train_step_fn(image_tensors, gt_boxes_list, gt_classes_list)\n",
    "\n",
    "  if idx % 10 == 0:\n",
    "    print('batch ' + str(idx) + ' of ' + str(num_batches)\n",
    "    + ', loss=' +  str(total_loss.numpy()), flush=True)\n",
    "\n",
    "print('Done fine-tuning!')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load test images and run inference with new model!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_image_dir = 'models/research/object_detection/test_images/ducky/test/'\n",
    "test_images_np = []\n",
    "for i in range(1, 50):\n",
    "  image_path = os.path.join(test_image_dir, 'out' + str(i) + '.jpg')\n",
    "  test_images_np.append(np.expand_dims(\n",
    "      load_image_into_numpy_array(image_path), axis=0))\n",
    "\n",
    "# Again, uncomment this decorator if you want to run inference eagerly\n",
    "@tf.function\n",
    "def detect(input_tensor):\n",
    "  \"\"\"Run detection on an input image.\n",
    "\n",
    "  Args:\n",
    "    input_tensor: A [1, height, width, 3] Tensor of type tf.float32.\n",
    "      Note that height and width can be anything since the image will be\n",
    "      immediately resized according to the needs of the model within this\n",
    "      function.\n",
    "\n",
    "  Returns:\n",
    "    A dict containing 3 Tensors (`detection_boxes`, `detection_classes`,\n",
    "      and `detection_scores`).\n",
    "  \"\"\"\n",
    "  preprocessed_image, shapes = detection_model.preprocess(input_tensor)\n",
    "  prediction_dict = detection_model.predict(preprocessed_image, shapes)\n",
    "  return detection_model.postprocess(prediction_dict, shapes)\n",
    "\n",
    "# Note that the first frame will trigger tracing of the tf.function, which will\n",
    "# take some time, after which inference should be fast.\n",
    "\n",
    "label_id_offset = 1\n",
    "for i in range(len(test_images_np)):\n",
    "  input_tensor = tf.convert_to_tensor(test_images_np[i], dtype=tf.float32)\n",
    "  detections = detect(input_tensor)\n",
    "\n",
    "  plot_detections(\n",
    "      test_images_np[i][0],\n",
    "      detections['detection_boxes'][0].numpy(),\n",
    "      detections['detection_classes'][0].numpy().astype(np.uint32)\n",
    "      + label_id_offset,\n",
    "      detections['detection_scores'][0].numpy(),\n",
    "      category_index, figsize=(15, 20), image_name=\"gif_frame_\" + ('%02d' % i) + \".jpg\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Let's take a look at the result!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "imageio.plugins.freeimage.download()\n",
    "\n",
    "anim_file = 'duckies_test.gif'\n",
    "\n",
    "filenames = glob.glob('gif_frame_*.jpg')\n",
    "filenames = sorted(filenames)\n",
    "last = -1\n",
    "images = []\n",
    "for filename in filenames:\n",
    "  image = imageio.imread(filename)\n",
    "  images.append(image)\n",
    "\n",
    "imageio.mimsave(anim_file, images, 'GIF-FI', fps=5)\n",
    "\n",
    "display(IPyImage(open(anim_file, 'rb').read()))"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "ds4_04.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "file_extension": ".py",
  "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.6.9"
  },
  "mimetype": "text/x-python",
  "name": "python",
  "npconvert_exporter": "python",
  "pygments_lexer": "ipython3",
  "version": 3
 },
 "nbformat": 4,
 "nbformat_minor": 4
}