Efficient way to map a numpy array of values to an image-CodePudding

I'm looking for the most efficient way to perform the following task.

I have a numpy array with integer values and I have a color map which is a dictionary mapping integers to rgb colors.

What I need is to create for each width by heigth numpy array a width by height by 3 numpy array that can be interpreted as a color image.

For example

import numpy as np

colorMap = {1: [22,244,23], 2: [220,244,23], 3: [22,244,230]}

x = np.array([[1,2,2],[2,2,3],[3,1,2] ])

#I need a very efficient function to create a color image from these two components
image = f(x, colorMap)

My current approach is as follows

import numpy as np

colorMap = {1: [22,244,23], 2: [220,244,23], 3: [22,244,230]}

x = np.array([[1,2,2],[2,2,3],[3,1,2] ])


def f(x):
    return colorMap[x]

x = x.flatten()

image = np.reshape(np.array(list(map(f, x))) , (3,3,3))

But when I time this it is rather slow when compared to numpy inbuilt functions. I'm wondering if anyone knows a way to do this using numpy built in functions that would speed up the procedure.

The above is a dummy example, in reality I need to map large rasters to a visualization in real time. The problem is that the colorMap can be rather long (length between 1 and 100) so that looping over the color map is not a really good option. (If I could loop over the colorMap I would see how to do this with numpy built in functions)

CodePudding user response：

Since the keys of the dictionary are small positive integers, you can convert the colorMap from the dictionary to a numpy array, and then directly use the index:

>>> color_map = [[0] * 3]
>>> color_map  = colorMap.values()
>>> color_map = np.array(color_map)
>>> color_map
array([[  0,   0,   0],
       [ 22, 244,  23],
       [220, 244,  23],
       [ 22, 244, 230]])
>>> color_map[x]
array([[[ 22, 244,  23],
        [220, 244,  23],
        [220, 244,  23]],

       [[220, 244,  23],
        [220, 244,  23],
        [ 22, 244, 230]],

       [[ 22, 244, 230],
        [ 22, 244,  23],
        [220, 244,  23]]])

CodePudding user response：

You can apply the mapping with a simple for loop:

def int_to_rgb(arr, mapping):
    res = arr.copy()[..., None].repeat(3, axis=-1)
    for old_value, new_value in mapping.items():
        res[arr == old_value] = new_value
    return res

To be used as follows:

import numpy as np

mapping = {
    10: (255, 102, 102),
    20: (255, 178, 102),
    30: (255, 255, 102),
    40: (178, 255, 102),
    50: (102, 255, 102),
}
orig_values = list(mapping.keys())

arr = np.random.choice(orig_values, size=(15, 15))
res = int_to_rgb(arr, mapping)

Looking at the result with an imshow:

import matplotlib.pyplot as plt

fig, (ax_orig, ax_res) = plt.subplots(ncols=2)

ax_orig.set_title("Original array")
ax_orig.imshow(arr, cmap="gray")

ax_res.set_title("RGB mapped")
ax_res.imshow(res)

plt.show()