colabs/pix2seq_inference_multitask.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "#@title License\n",
        "# Copyright 2022 The Pix2Seq Authors.\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",
        "# =============================================================================="
      ],
      "metadata": {
        "cellView": "form",
        "id": "0BkmPh5WJDHX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## A Unified Sequence Interface for Vision Tasks\n",
        "<a href=\"https://colab.research.google.com/github/google-research/pix2seq/blob/master/colabs/pix2seq_inference_multitask.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
        "\n",
        "\n",
        "This colab presents a demo for multi-task inference with Pix2seq. The table below provides a summary and model location for fine-tuned models on MSCOCO dataset.\n",
        "\n",
        "Backbone       | Total params (M) | Image size | COCO AP   | Google cloud storage location\n",
        "-------------: | ---------------: | ---------: | --------: | -----------:\n",
        "ViT-B          | 115.2            | 640x640    | 44.2      | [gs://pix2seq/multi_task/ckpt/vit_b_640x640](https://console.cloud.google.com/storage/browser/pix2seq/multi_task/ckpt/vit_b_640x640)\n",
        "ViT-B          | 115.2            | 1024x1024  | 46.5      | [gs://pix2seq/multi_task/ckpt/vit_b_1024x1024](https://console.cloud.google.com/storage/browser/pix2seq/multi_task/ckpt/vit_b_1024x1024)"
      ],
      "metadata": {
        "id": "iGuRf1kFH0-H"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kEdYtgV-ZTF5"
      },
      "outputs": [],
      "source": [
        "#@title Imports.\n",
        "import os\n",
        "import sys\n",
        "\n",
        "!pip install tensorflow\n",
        "!pip install ml_collections\n",
        "!pip install tensorflow-addons\n",
        "!pip install tensorflow-text\n",
        "!git clone https://github.com/google-research/pix2seq.git\n",
        "sys.path.append(os.getcwd())\n",
        "root_dir = os.getcwd()\n",
        "sys.path.insert(1, 'pix2seq')\n",
        "\n",
        "import tensorflow as tf\n",
        "from PIL import Image\n",
        "import numpy as np\n",
        "import google.colab\n",
        "import ml_collections\n",
        "import json\n",
        "import copy\n",
        "\n",
        "from models import ar_model as model_lib\n",
        "from tasks import instance_segmentation\n",
        "from tasks import keypoint_detection\n",
        "from tasks import object_detection\n",
        "from tasks import captioning\n",
        "from metrics import coco_metrics\n",
        "from tasks import task as task_lib\n",
        "from data import data_utils\n",
        "import utils\n",
        "from tasks.visualization import vis_utils\n",
        "from configs import config_multi_task"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "## Download coco annotations\n",
        "!mkdir /tmp/coco_annotations\n",
        "!wget https://storage.googleapis.com/pix2seq/multi_task/data/coco/json/captions_train2017_eval_compatible.json -P /tmp/coco_annotations/\n",
        "!wget https://storage.googleapis.com/pix2seq/multi_task/data/coco/json/captions_val2017_eval_compatible.json -P /tmp/coco_annotations/\n",
        "!wget https://storage.googleapis.com/pix2seq/multi_task/data/coco/json/instances_train2017.json -P /tmp/coco_annotations/\n",
        "!wget https://storage.googleapis.com/pix2seq/multi_task/data/coco/json/instances_val2017.json -P /tmp/coco_annotations/\n",
        "!wget https://storage.googleapis.com/pix2seq/multi_task/data/coco/json/person_keypoints_train2017.json -P /tmp/coco_annotations/\n",
        "!wget https://storage.googleapis.com/pix2seq/multi_task/data/coco/json/person_keypoints_val2017.json -P /tmp/coco_annotations/"
      ],
      "metadata": {
        "id": "MUY8cSCOVhb2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iRPSRPl49QV0"
      },
      "outputs": [],
      "source": [
        "#@title Load pix2seq multitask model.\n",
        "config = config_multi_task.get_config('object_detection@coco/2017_object_detection+instance_segmentation@coco/2017_instance_segmentation+keypoint_detection@coco/2017_keypoint_detection+captioning@coco/2017_captioning,vit-b')\n",
        "config.training = False\n",
        "\n",
        "# Restore checkpoint.\n",
        "model = model_lib.Model(config)\n",
        "checkpoint = tf.train.Checkpoint(\n",
        "    model=model, global_step=tf.Variable(0, dtype=tf.int64))\n",
        "model_dir = 'gs://pix2seq/multi_task/ckpt/vit_b_640x640'\n",
        "ckpt = tf.train.latest_checkpoint(model_dir)\n",
        "checkpoint.restore(ckpt).expect_partial()\n",
        "global_step = checkpoint.global_step"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lzK2JNQeAgBl"
      },
      "outputs": [],
      "source": [
        "#@title Get task and infer_fn.\n",
        "def get_task_config_and_infer_fn(config, task_name):\n",
        "  tconfig = copy.deepcopy(config)\n",
        "  for t in config.tasks:\n",
        "    if t['name'] == task_name:\n",
        "      tconfig.task = t\n",
        "  for d in config.datasets:\n",
        "    if d['name'] == task_name:\n",
        "      tconfig.dataset = d\n",
        "  tconfig.model_dir = ''\n",
        "\n",
        "  assert tconfig.task['name'] == task_name\n",
        "  task = task_lib.TaskRegistry.lookup(task_name)(tconfig)\n",
        "\n",
        "  if task_name == \"object_detection\":\n",
        "    tconfig.eval.batch_size = 1\n",
        "    tconfig.task.max_instances_per_image_test = 10\n",
        "  elif task_name == \"instance_segmentation\":\n",
        "    tconfig.task.use_gt_box_at_test = True\n",
        "    tconfig.task.ensemble_num_samples = 8\n",
        "    # For faster inference use smaller number of samples.\n",
        "    # tconfig.task.ensemble_num_samples = 1\n",
        "    tconfig.task.ensemble_threshold = 0.5\n",
        "    tconfig.eval.batch_size = 1\n",
        "    tconfig.task.max_instances_per_image_test = 1\n",
        "  elif task_name == \"keypoint_detection\":\n",
        "    tconfig.task.use_gt_box_at_test = True\n",
        "    tconfig.task.eval_suppress_invisible_token = True\n",
        "    tconfig.task.unbatch=True\n",
        "    tconfig.task.crop_to_bbox=True\n",
        "    tconfig.task.crop_to_bbox_pad_scale=0.5\n",
        "    tconfig.task.keypoint_score_weight=0.1\n",
        "    tconfig.task.points_score_weight = 0.1\n",
        "    tconfig.task.max_instances_per_image_test = 1\n",
        "    tconfig.eval.batch_size = 1\n",
        "    tconfig.task.top_p = 0.2\n",
        "  elif task_name == \"captioning\":\n",
        "    tconfig.task.captions_per_image = 1\n",
        "    tconfig.task.max_instances_per_image = 1\n",
        "\n",
        "  @tf.function\n",
        "  def infer(model, preprocessed_outputs):\n",
        "    return task.infer(model, preprocessed_outputs)\n",
        "\n",
        "  return tconfig, task, infer\n",
        "\n",
        "config_det, task_det, infer_det = get_task_config_and_infer_fn(config, \"object_detection\")\n",
        "config_seg, task_seg, infer_seg = get_task_config_and_infer_fn(config, \"instance_segmentation\")\n",
        "config_key, task_key, infer_key = get_task_config_and_infer_fn(config, \"keypoint_detection\")\n",
        "config_cap, task_cap, infer_cap = get_task_config_and_infer_fn(config, \"captioning\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4BCjazj20EWK"
      },
      "outputs": [],
      "source": [
        "#@title Functions to contruct inference examples and visualizations.\n",
        "def get_preprocessed_outputs_detection(image, image_id, config):\n",
        "  im = image\n",
        "  h, w = np.shape(im)[0:2]\n",
        "  features = {\n",
        "      'image': tf.image.convert_image_dtype(np.array(im), tf.float32),\n",
        "      'image/id': image_id,\n",
        "      'orig_image_size': tf.shape(im)[0:2],\n",
        "      'image_size_before_cropping': tf.shape(im)[0:2],\n",
        "  }\n",
        "  n_instance = config.task.max_instances_per_image_test\n",
        "  labels = {\n",
        "      'label': tf.zeros([n_instance], tf.int64),\n",
        "      'bbox': tf.zeros([n_instance, 4], tf.float32),\n",
        "      'is_crowd': tf.zeros([n_instance], tf.bool),\n",
        "      'area': tf.zeros([n_instance]),\n",
        "  }\n",
        "\n",
        "  features, labels = data_utils.preprocess_eval(\n",
        "      features,\n",
        "      labels,\n",
        "      max_image_size=config.model.image_size,\n",
        "      max_instances_per_image=config.task.max_instances_per_image_test,\n",
        "      object_coordinate_keys=('bbox',))\n",
        "\n",
        "  # Batch features and labels.\n",
        "  features = {\n",
        "      k: tf.expand_dims(v, 0) for k, v in features.items()\n",
        "  }\n",
        "  labels = {\n",
        "      k: tf.expand_dims(v, 0) for k, v in labels.items()\n",
        "  }\n",
        "\n",
        "  preprocessed_outputs = (features['image'], None, (features, labels))\n",
        "  return preprocessed_outputs\n",
        "\n",
        "\n",
        "def visualize_detection(results, image, categeory_names):\n",
        "  (images, image_ids, pred_bboxes, pred_bboxes_rescaled, pred_classes,\n",
        "   scores, gt_classes, gt_bboxes, gt_bboxes_rescaled, area, is_crowd) = results\n",
        "  n_instance = config.task.max_instances_per_image_test\n",
        "\n",
        "  pred_bboxes_rescaled = tf.reshape(pred_bboxes_rescaled, [-1, n_instance, 4])\n",
        "  pred_classes = tf.reshape(pred_classes, [-1, n_instance])\n",
        "  scores = tf.reshape(scores, [-1, n_instance])\n",
        "\n",
        "  bboxes_ = pred_bboxes_rescaled.numpy()\n",
        "  classes_ = pred_classes.numpy()\n",
        "  scores_ = scores.numpy()\n",
        "  images_ = np.copy(tf.image.convert_image_dtype(image, tf.uint8))\n",
        "\n",
        "  keep_indices = np.where(classes_[0] > 0)\n",
        "  vis = vis_utils.visualize_boxes_and_labels_on_image_array(\n",
        "      image=images_,\n",
        "      boxes=bboxes_[0][keep_indices],\n",
        "      classes=classes_[0][keep_indices],\n",
        "      scores=scores_[0][keep_indices],\n",
        "      category_index=categeory_names,\n",
        "      use_normalized_coordinates=False,\n",
        "      min_score_thresh=0.9,\n",
        "      skip_labels=False,\n",
        "      skip_scores=False,\n",
        "      line_thickness=4)\n",
        "\n",
        "  return vis\n",
        "\n",
        "\n",
        "def get_preprocessed_outputs_segmentation(image, image_id, bbox, config, label_id=1):\n",
        "  im = image\n",
        "  ymin, xmin, ymax, xmax = bbox\n",
        "  h, w = np.shape(im)[0:2]\n",
        "  features = {\n",
        "      'image': tf.image.convert_image_dtype(np.array(im), tf.float32),\n",
        "      'image/id': image_id,\n",
        "      'orig_image_size': tf.shape(im)[0:2],\n",
        "      'image_size_before_cropping': tf.shape(im)[0:2],\n",
        "  }\n",
        "  labels = {\n",
        "      'label': tf.convert_to_tensor([label_id], tf.int64),\n",
        "      'bbox': tf.convert_to_tensor([[\n",
        "          ymin/h, xmin/w, ymax/h, xmax/w\n",
        "      ]], tf.float32),\n",
        "      'scores': tf.ones([1]),\n",
        "      'is_crowd': tf.zeros([1], tf.bool),\n",
        "      'area': tf.zeros([1]),\n",
        "  }\n",
        "\n",
        "  features, labels = data_utils.preprocess_eval(\n",
        "      features,\n",
        "      labels,\n",
        "      max_image_size=config.model.image_size,\n",
        "      max_instances_per_image=config.task.max_instances_per_image_test,\n",
        "      object_coordinate_keys=('bbox',))\n",
        "\n",
        "  # Batch features and labels.\n",
        "  features = {\n",
        "      k: tf.expand_dims(v, 0) for k, v in features.items()\n",
        "  }\n",
        "  labels = {\n",
        "      k: tf.expand_dims(v, 0) for k, v in labels.items()\n",
        "  }\n",
        "\n",
        "  preprocessed_outputs = (features['image'], None, (features, labels))\n",
        "  return preprocessed_outputs\n",
        "\n",
        "def visualize_segmentation(results, config, category_names):\n",
        "  (images, image_ids, orig_image_size, unpadded_image_size, n_instances,  # pylint: disable=unbalanced-tuple-unpacking\n",
        "    pred_points, pred_points_rescaled, pred_classes, pred_bboxes,\n",
        "    scores) = results\n",
        "  n_instances = n_instances[0]\n",
        "  num_samples = config.task.ensemble_num_samples\n",
        "  threshold = config.task.ensemble_threshold\n",
        "  bsz = tf.shape(images)[0]\n",
        "\n",
        "  # Log/accumulate metrics.\n",
        "  mask_size = orig_image_size\n",
        "  pred_masks_rle = instance_segmentation.segment_to_mask_rle(\n",
        "      pred_points_rescaled.numpy(), mask_size.numpy())\n",
        "  # Ensemble mask for metrics.\n",
        "  pred_masks_np = instance_segmentation.mask_rle_to_mask_np(\n",
        "      pred_masks_rle, mask_size.numpy())\n",
        "  pred_masks_rle = instance_segmentation.ensemble_mask_np(\n",
        "      pred_masks_np, (bsz * n_instances).numpy(), num_samples, threshold)\n",
        "\n",
        "  # Image summary.\n",
        "  image_size = images.shape[1:3].as_list()\n",
        "  pred_points_rescaled = utils.scale_points(pred_points, image_size)\n",
        "  pred_masks_rle = instance_segmentation.segment_to_mask_rle(\n",
        "      pred_points_rescaled.numpy(), image_size, is_single_shape=True)\n",
        "  pred_masks_np = instance_segmentation.mask_rle_to_mask_np(\n",
        "      pred_masks_rle, image_size, is_single_shape=True)\n",
        "  # Ensemble masks for image summary.\n",
        "  pred_masks_np = instance_segmentation.ensemble_mask_np(\n",
        "      pred_masks_np, (bsz * n_instances).numpy(), num_samples, threshold)\n",
        "  pred_masks_np = np.asarray(pred_masks_np, np.uint8)\n",
        "  mask_new_shape = [pred_masks_np.shape[0]//n_instances, n_instances] + list(\n",
        "      pred_masks_np.shape[1:])\n",
        "  pred_masks_np = np.reshape(pred_masks_np, mask_new_shape)\n",
        "  pred_bboxes = tf.reshape(pred_bboxes, [-1, n_instances, 4])\n",
        "  pred_classes = tf.reshape(pred_classes, [-1, n_instances])\n",
        "  scores = tf.reshape(scores, [-1, n_instances])\n",
        "\n",
        "  bboxes_ = pred_bboxes.numpy()\n",
        "  classes_ = pred_classes.numpy()\n",
        "  scores_ = scores.numpy()\n",
        "  images_ = np.copy(tf.image.convert_image_dtype(images, tf.uint8))\n",
        "\n",
        "  keep_indices = np.where(classes_[0] > 0)[0]\n",
        "  vis = vis_utils.visualize_boxes_and_labels_on_image_array(\n",
        "      image=images_[0],\n",
        "      boxes=bboxes_[0][keep_indices],\n",
        "      classes=classes_[0][keep_indices],\n",
        "      scores=scores_[0][keep_indices],\n",
        "      category_index=category_names,\n",
        "      instance_masks=pred_masks_np[0][keep_indices],\n",
        "      use_normalized_coordinates=True,\n",
        "      max_boxes_to_draw=20,\n",
        "      min_score_thresh=0.9,\n",
        "      skip_boxes=False,\n",
        "      skip_scores=True,\n",
        "      skip_labels=False)\n",
        "  \n",
        "  return vis\n",
        "\n",
        "def get_preprocessed_outputs_keypoint(image, image_id, bbox, config, label_id=1):\n",
        "  im = image\n",
        "  ymin, xmin, ymax, xmax = bbox\n",
        "  h, w = np.shape(im)[0:2]\n",
        "  features = {\n",
        "      'image': tf.image.convert_image_dtype(np.array(im), tf.float32),\n",
        "      'image/id': image_id,\n",
        "      'orig_image_size': tf.shape(im)[0:2],\n",
        "      'image_size_before_cropping': tf.shape(im)[0:2],\n",
        "  }\n",
        "  labels = {\n",
        "      'label': tf.convert_to_tensor([label_id], tf.int64),\n",
        "      'bbox': tf.convert_to_tensor([[\n",
        "          ymin/h, xmin/w, ymax/h, xmax/w\n",
        "      ]], tf.float32),\n",
        "      'scores': tf.ones([1]),\n",
        "      'is_crowd': tf.zeros([1], tf.bool),\n",
        "      'area': tf.zeros([1]),\n",
        "      'keypoints': tf.zeros([1, 14])\n",
        "  }\n",
        "\n",
        "  @tf.autograph.experimental.do_not_convert\n",
        "  def crop_to_bbox(features, labels, config):\n",
        "    tconfig = config.task\n",
        "    per_instance_points_keys=('polygon', 'keypoints')\n",
        "    bbox = labels['bbox'][0]\n",
        "    image_shape = tf.shape(features['image'])[:2]\n",
        "\n",
        "    # Normalized bbox coords.\n",
        "    ymin, xmin, ymax, xmax = [tf.squeeze(t) for t in tf.split(bbox, 4)]\n",
        "    image_h, image_w = [\n",
        "        tf.squeeze(t) for t in tf.split(utils.tf_float32(image_shape), 2)\n",
        "    ]\n",
        "    \n",
        "    def i32(t):\n",
        "      return tf.cast(t, tf.int32)\n",
        "    def f32(t):\n",
        "      return tf.cast(t, tf.float32)\n",
        "\n",
        "    floor = tf.math.floor\n",
        "    ceil = tf.math.ceil\n",
        "    ymin = i32(floor(ymin * (image_h - 1)))\n",
        "    xmin = i32(floor(xmin * (image_w - 1)))\n",
        "    ymax = i32(ceil(ymax * (image_h - 1)))\n",
        "    xmax = i32(ceil(xmax * (image_w - 1)))\n",
        "\n",
        "    crop_to_bbox_pad_scale = tconfig.crop_to_bbox_pad_scale\n",
        "    bbox_h = ymax - ymin + 1\n",
        "    bbox_w = xmax - xmin + 1\n",
        "    ypad = i32(f32(bbox_h) * crop_to_bbox_pad_scale)\n",
        "    xpad = i32(f32(bbox_w) * crop_to_bbox_pad_scale)\n",
        "    ymin = tf.math.maximum(0, ymin - ypad)\n",
        "    ymax = tf.math.minimum(image_shape[0] - 1, ymax + ypad)\n",
        "    xmin = tf.math.maximum(0, xmin - xpad)\n",
        "    xmax = tf.math.minimum(image_shape[1] - 1, xmax + xpad)\n",
        "\n",
        "    region = [ymin, xmin, ymax - ymin + 1, xmax - xmin + 1]\n",
        "    points_orig = {\n",
        "        key: labels[key] for key in per_instance_points_keys if key in labels\n",
        "    }\n",
        "    features, labels = data_utils.crop(features, labels, region)\n",
        "    for key in points_orig:\n",
        "      labels[key] = utils.preserve_reserved_tokens(\n",
        "          labels[key], points_orig[key])\n",
        "    return features, labels\n",
        "\n",
        "  # Crop.\n",
        "  features, labels = crop_to_bbox(features, labels, config_key)\n",
        "  features['orig_image_size'] = tf.shape(features['image'])[0:2]\n",
        "\n",
        "  features, labels = data_utils.preprocess_eval(\n",
        "      features, labels,\n",
        "      max_image_size=config.model.image_size,\n",
        "      max_instances_per_image=1,\n",
        "      object_coordinate_keys=('bbox',))\n",
        "\n",
        "  # Batch features and labels.\n",
        "  features = {\n",
        "      k: tf.expand_dims(v, 0) for k, v in features.items()\n",
        "  }\n",
        "  labels = {\n",
        "      k: tf.expand_dims(v, 0) for k, v in labels.items()\n",
        "  }\n",
        "\n",
        "  preprocessed_outputs = (features['image'], None, (features, labels))\n",
        "  return preprocessed_outputs\n",
        "\n",
        "def visualize_keypoint(results, image):\n",
        "  (images, image_ids, n_instances, pred_points, pred_points_rescaled,  # pylint: disable=unbalanced-tuple-unpacking\n",
        "   pred_classes, pred_bboxes, scores) = results\n",
        "  n_instances = tf.shape(images)[0]\n",
        "\n",
        "  pred_bboxes = tf.reshape(pred_bboxes, [-1, n_instances, 4])\n",
        "  pred_classes = tf.reshape(pred_classes, [-1, n_instances])\n",
        "  scores = tf.reshape(scores, [-1, n_instances])\n",
        "  points_shape = [-1, n_instances, pred_points.shape[-1] // 2, 2]\n",
        "  pred_points = tf.reshape(pred_points_rescaled, points_shape)\n",
        "\n",
        "  bboxes_ = pred_bboxes.numpy()\n",
        "  classes_ = pred_classes.numpy()\n",
        "  scores_ = scores.numpy()\n",
        "  keypoints_ = pred_points.numpy()\n",
        "  images_ = np.copy(tf.image.convert_image_dtype(image, tf.uint8))\n",
        "\n",
        "  keypoint_edges = [[16, 14], [14, 12], [17, 15], [15, 13], [12, 13], [6, 12],\n",
        "                    [7, 13], [6, 7], [6, 8], [7, 9], [8, 10], [9, 11], [2, 3],\n",
        "                    [1, 2], [1, 3], [2, 4], [3, 5], [4, 6], [5, 7]]\n",
        "  for i in range(len(keypoint_edges)):\n",
        "    for j in range(len(keypoint_edges[i])):\n",
        "      keypoint_edges[i][j] -= 1\n",
        "  person_class_id = 1\n",
        "\n",
        "  keep_indices = np.where(classes_[0] == person_class_id)[0]\n",
        "  vis = vis_utils.visualize_boxes_and_labels_on_image_array(\n",
        "      image=images_,\n",
        "      boxes=bboxes_[0][keep_indices],\n",
        "      classes=classes_[0][keep_indices],\n",
        "      scores=scores_[0][keep_indices],\n",
        "      category_index={},\n",
        "      keypoints=keypoints_[0][keep_indices],\n",
        "      keypoint_edges=keypoint_edges,\n",
        "      use_normalized_coordinates=False,\n",
        "      max_boxes_to_draw=1,\n",
        "      min_score_thresh=0.9,\n",
        "      skip_labels=True,\n",
        "      skip_scores=True,\n",
        "      line_thickness=4)\n",
        "\n",
        "  return vis\n",
        "\n",
        "def get_preprocessed_outputs_captioning(image, image_id, config):\n",
        "  im = image\n",
        "  ymin, xmin, ymax, xmax = bbox\n",
        "  h, w = np.shape(im)[0:2]\n",
        "\n",
        "  features = {\n",
        "      'image': tf.image.convert_image_dtype(np.array(im), tf.float32),\n",
        "      'image/id': image_id,\n",
        "      'orig_image_size': tf.shape(im)[0:2],\n",
        "  }\n",
        "  labels = {\n",
        "      'label': tf.zeros([1], tf.int64),\n",
        "      'bbox': tf.zeros([1, 4]),\n",
        "      'area': tf.zeros([1]),\n",
        "      'is_crowd': tf.zeros([1]),\n",
        "      'captions': tf.convert_to_tensor([\"dummy\"], tf.string)\n",
        "  }\n",
        "\n",
        "  features, labels = data_utils.preprocess_eval(\n",
        "      features,\n",
        "      labels,\n",
        "      max_image_size=config.model.image_size,\n",
        "      max_instances_per_image=1)\n",
        "    \n",
        "  # Batch features and labels.\n",
        "  features = {\n",
        "      k: tf.expand_dims(v, 0) for k, v in features.items()\n",
        "  }\n",
        "  labels = {\n",
        "      k: tf.expand_dims(v, 0) for k, v in labels.items()\n",
        "  }\n",
        "\n",
        "  preprocessed_outputs = (features['image'], None, (features, labels))\n",
        "  return preprocessed_outputs\n",
        "\n",
        "def print_captioning_result(outputs, tokenizer):\n",
        "  pred_seq = outputs[3]\n",
        "  pred_seq = tf.where(pred_seq == 0, 0, pred_seq - config_cap.model.text_vocab_shift)\n",
        "  print(tokenizer.detokenize([int(j) for j in pred_seq[0].numpy()]).numpy())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5W0tukYSBy1O"
      },
      "outputs": [],
      "source": [
        "# Download image.\n",
        "import requests\n",
        "def get_image(image_id, train=False):\n",
        "  image_id = \"{:0>6d}\".format(image_id)\n",
        "  split = 'train' if train else 'val'\n",
        "  url = f'http://images.cocodataset.org/{split}2017/000000{image_id}.jpg'\n",
        "  with tf.io.gfile.GFile(url) as f:\n",
        "    im = Image.open(requests.get(url, stream=True).raw)\n",
        "  return im"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VvwcagE9nMzD"
      },
      "outputs": [],
      "source": [
        "image_id = 230983\n",
        "im = get_image(image_id, train=False)\n",
        "im"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# Object detection.\n",
        "preprocessed_outputs = get_preprocessed_outputs_detection(\n",
        "    im, image_id, config_det)\n",
        "infer_outputs_det = infer_det(model, preprocessed_outputs)\n",
        "results_det = task_det.postprocess_tpu(*infer_outputs_det)\n",
        "\n",
        "vis_det = visualize_detection(results_det, np.asarray(im), task_det._category_names)\n",
        "Image.fromarray(vis_det)"
      ],
      "metadata": {
        "id": "DQBbSG4fmVcb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Get a bbox for person and use it for keypoint detection and instance segmentation.\n",
        "pred_bboxes_rescaled = results_det[3].numpy().reshape((-1, 4))\n",
        "pred_classes = results_det[4].numpy().reshape((-1))\n",
        "person_idx = np.where(pred_classes == 1)[0][0]\n",
        "bbox = list(pred_bboxes_rescaled[person_idx])\n",
        "# You can also specify a custom box in (ymin, xmin, ymax, xmax) format.\n",
        "# bbox = [148.5505782847104, 149.26714806698848, 325.032441173339, 294.49976727245297]"
      ],
      "metadata": {
        "id": "dZc6CmgeusrO"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eVd5ilgSC6_r"
      },
      "outputs": [],
      "source": [
        "# Instance segmentation.\n",
        "preprocessed_outputs = get_preprocessed_outputs_segmentation(\n",
        "    im, image_id, bbox, config_seg)\n",
        "infer_outputs_seg = infer_seg(model, preprocessed_outputs)\n",
        "results_seg = task_seg.postprocess_tpu(*infer_outputs_seg)\n",
        "\n",
        "vis_seg = visualize_segmentation(results_seg, config_seg, task_seg._category_names)\n",
        "Image.fromarray(vis_seg)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FwkhCLsB6V6T"
      },
      "outputs": [],
      "source": [
        "# Keypoint detection.\n",
        "preprocessed_outputs = get_preprocessed_outputs_keypoint(\n",
        "    im, image_id, bbox, config_key)\n",
        "infer_outputs_key = infer_key(model, preprocessed_outputs)\n",
        "results_key = task_key.postprocess_tpu(*infer_outputs_key)\n",
        "\n",
        "vis_key = visualize_keypoint(results_key, tf.image.convert_image_dtype(np.array(im), tf.float32))\n",
        "Image.fromarray(vis_key)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m39ybc2u68c_"
      },
      "outputs": [],
      "source": [
        "# Captioning.\n",
        "preprocessed_outputs = get_preprocessed_outputs_captioning(\n",
        "    im, image_id, config_cap)\n",
        "infer_outputs = infer_cap(model, preprocessed_outputs)\n",
        "outputs = task_cap.postprocess_tpu(*infer_outputs)\n",
        "print_captioning_result(outputs, task_cap._tokenizer.tokenizer)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "TPU",
    "gpuClass": "standard"
  },
  "nbformat": 4,
  "nbformat_minor": 0
}