visualization.py

import numpy as np

def mask2rgb(mask, palette):
    img = np.zeros((mask.shape[0], mask.shape[1], 3))
    for l in np.unique(mask):
        img[mask == int(l)] = palette[int(l)]
    return img.astype(int)