def test_draw_bounding_boxes_on_image_tensors_with_track_ids(self):
    """Tests that bounding box utility produces reasonable results."""
    category_index = {1: {'id': 1, 'name': 'dog'}, 2: {'id': 2, 'name': 'cat'}}

    fname = os.path.join(_TESTDATA_PATH, 'image1.jpg')
    image_np = np.array(Image.open(fname))
    images_np = np.stack((image_np, image_np), axis=0)
    original_image_shape = [[636, 512], [636, 512]]

    def graph_fn():
      images_tensor = tf.constant(value=images_np, dtype=tf.uint8)
      image_shape = tf.constant(original_image_shape, dtype=tf.int32)
      boxes = tf.constant([[[0.4, 0.25, 0.75, 0.75],
                            [0.5, 0.3, 0.7, 0.9],
                            [0.7, 0.5, 0.8, 0.9]],
                           [[0.41, 0.25, 0.75, 0.75],
                            [0.51, 0.3, 0.7, 0.9],
                            [0.75, 0.5, 0.8, 0.9]]])
      classes = tf.constant([[1, 1, 2], [1, 1, 2]], dtype=tf.int64)
      scores = tf.constant([[0.8, 0.5, 0.7], [0.6, 0.5, 0.8]])
      track_ids = tf.constant([[3, 9, 7], [3, 9, 144]], dtype=tf.int32)
      images_with_boxes = (
          visualization_utils.draw_bounding_boxes_on_image_tensors(
              images_tensor,
              boxes,
              classes,
              scores,
              category_index,
              original_image_spatial_shape=image_shape,
              true_image_shape=image_shape,
              track_ids=track_ids,
              min_score_thresh=0.2))
      return images_with_boxes

    # Write output images for visualization.
    images_with_boxes_np = self.execute(graph_fn, [])
    self.assertEqual(images_np.shape[0], images_with_boxes_np.shape[0])
    self.assertEqual(images_np.shape[3], images_with_boxes_np.shape[3])
    self.assertEqual(
        tuple(original_image_shape[0]), images_with_boxes_np.shape[1:3])
    for i in range(images_with_boxes_np.shape[0]):
      img_name = 'image_with_track_ids_' + str(i) + '.png'
      output_file = os.path.join(self.get_temp_dir(), img_name)
      logging.info('Writing output image %d to %s', i, output_file)
      image_pil = Image.fromarray(images_with_boxes_np[i, ...])
      image_pil.save(output_file)
Tests that bounding box utility produces reasonable results.
inference_api/src/main/object_detection/utils/visualization_utils_test.py