LeapImageWithHeatmap

code_loader.contract.visualizer_classes.LeapImageWithHeatmap

Used to visualize attention maps or a similar float-based mask over an image.

import numpy.typing as npt
from code_loader.contract.enums import LeapDataType

@dataclass
class LeapImageWithHeatmap:
    mask: npt.NDArray[np.float32]
    image: npt.NDArray[np.float32]
    labels: List[str]
    type: LeapDataType = LeapDataType.ImageMask
Args

image

np.ndarray float32 representation of the image. The expected image format is [H,W,1] OR [H,W,3] and is expected to be in [0,255].

mask

np.ndarray float32 representation of the mask. The expected shape is [C,H,W], where C is the amount of heatmaps that are provided for the image.

labels

A list of labels of C length for the heatmap labels

Examples

Basic Usage

import numpy as np
from code_loader.contract.visualizer_classes import LeapImageWithHeatmap
from code_loader import leap_binder
from code_loader.contract.enums import LeapDataType
...

@tensorleap_custom_visualizer("heatmap", LeapDataType.ImageWithHeatmap)
def heatmap(image_data: np.ndarray, heatmap_data: np.ndarray):
    labels = ["heatmap1"]
    return LeapImageWithHeatmap(image=image_data.squeeze(0), heatmaps=heatmap_data), labels=labels)

Last updated

Was this helpful?