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!

iz_ · Accepted Answer · 2019-01-22 02:29:50Z

4

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.]]])

edited Jan 22, 2019 at 2:29

answered Jan 21, 2019 at 23:26

iz_

16.7k4 gold badges29 silver badges43 bronze badges

Sign up to request clarification or add additional context in comments.

3 Comments

Jane Sully Over a year ago

This is a nice idea! The only thing is the output new is in the format dtype=[('r', '<i4'), ('g', '<i4'), ('b', '<i4')]. I should have mentioned this, but my goal is to call Maplotlib's imshow so when I do plt.imshow(new) I get the error Image data cannot be converted to float. Do you have any ideas how to fix this?

Eric Over a year ago

Use img.view((np.uint32, 3)) to get in in a form matplotlib understands

iz_ Over a year ago

@JaneSully In addition to Eric's option, I added another possibility to my answer.

Collectives™ on Stack Overflow

Mapping values in a numpy array

1 Answer 1

3 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

3 Comments

Your Answer

Sign up or log in

Post as a guest

Related