utils.py 文件源码

def colorize_raster(masks):
    ''' (H, W, 10) -> (H, W, 3)
    '''
    assert masks.shape[2] == 10
    palette = np.array([(180, 180, 180), (100, 100, 100),  # Buildings, Misc.
                        (6, 88, 179), (125, 194, 223),  # Road, Track
                        (55, 120, 27), (160, 219, 166),  # Trees, Crops
                        (209, 173, 116), (180, 117, 69),  # Waterway, Standing
                        (67, 109, 244), (39, 48, 215)], dtype=np.uint8)  # Car

    r = []
    for obj_type in range(10):
        c = palette[obj_type]
        result = np.stack([masks[:, :, obj_type]] * 3, axis=2)
        r.append(result * c)
    r = np.stack(r)
    r = np.max(r, axis=0)
    return r