With keypoints: Inference TF2 colab (#8819)

research/object_detection/colab_tutorials/inference_tf2_colab.ipynb

@@ -167,7 +167,39 @@
         "  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)"
+        "      (im_height, im_width, 3)).astype(np.uint8)\n",
+        "\n",
+        "def get_keypoint_tuples(eval_config):\n",
+        "  \"\"\"Return a tuple list of keypoint edges from the eval config.\n",
+        "  \n",
+        "  Args:\n",
+        "    eval_config: an eval config containing the keypoint edges\n",
+        "  \n",
+        "  Returns:\n",
+        "    a list of edge tuples, each in the format (start, end)\n",
+        "  \"\"\"\n",
+        "  tuple_list = []\n",
+        "  kp_list = eval_config.keypoint_edge\n",
+        "  for edge in kp_list:\n",
+        "    tuple_list.append((edge.start, edge.end))\n",
+        "  return tuple_list"
+      ],
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "metadata": {
+        "id": "R4YjnOjME1gy",
+        "colab_type": "code",
+        "colab": {}
+      },
+      "source": [
+        "# @title Choose the model to use, then evaluate the cell.\n",
+        "MODELS = {'centernet_with_keypoints': 'center_net_resnet101_v1_fpn_512x512_kpts_coco17_tpu-8', 'centernet_without_keypoints': 'center_net_resnet101_v1_fpn_512x512_coco17_tpu-8'}\n",
+        "\n",
+        "model_display_name = 'centernet_with_keypoints' # @param ['centernet_with_keypoints', 'centernet_without_keypoints']\n",
+        "model_name = MODELS[model_display_name]"
       ],
       "execution_count": null,
       "outputs": []
@@ -205,7 +237,8 @@
         "colab": {}
       },
       "source": [
-        "pipeline_config = 'models/research/object_detection/configs/tf2/center_net_resnet101_v1_fpn_512x512_coco17_tpu-8.config'\n",
+        "pipeline_config = os.path.join('models/research/object_detection/configs/tf2/',\n",
+        "                                model_name + '.config')\n",
         "model_dir = 'models/research/object_detection/test_data/checkpoint/'\n",
         "\n",
         "# Load pipeline config and build a detection model\n",
@@ -280,7 +313,7 @@
       "source": [
         "### Putting everything together!\n",
         "\n",
-        "Run the below code which loads an image, runs it through the detection model and visualizes the detection results.\n",
+        "Run the below code which loads an image, runs it through the detection model and visualizes the detection results, including the keypoints.\n",
         "\n",
         "Note that this will take a long time (several minutes) the first time you run this code due to tf.function's trace-compilation --- on subsequent runs (e.g. on new images), things will be faster.\n",
         "\n",
@@ -302,7 +335,7 @@
       },
       "source": [
         "image_dir = 'models/research/object_detection/test_images/'\n",
-        "image_path = os.path.join(image_dir, 'image1.jpg')\n",
+        "image_path = os.path.join(image_dir, 'image2.jpg')\n",
         "image_np = load_image_into_numpy_array(image_path)\n",
         "\n",
         "# Things to try:\n",
@@ -319,6 +352,13 @@
         "\n",
         "label_id_offset = 1\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 detections:\n",
+        "  keypoints = detections['detection_keypoints'][0].numpy()\n",
+        "  keypoint_scores = detections['detection_keypoint_scores'][0].numpy()\n",
+        "\n",
         "viz_utils.visualize_boxes_and_labels_on_image_array(\n",
         "      image_np_with_detections,\n",
         "      detections['detection_boxes'][0].numpy(),\n",
@@ -328,7 +368,10 @@
         "      use_normalized_coordinates=True,\n",
         "      max_boxes_to_draw=200,\n",
         "      min_score_thresh=.30,\n",
-        "      agnostic_mode=False)\n",
+        "      agnostic_mode=False,\n",
+        "      keypoints=keypoints,\n",
+        "      keypoint_scores=keypoint_scores,\n",
+        "      keypoint_edges=get_keypoint_tuples(configs['eval_config']))\n",
         "\n",
         "plt.figure(figsize=(12,16))\n",
         "plt.imshow(image_np_with_detections)\n",
@@ -409,7 +452,7 @@
         "      method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)[:,:,0]\n",
         "\n",
         "\n",
-        "class_name = 'dog'\n",
+        "class_name = 'kite'\n",
         "heatmap = get_heatmap(predictions_dict, class_name)\n",
         "resized_heatmap_unpadded = unpad_heatmap(heatmap, image_np)\n",
         "plt.figure(figsize=(12,16))\n",