site/en/hub/tutorials/movenet.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "toCy3v03Dwx7"
      },
      "source": [
        "##### Copyright 2021 The TensorFlow Hub Authors.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QKe-ubNcDvgv"
      },
      "outputs": [],
      "source": [
        "# Copyright 2021 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": "KqtQzBCpIJ7Y"
      },
      "source": [
        "# MoveNet: Ultra fast and accurate pose detection model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MCmFOosnSkCd"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/hub/tutorials/movenet\"><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/movenet.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/movenet.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/movenet.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/s?q=movenet\"><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": "6x99e0aEY_d6"
      },
      "source": [
        "**[MoveNet](https://t.co/QpfnVL0YYI?amp=1)** is an ultra fast and accurate model that detects 17 keypoints of a body. The model is offered on [TF Hub](https://tfhub.dev/s?q=movenet) with two variants, known as Lightning and Thunder. Lightning is intended for latency-critical applications, while Thunder is intended for applications that require high accuracy. Both models run faster than real time (30+ FPS) on most modern desktops, laptops, and phones, which proves crucial for live fitness, health, and wellness applications.\n",
        "\n",
        "\n",
        "<img src=\"https://github.com/tensorflow/tfjs-models/raw/master/pose-detection/assets/combined_squat_dance.gif\" alt=\"drawing\"/>\n",
        "\n",
        "*Images downloaded from Pexels (https://www.pexels.com/)\n",
        "\n",
        "This Colab walks you through the details of how to load MoveNet, and run inference on the input image and video below.\n",
        "\n",
        "Note: check out the [live demo](https://storage.googleapis.com/tfjs-models/demos/pose-detection/index.html?model=movenet) for how the model works!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "10_zkgbZBkIE"
      },
      "source": [
        "# Human Pose Estimation with MoveNet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9u_VGR6_BmbZ"
      },
      "source": [
        "## Visualization libraries & Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TtcwSIcgbIVN"
      },
      "outputs": [],
      "source": [
        "!pip install -q imageio\n",
        "!pip install -q opencv-python\n",
        "!pip install -q git+https://github.com/tensorflow/docs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9BLeJv-pCCld"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow_hub as hub\n",
        "from tensorflow_docs.vis import embed\n",
        "import numpy as np\n",
        "import cv2\n",
        "\n",
        "# Import matplotlib libraries\n",
        "from matplotlib import pyplot as plt\n",
        "from matplotlib.collections import LineCollection\n",
        "import matplotlib.patches as patches\n",
        "\n",
        "# Some modules to display an animation using imageio.\n",
        "import imageio\n",
        "from IPython.display import HTML, display"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "bEJBMeRb3YUy"
      },
      "outputs": [],
      "source": [
        "#@title Helper functions for visualization\n",
        "\n",
        "# Dictionary that maps from joint names to keypoint indices.\n",
        "KEYPOINT_DICT = {\n",
        "    'nose': 0,\n",
        "    'left_eye': 1,\n",
        "    'right_eye': 2,\n",
        "    'left_ear': 3,\n",
        "    'right_ear': 4,\n",
        "    'left_shoulder': 5,\n",
        "    'right_shoulder': 6,\n",
        "    'left_elbow': 7,\n",
        "    'right_elbow': 8,\n",
        "    'left_wrist': 9,\n",
        "    'right_wrist': 10,\n",
        "    'left_hip': 11,\n",
        "    'right_hip': 12,\n",
        "    'left_knee': 13,\n",
        "    'right_knee': 14,\n",
        "    'left_ankle': 15,\n",
        "    'right_ankle': 16\n",
        "}\n",
        "\n",
        "# Maps bones to a matplotlib color name.\n",
        "KEYPOINT_EDGE_INDS_TO_COLOR = {\n",
        "    (0, 1): 'm',\n",
        "    (0, 2): 'c',\n",
        "    (1, 3): 'm',\n",
        "    (2, 4): 'c',\n",
        "    (0, 5): 'm',\n",
        "    (0, 6): 'c',\n",
        "    (5, 7): 'm',\n",
        "    (7, 9): 'm',\n",
        "    (6, 8): 'c',\n",
        "    (8, 10): 'c',\n",
        "    (5, 6): 'y',\n",
        "    (5, 11): 'm',\n",
        "    (6, 12): 'c',\n",
        "    (11, 12): 'y',\n",
        "    (11, 13): 'm',\n",
        "    (13, 15): 'm',\n",
        "    (12, 14): 'c',\n",
        "    (14, 16): 'c'\n",
        "}\n",
        "\n",
        "def _keypoints_and_edges_for_display(keypoints_with_scores,\n",
        "                                     height,\n",
        "                                     width,\n",
        "                                     keypoint_threshold=0.11):\n",
        "  \"\"\"Returns high confidence keypoints and edges for visualization.\n",
        "\n",
        "  Args:\n",
        "    keypoints_with_scores: A numpy array with shape [1, 1, 17, 3] representing\n",
        "      the keypoint coordinates and scores returned from the MoveNet model.\n",
        "    height: height of the image in pixels.\n",
        "    width: width of the image in pixels.\n",
        "    keypoint_threshold: minimum confidence score for a keypoint to be\n",
        "      visualized.\n",
        "\n",
        "  Returns:\n",
        "    A (keypoints_xy, edges_xy, edge_colors) containing:\n",
        "      * the coordinates of all keypoints of all detected entities;\n",
        "      * the coordinates of all skeleton edges of all detected entities;\n",
        "      * the colors in which the edges should be plotted.\n",
        "  \"\"\"\n",
        "  keypoints_all = []\n",
        "  keypoint_edges_all = []\n",
        "  edge_colors = []\n",
        "  num_instances, _, _, _ = keypoints_with_scores.shape\n",
        "  for idx in range(num_instances):\n",
        "    kpts_x = keypoints_with_scores[0, idx, :, 1]\n",
        "    kpts_y = keypoints_with_scores[0, idx, :, 0]\n",
        "    kpts_scores = keypoints_with_scores[0, idx, :, 2]\n",
        "    kpts_absolute_xy = np.stack(\n",
        "        [width * np.array(kpts_x), height * np.array(kpts_y)], axis=-1)\n",
        "    kpts_above_thresh_absolute = kpts_absolute_xy[\n",
        "        kpts_scores > keypoint_threshold, :]\n",
        "    keypoints_all.append(kpts_above_thresh_absolute)\n",
        "\n",
        "    for edge_pair, color in KEYPOINT_EDGE_INDS_TO_COLOR.items():\n",
        "      if (kpts_scores[edge_pair[0]] > keypoint_threshold and\n",
        "          kpts_scores[edge_pair[1]] > keypoint_threshold):\n",
        "        x_start = kpts_absolute_xy[edge_pair[0], 0]\n",
        "        y_start = kpts_absolute_xy[edge_pair[0], 1]\n",
        "        x_end = kpts_absolute_xy[edge_pair[1], 0]\n",
        "        y_end = kpts_absolute_xy[edge_pair[1], 1]\n",
        "        line_seg = np.array([[x_start, y_start], [x_end, y_end]])\n",
        "        keypoint_edges_all.append(line_seg)\n",
        "        edge_colors.append(color)\n",
        "  if keypoints_all:\n",
        "    keypoints_xy = np.concatenate(keypoints_all, axis=0)\n",
        "  else:\n",
        "    keypoints_xy = np.zeros((0, 17, 2))\n",
        "\n",
        "  if keypoint_edges_all:\n",
        "    edges_xy = np.stack(keypoint_edges_all, axis=0)\n",
        "  else:\n",
        "    edges_xy = np.zeros((0, 2, 2))\n",
        "  return keypoints_xy, edges_xy, edge_colors\n",
        "\n",
        "\n",
        "def draw_prediction_on_image(\n",
        "    image, keypoints_with_scores, crop_region=None, close_figure=False,\n",
        "    output_image_height=None):\n",
        "  \"\"\"Draws the keypoint predictions on image.\n",
        "\n",
        "  Args:\n",
        "    image: A numpy array with shape [height, width, channel] representing the\n",
        "      pixel values of the input image.\n",
        "    keypoints_with_scores: A numpy array with shape [1, 1, 17, 3] representing\n",
        "      the keypoint coordinates and scores returned from the MoveNet model.\n",
        "    crop_region: A dictionary that defines the coordinates of the bounding box\n",
        "      of the crop region in normalized coordinates (see the init_crop_region\n",
        "      function below for more detail). If provided, this function will also\n",
        "      draw the bounding box on the image.\n",
        "    output_image_height: An integer indicating the height of the output image.\n",
        "      Note that the image aspect ratio will be the same as the input image.\n",
        "\n",
        "  Returns:\n",
        "    A numpy array with shape [out_height, out_width, channel] representing the\n",
        "    image overlaid with keypoint predictions.\n",
        "  \"\"\"\n",
        "  height, width, channel = image.shape\n",
        "  aspect_ratio = float(width) / height\n",
        "  fig, ax = plt.subplots(figsize=(12 * aspect_ratio, 12))\n",
        "  # To remove the huge white borders\n",
        "  fig.tight_layout(pad=0)\n",
        "  ax.margins(0)\n",
        "  ax.set_yticklabels([])\n",
        "  ax.set_xticklabels([])\n",
        "  plt.axis('off')\n",
        "\n",
        "  im = ax.imshow(image)\n",
        "  line_segments = LineCollection([], linewidths=(4), linestyle='solid')\n",
        "  ax.add_collection(line_segments)\n",
        "  # Turn off tick labels\n",
        "  scat = ax.scatter([], [], s=60, color='#FF1493', zorder=3)\n",
        "\n",
        "  (keypoint_locs, keypoint_edges,\n",
        "   edge_colors) = _keypoints_and_edges_for_display(\n",
        "       keypoints_with_scores, height, width)\n",
        "\n",
        "  line_segments.set_segments(keypoint_edges)\n",
        "  line_segments.set_color(edge_colors)\n",
        "  if keypoint_edges.shape[0]:\n",
        "    line_segments.set_segments(keypoint_edges)\n",
        "    line_segments.set_color(edge_colors)\n",
        "  if keypoint_locs.shape[0]:\n",
        "    scat.set_offsets(keypoint_locs)\n",
        "\n",
        "  if crop_region is not None:\n",
        "    xmin = max(crop_region['x_min'] * width, 0.0)\n",
        "    ymin = max(crop_region['y_min'] * height, 0.0)\n",
        "    rec_width = min(crop_region['x_max'], 0.99) * width - xmin\n",
        "    rec_height = min(crop_region['y_max'], 0.99) * height - ymin\n",
        "    rect = patches.Rectangle(\n",
        "        (xmin,ymin),rec_width,rec_height,\n",
        "        linewidth=1,edgecolor='b',facecolor='none')\n",
        "    ax.add_patch(rect)\n",
        "\n",
        "  fig.canvas.draw()\n",
        "  image_from_plot = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)\n",
        "  image_from_plot = image_from_plot.reshape(\n",
        "      fig.canvas.get_width_height()[::-1] + (3,))\n",
        "  plt.close(fig)\n",
        "  if output_image_height is not None:\n",
        "    output_image_width = int(output_image_height / height * width)\n",
        "    image_from_plot = cv2.resize(\n",
        "        image_from_plot, dsize=(output_image_width, output_image_height),\n",
        "         interpolation=cv2.INTER_CUBIC)\n",
        "  return image_from_plot\n",
        "\n",
        "def to_gif(images, duration):\n",
        "  \"\"\"Converts image sequence (4D numpy array) to gif.\"\"\"\n",
        "  imageio.mimsave('./animation.gif', images, duration=duration)\n",
        "  return embed.embed_file('./animation.gif')\n",
        "\n",
        "def progress(value, max=100):\n",
        "  return HTML(\"\"\"\n",
        "      <progress\n",
        "          value='{value}'\n",
        "          max='{max}',\n",
        "          style='width: 100%'\n",
        "      >\n",
        "          {value}\n",
        "      </progress>\n",
        "  \"\"\".format(value=value, max=max))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UvrN0iQiOxhR"
      },
      "source": [
        "## Load Model from TF hub"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zeGHgANcT7a1"
      },
      "outputs": [],
      "source": [
        "model_name = \"movenet_lightning\" #@param [\"movenet_lightning\", \"movenet_thunder\", \"movenet_lightning_f16.tflite\", \"movenet_thunder_f16.tflite\", \"movenet_lightning_int8.tflite\", \"movenet_thunder_int8.tflite\"]\n",
        "\n",
        "if \"tflite\" in model_name:\n",
        "  if \"movenet_lightning_f16\" in model_name:\n",
        "    !wget -q -O model.tflite https://tfhub.dev/google/lite-model/movenet/singlepose/lightning/tflite/float16/4?lite-format=tflite\n",
        "    input_size = 192\n",
        "  elif \"movenet_thunder_f16\" in model_name:\n",
        "    !wget -q -O model.tflite https://tfhub.dev/google/lite-model/movenet/singlepose/thunder/tflite/float16/4?lite-format=tflite\n",
        "    input_size = 256\n",
        "  elif \"movenet_lightning_int8\" in model_name:\n",
        "    !wget -q -O model.tflite https://tfhub.dev/google/lite-model/movenet/singlepose/lightning/tflite/int8/4?lite-format=tflite\n",
        "    input_size = 192\n",
        "  elif \"movenet_thunder_int8\" in model_name:\n",
        "    !wget -q -O model.tflite https://tfhub.dev/google/lite-model/movenet/singlepose/thunder/tflite/int8/4?lite-format=tflite\n",
        "    input_size = 256\n",
        "  else:\n",
        "    raise ValueError(\"Unsupported model name: %s\" % model_name)\n",
        "\n",
        "  # Initialize the TFLite interpreter\n",
        "  interpreter = tf.lite.Interpreter(model_path=\"model.tflite\")\n",
        "  interpreter.allocate_tensors()\n",
        "\n",
        "  def movenet(input_image):\n",
        "    \"\"\"Runs detection on an input image.\n",
        "\n",
        "    Args:\n",
        "      input_image: A [1, height, width, 3] tensor represents the input image\n",
        "        pixels. Note that the height/width should already be resized and match the\n",
        "        expected input resolution of the model before passing into this function.\n",
        "\n",
        "    Returns:\n",
        "      A [1, 1, 17, 3] float numpy array representing the predicted keypoint\n",
        "      coordinates and scores.\n",
        "    \"\"\"\n",
        "    # TF Lite format expects tensor type of uint8.\n",
        "    input_image = tf.cast(input_image, dtype=tf.uint8)\n",
        "    input_details = interpreter.get_input_details()\n",
        "    output_details = interpreter.get_output_details()\n",
        "    interpreter.set_tensor(input_details[0]['index'], input_image.numpy())\n",
        "    # Invoke inference.\n",
        "    interpreter.invoke()\n",
        "    # Get the model prediction.\n",
        "    keypoints_with_scores = interpreter.get_tensor(output_details[0]['index'])\n",
        "    return keypoints_with_scores\n",
        "\n",
        "else:\n",
        "  if \"movenet_lightning\" in model_name:\n",
        "    module = hub.load(\"https://tfhub.dev/google/movenet/singlepose/lightning/4\")\n",
        "    input_size = 192\n",
        "  elif \"movenet_thunder\" in model_name:\n",
        "    module = hub.load(\"https://tfhub.dev/google/movenet/singlepose/thunder/4\")\n",
        "    input_size = 256\n",
        "  else:\n",
        "    raise ValueError(\"Unsupported model name: %s\" % model_name)\n",
        "\n",
        "  def movenet(input_image):\n",
        "    \"\"\"Runs detection on an input image.\n",
        "\n",
        "    Args:\n",
        "      input_image: A [1, height, width, 3] tensor represents the input image\n",
        "        pixels. Note that the height/width should already be resized and match the\n",
        "        expected input resolution of the model before passing into this function.\n",
        "\n",
        "    Returns:\n",
        "      A [1, 1, 17, 3] float numpy array representing the predicted keypoint\n",
        "      coordinates and scores.\n",
        "    \"\"\"\n",
        "    model = module.signatures['serving_default']\n",
        "\n",
        "    # SavedModel format expects tensor type of int32.\n",
        "    input_image = tf.cast(input_image, dtype=tf.int32)\n",
        "    # Run model inference.\n",
        "    outputs = model(input_image)\n",
        "    # Output is a [1, 1, 17, 3] tensor.\n",
        "    keypoints_with_scores = outputs['output_0'].numpy()\n",
        "    return keypoints_with_scores"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-h1qHYaqD9ap"
      },
      "source": [
        "## Single Image Example"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ymTVR2I9x22I"
      },
      "source": [
        "This session demonstrates the minimum working example of running the model on a **single image** to predict the 17 human keypoints."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5I3xBq80E3N_"
      },
      "source": [
        "### Load Input Image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GMO4B-wx5psP"
      },
      "outputs": [],
      "source": [
        "!curl -o input_image.jpeg https://images.pexels.com/photos/4384679/pexels-photo-4384679.jpeg --silent"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lJZYQ8KYFQ6x"
      },
      "outputs": [],
      "source": [
        "# Load the input image.\n",
        "image_path = 'input_image.jpeg'\n",
        "image = tf.io.read_file(image_path)\n",
        "image = tf.image.decode_jpeg(image)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S_UWRdQxE6WN"
      },
      "source": [
        "### Run Inference"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VHmTwACwFW-v"
      },
      "outputs": [],
      "source": [
        "# Resize and pad the image to keep the aspect ratio and fit the expected size.\n",
        "input_image = tf.expand_dims(image, axis=0)\n",
        "input_image = tf.image.resize_with_pad(input_image, input_size, input_size)\n",
        "\n",
        "# Run model inference.\n",
        "keypoints_with_scores = movenet(input_image)\n",
        "\n",
        "# Visualize the predictions with image.\n",
        "display_image = tf.expand_dims(image, axis=0)\n",
        "display_image = tf.cast(tf.image.resize_with_pad(\n",
        "    display_image, 1280, 1280), dtype=tf.int32)\n",
        "output_overlay = draw_prediction_on_image(\n",
        "    np.squeeze(display_image.numpy(), axis=0), keypoints_with_scores)\n",
        "\n",
        "plt.figure(figsize=(5, 5))\n",
        "plt.imshow(output_overlay)\n",
        "_ = plt.axis('off')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rKm-B0eMYeg8"
      },
      "source": [
        "## Video (Image Sequence) Example"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gdPFXabLyiKv"
      },
      "source": [
        "This section demonstrates how to apply intelligent cropping based on detections from the previous frame when the input is a sequence of frames. This allows the model to devote its attention and resources to the main subject, resulting in much better prediction quality without sacrificing the speed.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "SYFdK-JHYhrv"
      },
      "outputs": [],
      "source": [
        "#@title Cropping Algorithm\n",
        "\n",
        "# Confidence score to determine whether a keypoint prediction is reliable.\n",
        "MIN_CROP_KEYPOINT_SCORE = 0.2\n",
        "\n",
        "def init_crop_region(image_height, image_width):\n",
        "  \"\"\"Defines the default crop region.\n",
        "\n",
        "  The function provides the initial crop region (pads the full image from both\n",
        "  sides to make it a square image) when the algorithm cannot reliably determine\n",
        "  the crop region from the previous frame.\n",
        "  \"\"\"\n",
        "  if image_width > image_height:\n",
        "    box_height = image_width / image_height\n",
        "    box_width = 1.0\n",
        "    y_min = (image_height / 2 - image_width / 2) / image_height\n",
        "    x_min = 0.0\n",
        "  else:\n",
        "    box_height = 1.0\n",
        "    box_width = image_height / image_width\n",
        "    y_min = 0.0\n",
        "    x_min = (image_width / 2 - image_height / 2) / image_width\n",
        "\n",
        "  return {\n",
        "    'y_min': y_min,\n",
        "    'x_min': x_min,\n",
        "    'y_max': y_min + box_height,\n",
        "    'x_max': x_min + box_width,\n",
        "    'height': box_height,\n",
        "    'width': box_width\n",
        "  }\n",
        "\n",
        "def torso_visible(keypoints):\n",
        "  \"\"\"Checks whether there are enough torso keypoints.\n",
        "\n",
        "  This function checks whether the model is confident at predicting one of the\n",
        "  shoulders/hips which is required to determine a good crop region.\n",
        "  \"\"\"\n",
        "  return ((keypoints[0, 0, KEYPOINT_DICT['left_hip'], 2] >\n",
        "           MIN_CROP_KEYPOINT_SCORE or\n",
        "          keypoints[0, 0, KEYPOINT_DICT['right_hip'], 2] >\n",
        "           MIN_CROP_KEYPOINT_SCORE) and\n",
        "          (keypoints[0, 0, KEYPOINT_DICT['left_shoulder'], 2] >\n",
        "           MIN_CROP_KEYPOINT_SCORE or\n",
        "          keypoints[0, 0, KEYPOINT_DICT['right_shoulder'], 2] >\n",
        "           MIN_CROP_KEYPOINT_SCORE))\n",
        "\n",
        "def determine_torso_and_body_range(\n",
        "    keypoints, target_keypoints, center_y, center_x):\n",
        "  \"\"\"Calculates the maximum distance from each keypoints to the center location.\n",
        "\n",
        "  The function returns the maximum distances from the two sets of keypoints:\n",
        "  full 17 keypoints and 4 torso keypoints. The returned information will be\n",
        "  used to determine the crop size. See determineCropRegion for more detail.\n",
        "  \"\"\"\n",
        "  torso_joints = ['left_shoulder', 'right_shoulder', 'left_hip', 'right_hip']\n",
        "  max_torso_yrange = 0.0\n",
        "  max_torso_xrange = 0.0\n",
        "  for joint in torso_joints:\n",
        "    dist_y = abs(center_y - target_keypoints[joint][0])\n",
        "    dist_x = abs(center_x - target_keypoints[joint][1])\n",
        "    if dist_y > max_torso_yrange:\n",
        "      max_torso_yrange = dist_y\n",
        "    if dist_x > max_torso_xrange:\n",
        "      max_torso_xrange = dist_x\n",
        "\n",
        "  max_body_yrange = 0.0\n",
        "  max_body_xrange = 0.0\n",
        "  for joint in KEYPOINT_DICT.keys():\n",
        "    if keypoints[0, 0, KEYPOINT_DICT[joint], 2] < MIN_CROP_KEYPOINT_SCORE:\n",
        "      continue\n",
        "    dist_y = abs(center_y - target_keypoints[joint][0]);\n",
        "    dist_x = abs(center_x - target_keypoints[joint][1]);\n",
        "    if dist_y > max_body_yrange:\n",
        "      max_body_yrange = dist_y\n",
        "\n",
        "    if dist_x > max_body_xrange:\n",
        "      max_body_xrange = dist_x\n",
        "\n",
        "  return [max_torso_yrange, max_torso_xrange, max_body_yrange, max_body_xrange]\n",
        "\n",
        "def determine_crop_region(\n",
        "      keypoints, image_height,\n",
        "      image_width):\n",
        "  \"\"\"Determines the region to crop the image for the model to run inference on.\n",
        "\n",
        "  The algorithm uses the detected joints from the previous frame to estimate\n",
        "  the square region that encloses the full body of the target person and\n",
        "  centers at the midpoint of two hip joints. The crop size is determined by\n",
        "  the distances between each joints and the center point.\n",
        "  When the model is not confident with the four torso joint predictions, the\n",
        "  function returns a default crop which is the full image padded to square.\n",
        "  \"\"\"\n",
        "  target_keypoints = {}\n",
        "  for joint in KEYPOINT_DICT.keys():\n",
        "    target_keypoints[joint] = [\n",
        "      keypoints[0, 0, KEYPOINT_DICT[joint], 0] * image_height,\n",
        "      keypoints[0, 0, KEYPOINT_DICT[joint], 1] * image_width\n",
        "    ]\n",
        "\n",
        "  if torso_visible(keypoints):\n",
        "    center_y = (target_keypoints['left_hip'][0] +\n",
        "                target_keypoints['right_hip'][0]) / 2;\n",
        "    center_x = (target_keypoints['left_hip'][1] +\n",
        "                target_keypoints['right_hip'][1]) / 2;\n",
        "\n",
        "    (max_torso_yrange, max_torso_xrange,\n",
        "      max_body_yrange, max_body_xrange) = determine_torso_and_body_range(\n",
        "          keypoints, target_keypoints, center_y, center_x)\n",
        "\n",
        "    crop_length_half = np.amax(\n",
        "        [max_torso_xrange * 1.9, max_torso_yrange * 1.9,\n",
        "          max_body_yrange * 1.2, max_body_xrange * 1.2])\n",
        "\n",
        "    tmp = np.array(\n",
        "        [center_x, image_width - center_x, center_y, image_height - center_y])\n",
        "    crop_length_half = np.amin(\n",
        "        [crop_length_half, np.amax(tmp)]);\n",
        "\n",
        "    crop_corner = [center_y - crop_length_half, center_x - crop_length_half];\n",
        "\n",
        "    if crop_length_half > max(image_width, image_height) / 2:\n",
        "      return init_crop_region(image_height, image_width)\n",
        "    else:\n",
        "      crop_length = crop_length_half * 2;\n",
        "      return {\n",
        "        'y_min': crop_corner[0] / image_height,\n",
        "        'x_min': crop_corner[1] / image_width,\n",
        "        'y_max': (crop_corner[0] + crop_length) / image_height,\n",
        "        'x_max': (crop_corner[1] + crop_length) / image_width,\n",
        "        'height': (crop_corner[0] + crop_length) / image_height -\n",
        "            crop_corner[0] / image_height,\n",
        "        'width': (crop_corner[1] + crop_length) / image_width -\n",
        "            crop_corner[1] / image_width\n",
        "      }\n",
        "  else:\n",
        "    return init_crop_region(image_height, image_width)\n",
        "\n",
        "def crop_and_resize(image, crop_region, crop_size):\n",
        "  \"\"\"Crops and resize the image to prepare for the model input.\"\"\"\n",
        "  boxes=[[crop_region['y_min'], crop_region['x_min'],\n",
        "          crop_region['y_max'], crop_region['x_max']]]\n",
        "  output_image = tf.image.crop_and_resize(\n",
        "      image, box_indices=[0], boxes=boxes, crop_size=crop_size)\n",
        "  return output_image\n",
        "\n",
        "def run_inference(movenet, image, crop_region, crop_size):\n",
        "  \"\"\"Runs model inference on the cropped region.\n",
        "\n",
        "  The function runs the model inference on the cropped region and updates the\n",
        "  model output to the original image coordinate system.\n",
        "  \"\"\"\n",
        "  image_height, image_width, _ = image.shape\n",
        "  input_image = crop_and_resize(\n",
        "    tf.expand_dims(image, axis=0), crop_region, crop_size=crop_size)\n",
        "  # Run model inference.\n",
        "  keypoints_with_scores = movenet(input_image)\n",
        "  # Update the coordinates.\n",
        "  for idx in range(17):\n",
        "    keypoints_with_scores[0, 0, idx, 0] = (\n",
        "        crop_region['y_min'] * image_height +\n",
        "        crop_region['height'] * image_height *\n",
        "        keypoints_with_scores[0, 0, idx, 0]) / image_height\n",
        "    keypoints_with_scores[0, 0, idx, 1] = (\n",
        "        crop_region['x_min'] * image_width +\n",
        "        crop_region['width'] * image_width *\n",
        "        keypoints_with_scores[0, 0, idx, 1]) / image_width\n",
        "  return keypoints_with_scores"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L2JmA1xAEntQ"
      },
      "source": [
        "### Load Input Image Sequence"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CzJxbxDckWl2"
      },
      "outputs": [],
      "source": [
        "!wget -q -O dance.gif https://github.com/tensorflow/tfjs-models/raw/master/pose-detection/assets/dance_input.gif"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IxbMFZJUkd6W"
      },
      "outputs": [],
      "source": [
        "# Load the input image.\n",
        "image_path = 'dance.gif'\n",
        "image = tf.io.read_file(image_path)\n",
        "image = tf.image.decode_gif(image)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CJKeQ4siEtU9"
      },
      "source": [
        "### Run Inference with Cropping Algorithm"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9B57XS0NZPIy"
      },
      "outputs": [],
      "source": [
        "# Load the input image.\n",
        "num_frames, image_height, image_width, _ = image.shape\n",
        "crop_region = init_crop_region(image_height, image_width)\n",
        "\n",
        "output_images = []\n",
        "bar = display(progress(0, num_frames-1), display_id=True)\n",
        "for frame_idx in range(num_frames):\n",
        "  keypoints_with_scores = run_inference(\n",
        "      movenet, image[frame_idx, :, :, :], crop_region,\n",
        "      crop_size=[input_size, input_size])\n",
        "  output_images.append(draw_prediction_on_image(\n",
        "      image[frame_idx, :, :, :].numpy().astype(np.int32),\n",
        "      keypoints_with_scores, crop_region=None,\n",
        "      close_figure=True, output_image_height=300))\n",
        "  crop_region = determine_crop_region(\n",
        "      keypoints_with_scores, image_height, image_width)\n",
        "  bar.update(progress(frame_idx, num_frames-1))\n",
        "\n",
        "# Prepare gif visualization.\n",
        "output = np.stack(output_images, axis=0)\n",
        "to_gif(output, duration=100)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "9u_VGR6_BmbZ",
        "5I3xBq80E3N_",
        "L2JmA1xAEntQ"
      ],
      "name": "movenet.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}