How to efficiently vectorize (i.e. avoid explicit loops) a mapping of some colours into others in a Numpy array

Question

I have as input a Numpy matrix of rank three (i.e. an image: horizontal, vertical and 4 color channels). I want to read this matrix element-wise in their first two indices and map only certain colours into others, defined in respective arrays. The performance is very important, as this mapping will be applied many times, possible becoming a bottleneck of the program.

Precisely the code I have so far is:

# data is the rank-3 Numpy array with the image (obtained using the library PIL)
# palette has shape (11, 4) and defines the 11 colours to map
# palette_grey has shape (11,) and defines the 11 tones of grey to apply

for i in range(palette.shape[0]):  # loop in colours
    match = (data[:,:,0] == palette[i,0]) & (data[:,:,1] == palette[i,1]) & (data[:,:,2] == palette[i,2]) # build matrix only True when a given pixel has the right color
    for j in range(3): # loop to apply the mapping to the three channels (because it's just grey, so all channels are equal)
        data[:,:,j] = np.where(match, grey_palette[i], data[:,:,j]) # the mapping itself

Although the main task is vectorized (via np.where), there are still two explicit loops I'd like to avoid to improve performance.

Any idea to achieve this?

EDIT:

I have tried to remove the second loop (in channels) by defining the both palettes to have the same shape (11,4). Then, I have tried this:

for i in range(palette.shape[0]):
    match = (data[:,:,0] == palette[i,0]) & (data[:,:,1] == palette[i,1]) & (data[:,:,2] == palette[i,2])
    data[:,:,:] = np.where(match, grey_palette[i], data[:,:,:])

But it raises the error:

ValueError: operands could not be broadcast together with shapes (480,480) (4,) (480,480,4)

I guess this is the expected behaviour, but I thought the mapping I propose is unambiguous, and therefore doable by Numpy.

Answer 1

Let's say that you have an (M, N, 4) array of uint8 values. You can view it as an array of 32-bit integers like this:

idata = data.view(np.uint32)

Secondly, you can convert your gray map to a proper (K, 4) array from a (K,) array:

grey = np.repeat(grey[:, None], 4, axis=-1)

While you're at it, you can convert it to integers as well as the palette:

ipalette = palette.view(np.uint32).squeeze()
igrey = grey.view(np.uint32).squeeze()

Now you have an (M, N, 1) image and (K,) palette and grey map. For sufficiently small K, you can do something like this:

ind = np.nonzero(idata == ipalette)
idata[ind[0], ind[1], 0] = igrey[ind[-1]]

This will write the grey values through directly to the original data array. The mask idata == ipalette broadcasts to (M, N, K). np.nonzero returns the indices along each axis. The first two indices are the coordinates of matches in the image, while the third index is exactly the grey value that you want.

This will work for uint16 images as well, since you can use uint64 to aggregate the data.

For larger K, you can take a completely different approach:

1. Pass the raveled idata through np.unique with return_inverse=True. This will give you an array of sorted unique values, and an index that will replace them into image order.
2. Run np.searchsorted to place ipalette into the unique values.
3. Discard any in indices that don't represent an exact match.
4. Replace the exact matches with the grey values.
5. Rebuild the image from the unique values with grey replacements.

Something like this:

uniques, reverse_index = np.unique(idata.ravel(), return_inverse=True)
palette_index = np.searchsorted(uniques, ipalette)
# clipping won't affect the match but avoids errors
palette_mask = (ipalette == uniques[palette_index.clip(0, uniques.size - 1)])
uniques[palette_index[palette_mask]] = igrey[palette_mask]
new_data = uniques[reverse_index].view(np.uint8).reshape(data.shape)

Here is a full working example of the first method:

np.random.seed(0)
data = np.random.randint(0, 255, size=(100, 100, 4), dtype=np.uint8)
palette = np.array([[117, 166,  22, 183],
                    [ 38,  28,  93, 140],
                    [ 63, 214,   9,  84],
                    [185,  51,   1,  24],
                    [131, 210, 145,   3],
                    [111, 180, 165, 245],
                    [ 62, 220, 102, 144],
                    [177,  97, 158, 135],
                    [202,  67, 169,  10],
                    [ 23, 177,  26,  15],
                    [ 19, 100,  25,  66],
                    [ 86, 227, 222, 182],
                    [255, 255, 255, 255]], dtype=np.uint8)
grey = np.arange(10, 10 + len(palette), dtype=np.uint8)

idata = data.view(np.uint32)
ipalette = palette.view(np.uint32).squeeze()
igrey = np.repeat(grey[:, None], data.shape[-1], axis=-1).view(np.uint32).squeeze()

ind = np.nonzero(idata == ipalette)
idata[ind[0], ind[1], 0] = igrey[ind[-1]]

I've selected the palette to be the pixels that are displayed with my default settings when I do >>> data, just so that you can immediately see that the method works.

With the second method, you can see that clipping makes it possible to have the value 0xFFFFFFFF in the palette, which gets and insertion index greater than the size of the array, without crashing or having to introduce any special cases.

Answer 2

Create a dictionary d build from you palettes {palette[i]: grey_palette[i]} (completed by trivial entries {i:i} otherwise), vectorize it and apply to your data

numpy.vectorize(d.get)(data)

It should be fast but I have not tested it with your type of data.

Answer 3

When I compare your solution with this one:

for i in range(palette.shape[0]):
    new_data[data == palette[i]] = grey_palette[i]

using %%timeit in a notebook gives 87ms vs 218ms for yours for a 1000x1000x3 data.

EDIT: deleted comment about a 'problem' with your solution that I created by changing to new_data only in one place.