rgb2palette.py

import os
import numpy as np
import cv2

def mask_to_label(mask_rgb):
    pascal_palette = np.array([(0, 0, 0), (192, 128, 128)], dtype=np.uint8)
    mask_labels = np.ones(mask_rgb.shape[:2])
    pixel = 0
    
    for i in range(mask_rgb.shape[0]):
        for j in range(mask_rgb.shape[1]):
            try:
                color = list(mask_rgb[i, j, :].astype(np.uint8))
                pixel = pascal_palette.tolist().index(color)
                mask_labels[i, j] = pixel
            except Exception:
                mask_labels[i, j] = pixel
    return mask_labels

labels_path = ""
new_labels_path = ""


for i in os.listdir(labels_path):
    label = cv2.imread(labels_path + i)
    label = label[:,:,::-1]
    palette = mask_to_label(label)
    res = cv2.imwrite(new_labels_path+i, palette)
    print (res)