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?