[?] Unable to draw colored segmentation mask

[thiagoribeirodamotta]

Using fiftyone-1.7.2, fiftyone_brain-0.21.2 and fiftyone_db-1.3.0.

How do I use colored masks to correctly draw a segmentation mask with a custom color?

I've created the following code to display the segmentation masks:

#Create a new Dataset
dataset_name = "segmentation_dataset"
dataset = fo.Dataset()

for img_path, mask_path in tqdm(zip(data_list, gt_list)):
    # Read RGB image for display (not used for masks)
    img_bgr = cv2.imread(img_path)
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB) #Read base image

    # Read the label mask as single-channel (ids 0..N)
    # print(f"Reading mask: {mask_path}")
    #Read ground truth mask (each class is drawn with a unique ID, from 1 
    mask_ = cv2.imread(mask_path, cv2.IMREAD_UNCHANGED)
    mask_ = cv2.cvtColor(mask_, cv2.COLOR_BGR2GRAY) to 7)
                
    sample = fo.Sample(filepath=img_path)

    #Read prediction mask (each class is drawn with a unique ID, from 1 to 7)                
    pred_mask = cv2.imread(pred_mask_path, cv2.IMREAD_UNCHANGED)
    pred_mask = cv2.cvtColor(pred_mask, cv2.COLOR_BGR2GRAY) 

    # Create color-coded mask
    colored_mask = np.zeros((*mask_.shape, 3), dtype=np.uint8)
    for class_id, color in colors_gt.items():
        if class_id != 0:  # Skip background
            cls_mask = (mask_ == class_id)
            colored_mask[cls_mask] = color 
            # This is how I understood that colored segmentation mask works. 
            # A single RGB matrix with all segmentation masks, each with a different color.
            # Also tried with an individual RGB matrix for each colored segmentation mask, but got the same result.
    
    for i in range(1, 7):
        # Build GT binary mask for class i (0/255)
        cls_mask = (mask_ == i)
        tmp_mask = np.zeros(mask_.shape, dtype=np.uint8)
        tmp_mask[cls_mask] = 255
        sample[f"{classes[i]}"] = fo.Segmentation(mask=tmp_mask)

    for i in range(1, 7):
        # Build predicted binary mask for class i (0/255)
        pred_bin = np.zeros(mask_.shape, dtype=np.uint8)
        if pred.ndim == 3 and i < pred.shape[-1]:
            cls_pred = pred[:, :, i] > 0.5
            pred_bin[cls_pred] = 255

        sample[f"predict_{classes[i]}"] = fo.Segmentation(mask=pred_bin)
    dataset.add_sample(sample)

dataset.save()

But then, to display them with the desired colors, I had to hack my way around it:

session = fo.launch_app(dataset, auto=False) #port=5158
gt_colors = [rgb_to_hex(*v) for v in list(colors_gt.values())[1:]]
pred_colors = [rgb_to_hex(*v) for v in list(colors_gt.values())[1:]]
session.color_scheme.color_pool = session.color_scheme.color_pool[:7] + gt_colors + pred_colors

colors_gt is defined as:

colors_gt = {
    0: (255,255,255),
    1: (230, 25, 75),
    2: (245, 130, 49),
    3: (255, 225, 25),
    4: (67, 99, 216),
    5: (60, 180, 75),
    6: (139, 255, 218),  
}

How do I correctly create RGB segmentation masks? If I simply do the following without messing with the color_pool:
sample[f"{classes[i]}"] = fo.Segmentation(mask=colored_mask)
The custom colors aren't respected and the default color pool colors are used.

What am I doing wrong?