def draw_heatmaps_on_image_array()

in object-detection/visualization_utils.py [0:0]

• 8 lines of code
• 3 McCabe index (conditional complexity)

def draw_heatmaps_on_image_array(image, heatmaps):
  """Overlays heatmaps to an image (numpy array).
  The function overlays the heatmaps on top of image. The heatmap values will be
  painted with different colors depending on the channels. Similar to
  "draw_heatmaps_on_image_array" function except the inputs are numpy arrays.
  Args:
    image: a numpy array with shape [height, width, 3].
    heatmaps: a numpy array with shape [height, width, channel].
  Returns:
    An uint8 numpy array representing the input image painted with heatmap
    colors.
  """
  if not isinstance(image, np.ndarray):
    image = image.numpy()
  if not isinstance(heatmaps, np.ndarray):
    heatmaps = heatmaps.numpy()
  image_pil = Image.fromarray(np.uint8(image)).convert('RGB')
  draw_heatmaps_on_image(image_pil, heatmaps)
  return np.array(image_pil)