Segmentation class ID color lookup off by one for class IDs > 128

Describe the bug

When specifying the AnnotationContext for a SegmentationImage, the viewer visualizes the classes with class ID larger than 128 with colors from one class ID below (so class_id - 1). Labels are visualized correctly.

In the screenshot below, 11 different classes with unique colors are logged in a SegmentationImage. For class IDs up to (and including) 128 the colors and labels are visualized as expected. But for class ID 129 and above, the color look up is off by one: class ID 129 takes the color of 128, class ID 130 takes the color of 129 and so on. Two notes:

• unlike the color, the class label is looked up correctly (when hovering over pixels with mouse)
• annotation context on the sidepane shows the correct colors

To Reproduce Steps to reproduce the behavior:

1. Log a segmentation mask with class IDs larger than 128
2. Log annotation context and specify the labels and colors
3. Compare the visualized SegmentationImage colors with the AnnotationContext
4. Colors of class IDs larger than 128 are off from the AnnotationContext by one

Expected behavior The visualized colors of SegmentationImage should follow the AnnotationContext for all 16 bit integers.

Screenshots

Desktop (please complete the following information):

• OS: Ubuntu 24.04 LTS

Rerun version 0.38.1

Code snippet

import numpy as np
import rerun as rr


SEGMENTATION_CLASSES = {
    0: ("black", (0, 0, 0)),
    1: ("red", (255, 0, 0)),
    2: ("yellow", (255, 255, 0)),
    64: ("light_blue", (0, 128, 255)),
    126: ("light_green", (128, 255, 0)),
    127: ("magenta", (255, 0, 255)),
    128: ("blue", (0, 0, 255)),
    129: ("grey", (128, 128, 128)),
    130: ("orange", (255, 128, 0)),
    131: ("purple", (128, 0, 128)),
    132: ("green", (0, 255, 0)),
}

def main():
    rr.init("test", spawn=True)

    rr.log(
        "segmentation_mask",
        rr.AnnotationContext([
            (class_id, name, color)
            for class_id, (name, color) in SEGMENTATION_CLASSES.items()
        ]),
        static=True,
    )

    # Test mask: one equal-width vertical band per class, in ascending class id order.
    class_ids = sorted(SEGMENTATION_CLASSES)
    band_width = 8
    mask = np.zeros((8, band_width * len(class_ids)), dtype=np.uint8)
    for i, class_id in enumerate(class_ids):
        mask[:, i * band_width:(i + 1) * band_width] = class_id

    rr.log("segmentation_mask", rr.SegmentationImage(mask))


if __name__ == "__main__":
    main()