site/en/hub/tutorials/tf2_object_detection.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "98rds-2OU-Rd"
      },
      "source": [
        "##### Copyright 2020 The TensorFlow Hub Authors.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "1c95xMGcU5_Z"
      },
      "outputs": [],
      "source": [
        "#@title Copyright 2020 The TensorFlow Hub Authors. All Rights Reserved.\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "# =============================================================================="
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V1UUX8SUUiMO"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/hub/tutorials/tf2_object_detection\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/hub/tutorials/tf2_object_detection.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/hub/tutorials/tf2_object_detection.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/hub/tutorials/tf2_object_detection.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://tfhub.dev/tensorflow/collections/object_detection/1\"><img src=\"https://www.tensorflow.org/images/hub_logo_32px.png\" />See TF Hub models</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rOvvWAVTkMR7"
      },
      "source": [
        "# TensorFlow Hub Object Detection Colab\n",
        "\n",
        "Welcome to the TensorFlow Hub Object Detection Colab! This notebook will take you through the steps of running an \"out-of-the-box\" object detection model on images."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IRImnk_7WOq1"
      },
      "source": [
        "### More models\n",
        "[This](https://tfhub.dev/tensorflow/collections/object_detection/1) collection contains TF2 object detection models that have been trained on the COCO 2017 dataset. [Here](https://tfhub.dev/s?module-type=image-object-detection) you can find all object detection models that are currently hosted on [tfhub.dev](https://tfhub.dev/)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vPs64QA1Zdov"
      },
      "source": [
        "## Imports and Setup\n",
        "\n",
        "Let's start with the base imports."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Xk4FU-jx9kc3"
      },
      "outputs": [],
      "source": [
        "# This Colab requires a recent numpy version.\n",
        "!pip install numpy==1.24.3\n",
        "!pip install protobuf==3.20.3\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yn5_uV1HLvaz"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import pathlib\n",
        "\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import io\n",
        "import scipy.misc\n",
        "import numpy as np\n",
        "from six import BytesIO\n",
        "from PIL import Image, ImageDraw, ImageFont\n",
        "from six.moves.urllib.request import urlopen\n",
        "\n",
        "import tensorflow as tf\n",
        "import tensorflow_hub as hub\n",
        "\n",
        "tf.get_logger().setLevel('ERROR')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IogyryF2lFBL"
      },
      "source": [
        "## Utilities\n",
        "\n",
        "Run the following cell to create some utils that will be needed later:\n",
        "\n",
        "- Helper method to load an image\n",
        "- Map of Model Name to TF Hub handle\n",
        "- List of tuples with Human Keypoints for the COCO 2017 dataset. This is needed for models with keypoints."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "-y9R0Xllefec"
      },
      "outputs": [],
      "source": [
        "# @title Run this!!\n",
        "\n",
        "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",
        "ALL_MODELS = {\n",
        "'CenterNet HourGlass104 512x512' : 'https://tfhub.dev/tensorflow/centernet/hourglass_512x512/1',\n",
        "'CenterNet HourGlass104 Keypoints 512x512' : 'https://tfhub.dev/tensorflow/centernet/hourglass_512x512_kpts/1',\n",
        "'CenterNet HourGlass104 1024x1024' : 'https://tfhub.dev/tensorflow/centernet/hourglass_1024x1024/1',\n",
        "'CenterNet HourGlass104 Keypoints 1024x1024' : 'https://tfhub.dev/tensorflow/centernet/hourglass_1024x1024_kpts/1',\n",
        "'CenterNet Resnet50 V1 FPN 512x512' : 'https://tfhub.dev/tensorflow/centernet/resnet50v1_fpn_512x512/1',\n",
        "'CenterNet Resnet50 V1 FPN Keypoints 512x512' : 'https://tfhub.dev/tensorflow/centernet/resnet50v1_fpn_512x512_kpts/1',\n",
        "'CenterNet Resnet101 V1 FPN 512x512' : 'https://tfhub.dev/tensorflow/centernet/resnet101v1_fpn_512x512/1',\n",
        "'CenterNet Resnet50 V2 512x512' : 'https://tfhub.dev/tensorflow/centernet/resnet50v2_512x512/1',\n",
        "'CenterNet Resnet50 V2 Keypoints 512x512' : 'https://tfhub.dev/tensorflow/centernet/resnet50v2_512x512_kpts/1',\n",
        "'EfficientDet D0 512x512' : 'https://tfhub.dev/tensorflow/efficientdet/d0/1',\n",
        "'EfficientDet D1 640x640' : 'https://tfhub.dev/tensorflow/efficientdet/d1/1',\n",
        "'EfficientDet D2 768x768' : 'https://tfhub.dev/tensorflow/efficientdet/d2/1',\n",
        "'EfficientDet D3 896x896' : 'https://tfhub.dev/tensorflow/efficientdet/d3/1',\n",
        "'EfficientDet D4 1024x1024' : 'https://tfhub.dev/tensorflow/efficientdet/d4/1',\n",
        "'EfficientDet D5 1280x1280' : 'https://tfhub.dev/tensorflow/efficientdet/d5/1',\n",
        "'EfficientDet D6 1280x1280' : 'https://tfhub.dev/tensorflow/efficientdet/d6/1',\n",
        "'EfficientDet D7 1536x1536' : 'https://tfhub.dev/tensorflow/efficientdet/d7/1',\n",
        "'SSD MobileNet v2 320x320' : 'https://tfhub.dev/tensorflow/ssd_mobilenet_v2/2',\n",
        "'SSD MobileNet V1 FPN 640x640' : 'https://tfhub.dev/tensorflow/ssd_mobilenet_v1/fpn_640x640/1',\n",
        "'SSD MobileNet V2 FPNLite 320x320' : 'https://tfhub.dev/tensorflow/ssd_mobilenet_v2/fpnlite_320x320/1',\n",
        "'SSD MobileNet V2 FPNLite 640x640' : 'https://tfhub.dev/tensorflow/ssd_mobilenet_v2/fpnlite_640x640/1',\n",
        "'SSD ResNet50 V1 FPN 640x640 (RetinaNet50)' : 'https://tfhub.dev/tensorflow/retinanet/resnet50_v1_fpn_640x640/1',\n",
        "'SSD ResNet50 V1 FPN 1024x1024 (RetinaNet50)' : 'https://tfhub.dev/tensorflow/retinanet/resnet50_v1_fpn_1024x1024/1',\n",
        "'SSD ResNet101 V1 FPN 640x640 (RetinaNet101)' : 'https://tfhub.dev/tensorflow/retinanet/resnet101_v1_fpn_640x640/1',\n",
        "'SSD ResNet101 V1 FPN 1024x1024 (RetinaNet101)' : 'https://tfhub.dev/tensorflow/retinanet/resnet101_v1_fpn_1024x1024/1',\n",
        "'SSD ResNet152 V1 FPN 640x640 (RetinaNet152)' : 'https://tfhub.dev/tensorflow/retinanet/resnet152_v1_fpn_640x640/1',\n",
        "'SSD ResNet152 V1 FPN 1024x1024 (RetinaNet152)' : 'https://tfhub.dev/tensorflow/retinanet/resnet152_v1_fpn_1024x1024/1',\n",
        "'Faster R-CNN ResNet50 V1 640x640' : 'https://tfhub.dev/tensorflow/faster_rcnn/resnet50_v1_640x640/1',\n",
        "'Faster R-CNN ResNet50 V1 1024x1024' : 'https://tfhub.dev/tensorflow/faster_rcnn/resnet50_v1_1024x1024/1',\n",
        "'Faster R-CNN ResNet50 V1 800x1333' : 'https://tfhub.dev/tensorflow/faster_rcnn/resnet50_v1_800x1333/1',\n",
        "'Faster R-CNN ResNet101 V1 640x640' : 'https://tfhub.dev/tensorflow/faster_rcnn/resnet101_v1_640x640/1',\n",
        "'Faster R-CNN ResNet101 V1 1024x1024' : 'https://tfhub.dev/tensorflow/faster_rcnn/resnet101_v1_1024x1024/1',\n",
        "'Faster R-CNN ResNet101 V1 800x1333' : 'https://tfhub.dev/tensorflow/faster_rcnn/resnet101_v1_800x1333/1',\n",
        "'Faster R-CNN ResNet152 V1 640x640' : 'https://tfhub.dev/tensorflow/faster_rcnn/resnet152_v1_640x640/1',\n",
        "'Faster R-CNN ResNet152 V1 1024x1024' : 'https://tfhub.dev/tensorflow/faster_rcnn/resnet152_v1_1024x1024/1',\n",
        "'Faster R-CNN ResNet152 V1 800x1333' : 'https://tfhub.dev/tensorflow/faster_rcnn/resnet152_v1_800x1333/1',\n",
        "'Faster R-CNN Inception ResNet V2 640x640' : 'https://tfhub.dev/tensorflow/faster_rcnn/inception_resnet_v2_640x640/1',\n",
        "'Faster R-CNN Inception ResNet V2 1024x1024' : 'https://tfhub.dev/tensorflow/faster_rcnn/inception_resnet_v2_1024x1024/1',\n",
        "'Mask R-CNN Inception ResNet V2 1024x1024' : 'https://tfhub.dev/tensorflow/mask_rcnn/inception_resnet_v2_1024x1024/1'\n",
        "}\n",
        "\n",
        "IMAGES_FOR_TEST = {\n",
        "  'Beach' : 'models/research/object_detection/test_images/image2.jpg',\n",
        "  'Dogs' : 'models/research/object_detection/test_images/image1.jpg',\n",
        "  # By Heiko Gorski, Source: https://commons.wikimedia.org/wiki/File:Naxos_Taverna.jpg\n",
        "  'Naxos Taverna' : 'https://upload.wikimedia.org/wikipedia/commons/6/60/Naxos_Taverna.jpg',\n",
        "  # Source: https://commons.wikimedia.org/wiki/File:The_Coleoptera_of_the_British_islands_(Plate_125)_(8592917784).jpg\n",
        "  'Beatles' : 'https://upload.wikimedia.org/wikipedia/commons/1/1b/The_Coleoptera_of_the_British_islands_%28Plate_125%29_%288592917784%29.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",
        "  # Source: https://commons.wikimedia.org/wiki/File:The_smaller_British_birds_(8053836633).jpg\n",
        "  'Birds' : 'https://upload.wikimedia.org/wikipedia/commons/0/09/The_smaller_British_birds_%288053836633%29.jpg',\n",
        "}\n",
        "\n",
        "COCO17_HUMAN_POSE_KEYPOINTS = [(0, 1),\n",
        " (0, 2),\n",
        " (1, 3),\n",
        " (2, 4),\n",
        " (0, 5),\n",
        " (0, 6),\n",
        " (5, 7),\n",
        " (7, 9),\n",
        " (6, 8),\n",
        " (8, 10),\n",
        " (5, 6),\n",
        " (5, 11),\n",
        " (6, 12),\n",
        " (11, 12),\n",
        " (11, 13),\n",
        " (13, 15),\n",
        " (12, 14),\n",
        " (14, 16)]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "14bNk1gzh0TN"
      },
      "source": [
        "## Visualization tools\n",
        "\n",
        "To visualize the images with the proper detected boxes, keypoints and segmentation, we will use the TensorFlow Object Detection API. To install it we will clone the repo."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oi28cqGGFWnY"
      },
      "outputs": [],
      "source": [
        "# Clone the tensorflow models repository\n",
        "!git clone --depth 1 https://github.com/tensorflow/models"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yX3pb_pXDjYA"
      },
      "source": [
        "Installing the Object Detection API"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NwdsBdGhFanc"
      },
      "outputs": [],
      "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 .\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3yDNgIx-kV7X"
      },
      "source": [
        "Now we can import the dependencies we will need later"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2JCeQU3fkayh"
      },
      "outputs": [],
      "source": [
        "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",
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NKtD0IeclbL5"
      },
      "source": [
        "### Load label map data (for plotting).\n",
        "\n",
        "Label maps correspond index numbers to category names, so that when our convolution network predicts `5`, we know that this corresponds to `airplane`.  Here we use internal utility functions, but anything that returns a dictionary mapping integers to appropriate string labels would be fine.\n",
        "\n",
        "We are going, for simplicity, to load from the repository that we loaded the Object Detection API code"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5mucYUS6exUJ"
      },
      "outputs": [],
      "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)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6917xnUSlp9x"
      },
      "source": [
        "## Build a detection model and load pre-trained model weights\n",
        "\n",
        "Here we will choose which Object Detection model we will use.\n",
        "Select the architecture and it will be loaded automatically.\n",
        "If you want to change the model to try other architectures later, just change the next cell and execute following ones.\n",
        "\n",
        "**Tip:** if you want to read more details about the selected model, you can follow the link (model handle) and read additional documentation on TF Hub. After you select a model, we will print the handle to make it easier."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HtwrSqvakTNn"
      },
      "outputs": [],
      "source": [
        "#@title Model Selection { display-mode: \"form\", run: \"auto\" }\n",
        "model_display_name = 'CenterNet HourGlass104 Keypoints 512x512' # @param ['CenterNet HourGlass104 512x512','CenterNet HourGlass104 Keypoints 512x512','CenterNet HourGlass104 1024x1024','CenterNet HourGlass104 Keypoints 1024x1024','CenterNet Resnet50 V1 FPN 512x512','CenterNet Resnet50 V1 FPN Keypoints 512x512','CenterNet Resnet101 V1 FPN 512x512','CenterNet Resnet50 V2 512x512','CenterNet Resnet50 V2 Keypoints 512x512','EfficientDet D0 512x512','EfficientDet D1 640x640','EfficientDet D2 768x768','EfficientDet D3 896x896','EfficientDet D4 1024x1024','EfficientDet D5 1280x1280','EfficientDet D6 1280x1280','EfficientDet D7 1536x1536','SSD MobileNet v2 320x320','SSD MobileNet V1 FPN 640x640','SSD MobileNet V2 FPNLite 320x320','SSD MobileNet V2 FPNLite 640x640','SSD ResNet50 V1 FPN 640x640 (RetinaNet50)','SSD ResNet50 V1 FPN 1024x1024 (RetinaNet50)','SSD ResNet101 V1 FPN 640x640 (RetinaNet101)','SSD ResNet101 V1 FPN 1024x1024 (RetinaNet101)','SSD ResNet152 V1 FPN 640x640 (RetinaNet152)','SSD ResNet152 V1 FPN 1024x1024 (RetinaNet152)','Faster R-CNN ResNet50 V1 640x640','Faster R-CNN ResNet50 V1 1024x1024','Faster R-CNN ResNet50 V1 800x1333','Faster R-CNN ResNet101 V1 640x640','Faster R-CNN ResNet101 V1 1024x1024','Faster R-CNN ResNet101 V1 800x1333','Faster R-CNN ResNet152 V1 640x640','Faster R-CNN ResNet152 V1 1024x1024','Faster R-CNN ResNet152 V1 800x1333','Faster R-CNN Inception ResNet V2 640x640','Faster R-CNN Inception ResNet V2 1024x1024','Mask R-CNN Inception ResNet V2 1024x1024']\n",
        "model_handle = ALL_MODELS[model_display_name]\n",
        "\n",
        "print('Selected model:'+ model_display_name)\n",
        "print('Model Handle at TensorFlow Hub: {}'.format(model_handle))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "muhUt-wWL582"
      },
      "source": [
        "## Loading the selected model from TensorFlow Hub\n",
        "\n",
        "Here we just need the model handle that was selected and use the Tensorflow Hub library to load it to memory.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rBuD07fLlcEO"
      },
      "outputs": [],
      "source": [
        "print('loading model...')\n",
        "hub_model = hub.load(model_handle)\n",
        "print('model loaded!')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GIawRDKPPnd4"
      },
      "source": [
        "## Loading an image\n",
        "\n",
        "Let's try the model on a simple image. To help with this, we provide a list of test images.\n",
        "\n",
        "Here are some simple things to try out if you are curious:\n",
        "* Try running inference on your own images, just upload them to colab and load the same way it's done in the cell below.\n",
        "* Modify some of the input images and see if detection still works.  Some simple things to try out here include flipping the image horizontally, or converting to grayscale (note that we still expect the input image to have 3 channels).\n",
        "\n",
        "**Be careful:** when using images with an alpha channel, the model expect 3 channels images and the alpha will count as a 4th.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hX-AWUQ1wIEr"
      },
      "outputs": [],
      "source": [
        "#@title Image Selection (don't forget to execute the cell!) { display-mode: \"form\"}\n",
        "selected_image = 'Beach' # @param ['Beach', 'Dogs', 'Naxos Taverna', 'Beatles', 'Phones', 'Birds']\n",
        "flip_image_horizontally = False #@param {type:\"boolean\"}\n",
        "convert_image_to_grayscale = False #@param {type:\"boolean\"}\n",
        "\n",
        "image_path = IMAGES_FOR_TEST[selected_image]\n",
        "image_np = load_image_into_numpy_array(image_path)\n",
        "\n",
        "# Flip horizontally\n",
        "if(flip_image_horizontally):\n",
        "  image_np[0] = np.fliplr(image_np[0]).copy()\n",
        "\n",
        "# Convert image to grayscale\n",
        "if(convert_image_to_grayscale):\n",
        "  image_np[0] = np.tile(\n",
        "    np.mean(image_np[0], 2, keepdims=True), (1, 1, 3)).astype(np.uint8)\n",
        "\n",
        "plt.figure(figsize=(24,32))\n",
        "plt.imshow(image_np[0])\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FTHsFjR6HNwb"
      },
      "source": [
        "## Doing the inference\n",
        "\n",
        "To do the inference we just need to call our TF Hub loaded model.\n",
        "\n",
        "Things you can try:\n",
        "* Print out `result['detection_boxes']` and try to match the box locations to the boxes in the image.  Notice that coordinates are given in normalized form (i.e., in the interval [0, 1]).\n",
        "* inspect other output keys present in the result. A full documentation can be seen on the models documentation page (pointing your browser to the model handle printed earlier)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Gb_siXKcnnGC"
      },
      "outputs": [],
      "source": [
        "# running inference\n",
        "results = hub_model(image_np)\n",
        "\n",
        "# different object detection models have additional results\n",
        "# all of them are explained in the documentation\n",
        "result = {key:value.numpy() for key,value in results.items()}\n",
        "print(result.keys())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IZ5VYaBoeeFM"
      },
      "source": [
        "## Visualizing the results\n",
        "\n",
        "Here is where we will need the TensorFlow Object Detection API to show the squares from the inference step (and the keypoints when available).\n",
        "\n",
        "the full documentation of this method can be seen [here](https://github.com/tensorflow/models/blob/master/research/object_detection/utils/visualization_utils.py)\n",
        "\n",
        "Here you can, for example, set `min_score_thresh` to other values (between 0 and 1) to allow more detections in or to filter out more detections."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2O7rV8g9s8Bz"
      },
      "outputs": [],
      "source": [
        "label_id_offset = 0\n",
        "image_np_with_detections = image_np.copy()\n",
        "\n",
        "# Use keypoints if available in detections\n",
        "keypoints, keypoint_scores = None, None\n",
        "if 'detection_keypoints' in result:\n",
        "  keypoints = result['detection_keypoints'][0]\n",
        "  keypoint_scores = result['detection_keypoint_scores'][0]\n",
        "\n",
        "viz_utils.visualize_boxes_and_labels_on_image_array(\n",
        "      image_np_with_detections[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=200,\n",
        "      min_score_thresh=.30,\n",
        "      agnostic_mode=False,\n",
        "      keypoints=keypoints,\n",
        "      keypoint_scores=keypoint_scores,\n",
        "      keypoint_edges=COCO17_HUMAN_POSE_KEYPOINTS)\n",
        "\n",
        "plt.figure(figsize=(24,32))\n",
        "plt.imshow(image_np_with_detections[0])\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qaw6Xi08NpEP"
      },
      "source": [
        "## [Optional]\n",
        "\n",
        "Among the available object detection models there's Mask R-CNN and the output of this model allows instance segmentation.\n",
        "\n",
        "To visualize it we will use the same method we did before but adding an additional parameter: `instance_masks=output_dict.get('detection_masks_reframed', None)`\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zl3qdtR1OvM_"
      },
      "outputs": [],
      "source": [
        "# Handle models with masks:\n",
        "image_np_with_mask = image_np.copy()\n",
        "\n",
        "if 'detection_masks' in result:\n",
        "  # we need to 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 bbox 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",
        "  detection_masks_reframed = tf.cast(detection_masks_reframed > 0.5,\n",
        "                                      tf.uint8)\n",
        "  result['detection_masks_reframed'] = detection_masks_reframed.numpy()\n",
        "\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=200,\n",
        "      min_score_thresh=.30,\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()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "tf2_object_detection.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}