TensorFlow: Advanced Techniques/C3_W3_Lab_3_Mask_RCNN_ImageSegmentation.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "C3_W3_Lab_3_Mask-RCNN-ImageSegmentation.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "TPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t_iHs_wm2Mhh"
      },
      "source": [
        "# Ungraded Lab: Mask R-CNN Image Segmentation Demo\n",
        "\n",
        "In this lab, you will see how to use a [Mask R-CNN](https://arxiv.org/abs/1703.06870) model from Tensorflow Hub for object detection and instance segmentation. This means that aside from the bounding boxes, the model is also able to predict segmentation masks for each instance of a class in the image. You have already encountered most of the commands here when you worked with the Object Dection API and you will see how you can use it with instance segmentation models. Let's begin!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VNlKYq2x4fa0"
      },
      "source": [
        "*Note: You should use a TPU runtime for this colab because of the processing requirements for this model. We have already enabled it for you but if you'll be using it in another colab, you can change the runtime from `Runtime --> Change runtime type` then select `TPU`.*"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m2gVpI7s53Xr"
      },
      "source": [
        "## Installation\n",
        "\n",
        "As mentioned, you will be using the Tensorflow 2 [Object Detection API](https://github.com/tensorflow/models/tree/master/research/object_detection). You can do that by cloning the [Tensorflow Model Garden](https://github.com/tensorflow/models) and installing the object detection packages just like you did in Week 2."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Q2yTMLA0yjpP"
      },
      "source": [
        "# Clone the tensorflow models repository\n",
        "!git clone --depth 1 https://github.com/tensorflow/models"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5s2C_7vmypdD"
      },
      "source": [
        "%%bash\n",
        "sudo apt install -y protobuf-compiler\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 ."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RPbB7s8hPHfO"
      },
      "source": [
        "## Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6EtGbyNc8VgS"
      },
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow_hub as hub\n",
        "\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import numpy as np\n",
        "from six import BytesIO\n",
        "from PIL import Image\n",
        "from six.moves.urllib.request import urlopen\n",
        "\n",
        "from object_detection.utils import label_map_util\n",
        "from object_detection.utils import visualization_utils as viz_utils\n",
        "from object_detection.utils import ops as utils_ops\n",
        "\n",
        "tf.get_logger().setLevel('ERROR')\n",
        "\n",
        "%matplotlib inline"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zvxMtmtE5Pro"
      },
      "source": [
        "## Utilities\n",
        "\n",
        "For convenience, you will use a function to convert an image to a numpy array. You can pass in a relative path to an image (e.g. to a local directory) or a URL. You can see this in the `TEST_IMAGES` dictionary below. Some paths point to test images that come with the API package (e.g. `Beach`) while others are URLs that point to images online (e.g. `Street`)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A-JE7Rj6ybEz"
      },
      "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: the file path to the image\n",
        "\n",
        "  Returns:\n",
        "    uint8 numpy array with shape (img_height, img_width, 3)\n",
        "  \"\"\"\n",
        "  image = None\n",
        "  if(path.startswith('http')):\n",
        "    response = urlopen(path)\n",
        "    image_data = response.read()\n",
        "    image_data = BytesIO(image_data)\n",
        "    image = Image.open(image_data)\n",
        "  else:\n",
        "    image_data = tf.io.gfile.GFile(path, 'rb').read()\n",
        "    image = Image.open(BytesIO(image_data))\n",
        "\n",
        "  (im_width, im_height) = (image.size)\n",
        "  return np.array(image.getdata()).reshape(\n",
        "      (1, im_height, im_width, 3)).astype(np.uint8)\n",
        "\n",
        "\n",
        "# dictionary with image tags as keys, and image paths as values\n",
        "TEST_IMAGES = {\n",
        "  'Beach' : 'models/research/object_detection/test_images/image2.jpg',\n",
        "  'Dogs' : 'models/research/object_detection/test_images/image1.jpg',\n",
        "  # By Américo Toledano, Source: https://commons.wikimedia.org/wiki/File:Biblioteca_Maim%C3%B3nides,_Campus_Universitario_de_Rabanales_007.jpg\n",
        "  'Phones' : 'https://upload.wikimedia.org/wikipedia/commons/thumb/0/0d/Biblioteca_Maim%C3%B3nides%2C_Campus_Universitario_de_Rabanales_007.jpg/1024px-Biblioteca_Maim%C3%B3nides%2C_Campus_Universitario_de_Rabanales_007.jpg',\n",
        "  # By 663highland, Source: https://commons.wikimedia.org/wiki/File:Kitano_Street_Kobe01s5s4110.jpg\n",
        "  'Street' : 'https://upload.wikimedia.org/wikipedia/commons/thumb/0/08/Kitano_Street_Kobe01s5s4110.jpg/2560px-Kitano_Street_Kobe01s5s4110.jpg'\n",
        "}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Jf7jVN8N-DJP"
      },
      "source": [
        "## Load the Model\n",
        "\n",
        "Tensorflow Hub provides a Mask-RCNN model that is built with the Object Detection API. You can read about the details [here](https://tfhub.dev/tensorflow/mask_rcnn/inception_resnet_v2_1024x1024/1). Let's first load the model and see how to use it for inference in the next section."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UbVcXrMay38S"
      },
      "source": [
        "model_display_name = 'Mask R-CNN Inception ResNet V2 1024x1024'\n",
        "model_handle = 'https://tfhub.dev/tensorflow/mask_rcnn/inception_resnet_v2_1024x1024/1'\n",
        "\n",
        "print('Selected model:'+ model_display_name)\n",
        "print('Model Handle at TensorFlow Hub: {}'.format(model_handle))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F7nTrTOfzNsF"
      },
      "source": [
        "# This will take 10 to 15 minutes to finish\n",
        "print('loading model...')\n",
        "hub_model = hub.load(model_handle)\n",
        "print('model loaded!')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9HS4NLZa-JmU"
      },
      "source": [
        "## Inference\n",
        "\n",
        "You will use the model you just loaded to do instance segmentation on an image. First, choose one of the test images you specified earlier and load it into a numpy array.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Jyw4ktg2zX8W"
      },
      "source": [
        "# Choose one and use as key for TEST_IMAGES below: \n",
        "# ['Beach', 'Street', 'Dogs','Phones']\n",
        "\n",
        "image_path = TEST_IMAGES['Street']\n",
        "\n",
        "image_np = load_image_into_numpy_array(image_path)\n",
        "\n",
        "plt.figure(figsize=(24,32))\n",
        "plt.imshow(image_np[1])\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XiTl--ZhAJm0"
      },
      "source": [
        "You can run inference by simply passing the numpy array of a *single* image to the model. Take note that this model does not support batching. As you've seen in the notebooks in Week 2, this will output a dictionary containing the results. These are described in the `Outputs` section of the [documentation](https://tfhub.dev/tensorflow/mask_rcnn/inception_resnet_v2_1024x1024/1)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xHHfCU-QzdLp"
      },
      "source": [
        "# run inference\n",
        "results = hub_model(image_np)\n",
        "\n",
        "# output values are tensors and we only need the numpy() \n",
        "# parameter when we visualize the results\n",
        "result = {key:value.numpy() for key,value in results.items()}\n",
        "\n",
        "# print the keys\n",
        "for key in result.keys():\n",
        "  print(key)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9NXwqHumEtpp"
      },
      "source": [
        "## Visualizing the results\n",
        "\n",
        "You can now plot the results on the original image. First, you need to create the `category_index` dictionary that will contain the class IDs and names. The model was trained on the [COCO2017 dataset](https://cocodataset.org/) and the API package has the labels saved in a different format (i.e. `mscoco_label_map.pbtxt`). You can use the [create_category_index_from_labelmap](https://github.com/tensorflow/models/blob/5ee7a4627edcbbaaeb8a564d690b5f1bc498a0d7/research/object_detection/utils/label_map_util.py#L313) internal utility function to convert this to the required dictionary format."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "B4ibUChhyujW"
      },
      "source": [
        "PATH_TO_LABELS = './models/research/object_detection/data/mscoco_label_map.pbtxt'\n",
        "category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS, use_display_name=True)\n",
        "\n",
        "# sample output\n",
        "print(category_index[1])\n",
        "print(category_index[2])\n",
        "print(category_index[4])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hd1Nhy9bIK-M"
      },
      "source": [
        "Next, you will preprocess the masks then finally plot the results.\n",
        "\n",
        "* The result dictionary contains a `detection_masks` key containing segmentation masks for each box. That will be converted first to masks that will overlay to the full image size. \n",
        "* You will also select mask pixel values that are above a certain threshold. We picked a value of `0.6` but feel free to modify this and see what results you will get. If you pick something lower, then you'll most likely notice mask pixels that are outside the object.\n",
        "* As you've seen before, you can use `visualize_boxes_and_labels_on_image_array()` to plot the results on the image. The difference this time is the parameter `instance_masks` and you will pass in the reframed detection boxes to see the segmentation masks on the image.\n",
        "\n",
        "You can see how all these are handled in the code below.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Y-e89eLPzjF2"
      },
      "source": [
        "# Handle models with masks:\n",
        "label_id_offset = 0\n",
        "image_np_with_mask = image_np.copy()\n",
        "\n",
        "if 'detection_masks' in result:\n",
        "\n",
        "  # convert np.arrays to tensors\n",
        "  detection_masks = tf.convert_to_tensor(result['detection_masks'][0])\n",
        "  detection_boxes = tf.convert_to_tensor(result['detection_boxes'][0])\n",
        "\n",
        "  # reframe the the bounding box mask to the image size.\n",
        "  detection_masks_reframed = utils_ops.reframe_box_masks_to_image_masks(\n",
        "            detection_masks, detection_boxes,\n",
        "              image_np.shape[1], image_np.shape[2])\n",
        "\n",
        "  # filter mask pixel values that are above a specified threshold\n",
        "  detection_masks_reframed = tf.cast(detection_masks_reframed > 0.6,\n",
        "                                      tf.uint8)\n",
        "  \n",
        "  # get the numpy array\n",
        "  result['detection_masks_reframed'] = detection_masks_reframed.numpy()\n",
        "\n",
        "# overlay labeled boxes and segmentation masks on the image\n",
        "viz_utils.visualize_boxes_and_labels_on_image_array(\n",
        "      image_np_with_mask[0],\n",
        "      result['detection_boxes'][0],\n",
        "      (result['detection_classes'][0] + label_id_offset).astype(int),\n",
        "      result['detection_scores'][0],\n",
        "      category_index,\n",
        "      use_normalized_coordinates=True,\n",
        "      max_boxes_to_draw=100,\n",
        "      min_score_thresh=.70,\n",
        "      agnostic_mode=False,\n",
        "      instance_masks=result.get('detection_masks_reframed', None),\n",
        "      line_thickness=8)\n",
        "\n",
        "plt.figure(figsize=(24,32))\n",
        "plt.imshow(image_np_with_mask[0])\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Alf-ddwT2ef1"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}