def draw_heatmaps_on_image(image, heatmaps):
  """Draws heatmaps on an image.

  The heatmaps are handled channel by channel and different colors are used to
  paint different heatmap channels.

  Args:
    image: a PIL.Image object.
    heatmaps: a numpy array with shape [image_height, image_width, channel].
      Note that the image_height and image_width should match the size of input
      image.
  """
  draw = ImageDraw.Draw(image)
  channel = heatmaps.shape[2]
  for c in range(channel):
    heatmap = heatmaps[:, :, c] * 255
    heatmap = heatmap.astype('uint8')
    bitmap = Image.fromarray(heatmap, 'L')
    bitmap.convert('1')
    draw.bitmap(
        xy=[(0, 0)],
        bitmap=bitmap,
        fill=STANDARD_COLORS[c])
Draws heatmaps on an image.

  The heatmaps are handled channel by channel and different colors are used to
  paint different heatmap channels.

  Args:
    image: a PIL.Image object.
    heatmaps: a numpy array with shape [image_height, image_width, channel].
      Note that the image_height and image_width should match the size of input
      image.
  
inference_api/src/main/object_detection/utils/visualization_utils.py