Converting RGB label image to categorical image

[brudfors]

Hello,

Could you please tell me how to best convert one of your RGB encoded label images of shape (nx, ny, 3) to a categorical image of shape (nx, ny), which takes values between 0 and 5 (the number of classes).

Thank you

[elisim]

+1 This will be helpful

[brudfors]

Hi @elisim

I ended up using a MATLAB function for doing this -- my code is below. It could hopefully be ported to Python without too much effort.

% path to RGB label image
fname = 'V1/Test_Labels/1.bmp'
% load RGB image
img = imread(fname)  % (nx, ny, 3)
% color map corresponding to colors in RGB image
cmap = [255 0 0;
        127 0 0;
        255 255 0;
        0 0 255;
        0 255 255;
        0 255 0];
cmap = cmap./255;
% convert RGB to categorical image
cat = rgb2ind(img,cmap,'nodither');  % (nx, ny)

[elisim]

Thank you @brudfors.

I wrote the following code for converting mask to label and label to mask

LABEL_CODES = [(255,0,0),    # background
               (127,0,0),    # hair
               (255,255,0),  # skin
               (0,0,255),    # eyes
               (0,255,255),  # nose
               (0,255,0)]    # mouth

n_labels = 6

def mask2label(mask):
    height, width, ch = mask.shape

    mask_labels = np.zeros((height, width), dtype=np.float32)
    for h_ in range(height):
        for w_ in range(width):
            r_, g_, b_ = mask[h_, w_, :]
            color = (r_, g_, b_)
            color = closest_color(color, LABEL_CODES)
            mask_labels[h_, w_] = color2label[color]
    return mask_labels       

def label2mask(labelmask):
    height, width = labelmask.shape

    mask = np.zeros((height, width, 3), dtype=np.float32)
    for h_ in range(height):
        for w_ in range(width):
            label_ = labelmask[h_, w_]
            mask[h_, w_] = label2color[label_]
    return mask

where closest_color implemented as:

def closest_color(rgb, colors):
    """
    :param rgb: color in rgb format (tuple)
    :param colors: list of rgb colors
    :return: closest color to `rgb` from `colors`. "closest" determined as minimum Euclidean distance
    """
    r, g, b = rgb
    color_diffs = []
    for color in colors:
        cr, cg, cb = color
        color_diff = np.sqrt((r - cr) ** 2 + (g - cg) ** 2 + (b - cb) ** 2)
        color_diffs.append((color_diff, color))
    return min(color_diffs)[1]