Convert multi-dimensional Numpy array to 2-dimensional array based on color values

Question

I have an image which is read as a uint8 array with the shape (512,512,3). Now I would like to convert this array to a uint8 array of shape (512,512,1)

Answer 1

One way: separately create the boolean arrays with True values where the input's pixel value matches one of the palette values, and then use arithmetic to combine them. Thus:

palette = [
    [0, 0, 0], 
    [0, 0, 255], 
    [255, 0, 0],
    # etc.
]

def palettized(data, palette):
    # Initialize result array
    shape = list(data.shape)
    shape[-1] = 1
    result = np.zeros(shape)
    # Loop and add each palette index component.
    for value, colour in enumerate(palette, 1):
        result += (data == colour).all(axis=2) * value
    return result