Sandeep2803/traffic-segmentation-alexnet

1import numpy as np2 3# BDD100K / Cityscapes-like color palette4COLOR_MAP = np.array([5    [128, 64, 128],    # Road6    [244, 35, 232],    # Sidewalk7    [70, 70, 70],      # Building8    [102, 102, 156],   # Wall9    [190, 153, 153],   # Fence10    [153, 153, 153],   # Pole11    [250, 170, 30],    # Traffic Light12    [220, 220, 0],     # Traffic Sign13    [107, 142, 35],    # Vegetation14    [152, 251, 152],   # Terrain15    [70, 130, 180],    # Sky16    [220, 20, 60],     # Person17    [255, 0, 0],       # Rider18    [0, 0, 142],       # Car19    [0, 0, 70],        # Truck20    [0, 60, 100],      # Bus21    [0, 80, 100],      # Train22    [0, 0, 230],       # Motorcycle23    [119, 11, 32]      # Bicycle24], dtype=np.uint8)25 26 27def colorize_mask(mask):28    """29    Convert class IDs to RGB colors.30    """31 32    height, width = mask.shape33 34    color_mask = np.zeros((height, width, 3), dtype=np.uint8)35 36    for cls in range(len(COLOR_MAP)):37        color_mask[mask == cls] = COLOR_MAP[cls]38 39    return color_mask