def test_draw_bounding_boxes_on_image_tensors_with_additional_channels(self):
    """Tests the case where input image tensor has more than 3 channels."""
    category_index = {1: {'id': 1, 'name': 'dog'}}
    image_np = self.create_test_image_with_five_channels()
    images_np = np.stack((image_np, image_np), axis=0)

    def graph_fn():
      images_tensor = tf.constant(value=images_np, dtype=tf.uint8)
      boxes = tf.constant(0, dtype=tf.float32, shape=[2, 0, 4])
      classes = tf.constant(0, dtype=tf.int64, shape=[2, 0])
      scores = tf.constant(0, dtype=tf.float32, shape=[2, 0])
      images_with_boxes = (
          visualization_utils.draw_bounding_boxes_on_image_tensors(
              images_tensor,
              boxes,
              classes,
              scores,
              category_index,
              min_score_thresh=0.2))

      return images_with_boxes

    final_images_np = self.execute(graph_fn, [])
    self.assertEqual((2, 100, 200, 3), final_images_np.shape)
Tests the case where input image tensor has more than 3 channels.
inference_api/src/main/object_detection/utils/visualization_utils_test.py