在numpy数组中映射值

时间:2019-01-21 22:52:35

标签: python numpy lambda mapping key-value

如何从2D numpy数组中查找,那里只有三个不同的值:-1、0和1,并将它们映射到颜色red(255,0,0),{{1} }(0,255,0)和green(255,0,0)?数组很大,但是让您大致了解我要寻找的内容,想象一下我有输入

blue

我想要输出:

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

我可以循环并有条件,但是我想知道是否有一个或两个使用lambda函数的班轮可以完成此任务?谢谢!

1 个答案:

答案 0 :(得分:2)

您可能要考虑结构化数组,因为它允许元组而数据类型不是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)

输出:

[[(  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)]]

另一种选择是使用3D数组,最后一个维度是3。第一个“层”为红色,第二个“层”为绿色,第三个“层”为蓝色。此选项与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

输出:

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