PIL: generate an image from applying a gradient to a NumPy array

Question

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.

To clarify, instead of drawing bands of colors in a linear gradient as in this example, I'd like to specify apply a gradient colormap to an existing 2d array and turn it into a 3d array.

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.

Answer 1

Method 1: vectorization of your code

Your code is almost already vectorized. Almost all operations of it can work indifferently on a float or on an array of floats

Here is a vectorized version

def get_color_atArr(arr):
    assert (arr>=0).all() and (arr<=1).all()
    n=len(gradient)
    gradient.append(gradient[-1])
    gradient=np.array(gradient, dtype=np.uint8)
    pos = arr*(n-1)
    idx1 = pos.astype(np.uint8)
    idx2 = idx1+1
    frac = (pos - idx1)[:,:,None]
    color1 = gradient[idx1]
    color2 = gradient[idx2]
    color_in_between = np.round(color1*(1-frac) + color2*frac).astype(np.uint8)

Basically, the changes are,

• the assert (can't use a<b<c notation with numpy arrays). Note that this assert iterates all values of array to check for assertion. That is not for free. So I included it because you did. But you need to be aware that this is not a compile-time verification. It does run code to check all values, which is a non-negligible part of all execution time of the code.
• more an implementation choice than a vectorization step (a pure translation of your code would have translated that if x==1 into some np.where, or masks. But I am never comfortable with usage of == on floats any way. So I prefer my way. Which costs nothing. It is not another iteration on the image. It adds a sentinel (In Donald Kuth sense of "sentinel": a few bytes that avoid special cases) to the gradient color. So that, in the unlikely even that arr is really 1.0, the gradient happen between last color and last color).
• frac is broadcasted in 3D array, so that it can be used as a coefficient on 3d arrays color1 and color2
• Plus of course, int or floor can't be used on numpy arrays

Method 2: not reinventing the wheel

Matplotlib (and, I am certain, many other libraries) already have a whole colormap module to deal with this kind of transformations. Let's use it

thresh=np.linspace(0,1,len(gradient))
cmap=LinearSegmentedColormap.from_list('mycmap', list(zip(thresh, np.array(gradient)/255.0)), N=256*len(gradient))
arr2 = cmap(arr)[:,:,:3]

This is building a custom colormap, using LinearSegmentedColormap, which takes, as 2nd argument, a list of pair (threshold, color). Such as [(0, (0,0,0)), (0.3, (1,0,0)), (0.8, (0,1,0)), (1, (0,0,1))] for a color map that goes from black to red when x goes from 0 tom 0.3, then from red to green when x goes from 0.3 to 0.8, then from green to blue.

In this case, your gradient can be transformed to such a list, with just a zip with a linspace.

It takes a N= argument, since it creates a discretization of all possible colors (with interpolation in between). Here I take an exaggerated option (my N is more than the maximum number of different colors than can exist, once uint8d)

Also since it returns a RGBA array, and to remain strictly identical to what you did, I drop the A using [:,:,:3].

Of course, both method need the final translation into PIL, but you already know how to do that. For this one, it also needs mapping between 0 and 255, which I can do with your own code:
Image.fromarray((255 * arr).astype(np.uint8))

Note that, while using matplotlib colormap, you may want to take a tour at what that module has to offer. For example some of the zillions of already existing colormaps may suit you. Or some other way to build colors map non-linearly.