Mapping values in a numpy array

Question

How do I go from a 2D numpy array where I only have three distinct values: -1, 0, and 1 and map them to the colors red (255,0,0), green (0,255,0), and blue (255,0,0)? The array is quite large, but to give you an idea of what I am looking for, imagine I have the input

array([[ 1,  0, -1],
       [-1,  1,  1],
       [ 0,  0,  1]])

I want the output:

array([[(0, 0, 255), (0, 255, 0), (255, 0, 0)],
       [(255, 0, 0), (0, 0, 255), (0, 0, 255)],
       [(0, 255, 0), (0, 255, 0), (0, 0, 255)]])

I could for-loop and have conditions but I was wondering if there is a one or two liner using a lambda function that could accomplish this? Thanks!

Answer 1

You might want to consider a structured array, as it allows tuples without the datatype being object.

import numpy as np

replacements = {-1: (255, 0, 0), 0: (0, 255, 0), 1: (0, 0, 255)}

arr = np.array([[ 1,  0, -1],
                [-1,  1,  1],
                [ 0,  0,  1]])

new = np.zeros(arr.shape, dtype=np.dtype([('r', np.int32), ('g', np.int32), ('b', np.int32)]))

for n, tup in replacements.items():
    new[arr == n] = tup

print(new)

Output:

[[(  0,   0, 255) (  0, 255,   0) (255,   0,   0)]
 [(255,   0,   0) (  0,   0, 255) (  0,   0, 255)]
 [(  0, 255,   0) (  0, 255,   0) (  0,   0, 255)]]

---

Another option is using an 3D array, where the last dimension is 3. The first "layer" would be red, the second "layer" would be green, and the third "layer" blue. This option is compatible with plt.imshow().

import numpy as np

arr = np.array([[ 1,  0, -1],
                [-1,  1,  1],
                [ 0,  0,  1]])

new = np.zeros((*arr.shape, 3))

for i in range(-1, 2):
    new[i + 1, arr == i] = 255

Output:

array([[[  0.,   0., 255.],
        [255.,   0.,   0.],
        [  0.,   0.,   0.]],

       [[  0., 255.,   0.],
        [  0.,   0.,   0.],
        [255., 255.,   0.]],

       [[255.,   0.,   0.],
        [  0., 255., 255.],
        [  0.,   0., 255.]]])