在numpy数组中映射值

Question

如何将仅包含三个不同值（-1、0和1）的2D NumPy数组映射到颜色red(255,0,0)、green(0,255,0)和blue(0,0,255)？该数组非常大，但为了让您了解我要寻找什么，假设我有以下输入：

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

我希望得到以下输出：

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

我可以使用for循环和条件语句来实现，但是我想知道是否有一种使用lambda函数的一两行代码可以完成这个任务？谢谢！

- Jane Sully

Answer 1

你可能需要考虑使用结构化数组，因为它允许元组，而不需要数据类型为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)]]

---

另一种选择是使用一个三维数组，其中最后一个维度是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.]]])

- iz_

---

这是一个不错的想法！唯一的问题是输出 new 的格式为 dtype=[('r', '<i4'), ('g', '<i4'), ('b', '<i4')]。我应该提到这一点，但我的目标是调用Maplotlib的imshow，所以当我执行 plt.imshow(new) 时，会出现错误 Image data cannot be converted to float。你有任何想法如何解决这个问题吗？ - Jane Sully

使用 img.view((np.uint32, 3)) 将其转换为 matplotlib 可以理解的形式。 - Eric

@JaneSully 除了Eric的选项之外，我在我的答案中添加了另一个可能性。 - iz_