utils2.py

import torch

def print_mask(mask:torch.Tensor, numClasses:int=19)->None:
    """
        Visualizes the segmentation mask by mapping each class to a specific color.
    
    Args:
        mask (torch.Tensor): The segmentation mask to visualize.
        numClasses (int, optional): Number of classes in the segmentation mask. Defaults to 19.
    """
    colors = [
        (128, 64, 128),  # 0: road
        (244, 35, 232),  # 1: sidewalk
        (70, 70, 70),    # 2: building
        (102, 102, 156), # 3: wall
        (190, 153, 153), # 4: fence
        (153, 153, 153), # 5: pole
        (250, 170, 30),  # 6: traffic light
        (220, 220, 0),   # 7: traffic sign
        (107, 142, 35),  # 8: vegetation
        (152, 251, 152), # 9: terrain
        (70, 130, 180),  # 10: sky
        (220, 20, 60),   # 11: person
        (255, 0, 0),     # 12: rider
        (0, 0, 142),     # 13: car
        (0, 0, 70),      # 14: truck
        (0, 60, 100),    # 15: bus
        (0, 80, 100),    # 16: train
        (0, 0, 230),     # 17: motorcycle
        (119, 11, 32)    # 18: bicycle
    ]

    new_mask = torch.zeros((mask.shape[0], mask.shape[1], 3),dtype=torch.uint8)
    new_mask[mask == 255] = (0,0,0)
    for i in range (numClasses):
        new_mask[mask == i] = colors[i][:3]
    return new_mask.permute(2,0,1)