LeapImageMask

code_loader.contract.visualizer_classes.LeapImageMask

Used to visualize a semantic segmentation mask

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

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

Args
image	np.ndarray uint8/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 uint8 representation of the mask. The expected format is [H,W] where each pixel value signifies its label
labels	A list of strings mapping pixel values to their labels

Examples

Basic Usage

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

@tensorleap_custom_visualizer(name='seg_visualizer',
                              visualizer_type=LeapDataType.ImageMask,
                              heatmap_function=None)
def segmentation_visualizer(image: np.ndarray, segmentation_prediction: np.ndarray) -> LeapImageMask:
    labels = ['background', 'vehicle', 'tree', 'road', 'pavement']
    return LeapImageMask(mask=segmentation_prediction.squeeze(0), image=image.squeeze(0), labels=labels)