I have a 2d NumPy array with values from 0 to 1. I want to turn this array into a Pillow image. I can do the following, which gives me a nice greyscale image:

arr = np.random.rand(100,100)
img = Image.fromarray((255 * arr).astype(np.uint8))

Now, instead of making a greyscale image, I’d like to apply a custom gradient.

Example: If my gradient is [color1, color2, color3], then all 0s should be color1, all 1s should be color3, and 0.25 should be somewhere in between color1 and color2. I was already able to write a simple function that does this:

gradient = [(0, 0, 0), (255, 80, 0), (0, 200, 255)] # black -> orange -> blue

def get_color_at(x):
    assert 0 <= x <= 1
    n = len(gradient)
    if x == 1:
        return gradient[-1]
    pos = x * (n - 1)
    idx1 = int(pos)
    idx2 = idx1 + 1
    frac = pos - idx1
    color1 = gradient[idx1]
    color2 = gradient[idx2]
    color_in_between = [round(color1[i] * (1 - frac) + color2[i] * frac) for i in range(3)]
    return tuple(color_in_between)

So get_color_at(0) returns (0,0,0) and get_color_at(0.75) equals (153, 128, 102), which is this tan/brownish color in between orange and blue.

Now, how can I apply this to the original NumPy array? I shouldn’t apply get_color_at directly to the NumPy array, since that would still give a 2d array, where each element is a 3-tuple. Instead, I think I want an array whose shape is (n, m, 3), so I can feed that to Pillow and create an RGB image.

If possible, I’d prefer to use vectorized operations whenever possible – my input arrays are quite large. If there is builtin-functionality to use a custom gradient, I would also love to use that instead of my own get_color_at function, since my implementation is pretty naive.

Thanks in advance.