Displaying image and mask from logits

[sumitsoman]

If we use palette like

palette= {
(0, 0, 0) : 0 , # background (black)
(255, 255, 255) : 1  # line (white)
}

and read in the segmentation images and masks

dataset = DataLoader(image_paths=image_paths,
                 mask_paths=mask_paths,
                 image_size=(IMG_SIZE, IMG_SIZE),
                 crop_percent=0.1,
                 channels=(3, 1),
                 augment=False,
                 compose=False,
                 one_hot_encoding=True,
                 palette = palette,
                 seed=SEED)

How can we access the images and masks to plot them? If we use

for images, masks in dataset.take(1):
     sample_image, sample_mask = images[0], masks[0]

It gives the error

InvalidArgumentError: ValueError: Tensor conversion requested dtype uint8 for Tensor with dtype float32: <tf.Tensor: shape=(128, 128, 2), dtype=float32, numpy=
array([[[1., 0.],
    [1., 0.],
    [1., 0.],
    ...,

[sumitsoman]

Had to change mask type to uint8 in DataLoader.