The docs for keras.visualization.draw_segmentation_masks say segmentation mask values should start at 1 and go up to num_classes with class 0 reserved for bg. They also say color_mapping keys should start at 1.

However, keras/src/visualization/draw_segmentation_masks.py treats classes as zero-based:

• num_classes is inferred from max(segmentation_masks)
• color_mapping is read with range(num_classes), requiring keys 0..num_classes - 1
• ops.one_hot(segmentation_masks, num_classes) makes label num_classes out of range

This means documented usage like color_mapping={1: ..., 2: ...} can raise KeyError: 0 and a documented mask label equal to num_classes can fail during one-hot encoding.

Minimal repro:

import numpy as np
from keras.visualization import draw_segmentation_masks

images = np.zeros((1, 2, 2, 3), dtype="uint8")
masks = np.array([[[1, 2], [0, 2]]], dtype="int32")[..., None]

draw_segmentation_masks(
    images,
    masks,
    num_classes=2,
    color_mapping={
        1: [255, 0, 0],
        2: [0, 255, 0],
    },
)

Actual behavior:

KeyError: 0

The failing line is:

colors = [color_mapping[i] for i in range(num_classes)]

Using the default also fails for the documented label value 2 when num_classes=2:

draw_segmentation_masks(images, masks, num_classes=2, blend=False)

With the NumPy as backend this raises:

IndexError: index 2 is out of bounds for axis 1 with size 2

Expected behavior (one of these should be there):

1. The implementation should match the docs and support labels 1..num_classes, with 0 reserved as background
2. The docs should be updated to say labels and color_mapping keys are zero-based: 0..num_classes - 1.

I can help with a fix but what out of those 2 options should be preferable?