如何使用 Numpy 将矩阵中的值替换为颜色 eg[0,255,0]

Question

I have a matrix which only contains 0 or 1. I also have a list of colors eg the below

class_colors = [[0,0,0], [0,255,0]]
m = np.random.random_integers(0,1,(5,5))

eg the m looks like:

0,1,0,1
0,1,0,1
0,1,0,1

How can I replace the 1-values in m with the class_colors[1] and 0-values in m with class_colors[0] , so the m will look something like:

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

I used to be able to do so with np.argmax() and np.take() but it requires the m looks like [class_num,w,h] and then I can do argmax with axis=0 .

I know I could do it with for loop, but is there any better and faster approach to do this?

Answer 1

Not surprisingly the intuitive way works (using advanced indexing):

class_colors = np.array([[0,0,0], [0,255,0]])
m = np.array([0,1,0,1,0,1,0,1,0,1,0,1]).reshape((3,4))
class_colors[m]

Answer 2

Is it what you expect:

>>> c[m]
array([[[  0,   0,   0],
        [  0, 255,   0],
        [  0,   0,   0],
        [  0, 255,   0]],

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

>>> c  # c = np.array(class_colors)
array([[  0,   0,   0],
       [  0, 255,   0]])

>>> m  # m = np.random.randint(0, 2, (2, 4))
array([[0, 1, 0, 1],
       [0, 1, 0, 1]])

By extension:

class_colors = [[0,0,0], [0,255,0], [255,0,0], [0,0,255]]
c = np.array(class_colors)
m = np.random.randint(0, 4, (2, 4))
# m <- array([[2, 1, 1, 3], [2, 3, 1, 0]])

>>> c[m]
array([[[255,   0,   0],
        [  0, 255,   0],
        [  0, 255,   0],
        [  0,   0, 255]],

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