如何将 2d numpy 数组保存并显示为 RGB 伪彩色图像？

Question

我有一个大尺寸的 2d numpy 数组 (size = (2000, 2000))，只有五个可能的值 1.0, 2.0, 3.0, 4.0 和 5.0。我想将此数组保存并显示为 RGB 彩色格式的图像，这里数组的每个唯一值都应该用不同的颜色表示。请帮助我是python的初学者。

Answer 1

您可以使用 PIL.Image 来执行此操作，但首先要转换您的数组。

例如，你可以说：

1. 1.0应该是红色的，表示为(255,0,0)
2. 2.0 应该是绿色 -> (0,255,0)
3. 3.0 应该是蓝色 -> (0,0,255)
4. 4.0 应该是黑色 -> (0,0,0)
5. 5.0 应该是白色 -> (255,255,255)

您当然可以将这些值更改为您选择的任何颜色，但这只是为了演示。话虽如此，您的二维数组也需要“展平”为一维，以便 PIL.Image 将其作为数据接受。

from PIL import Image
import numpy as np
your_2d_array = np.something() # Replace this line, obviously
img_array = []
for x in your_2d_array.reshape(2000*2000):
    if x == 1.0:
        img_array.append((255,0,0)) # RED
    elif x == 2.0:
        img_array.append((0,255,0)) # GREEN
    elif x == 3.0:
        img_array.append((0,0,255)) # BLUE
    elif x == 4.0:
        img_array.append((0,0,0)) # BLACK
    elif x == 5.0:
        img_array.append((255,255,255)) # WHITE

img = Image.new('RGB',(2000,2000))
img.putdata(img_array)
img.save('somefile.png')

虽然这应该可行，但我认为有更有效的方法可以做到这一点，我不知道，所以如果有人用更好的例子编辑这个答案，我会很高兴。但是，如果它是一个小型应用程序并且最高效率不会打扰您，那就是。

Answer 2

matplotlib 对于这样的任务很有用，尽管还有其他方法。 这是一个例子：

import numpy as np
import matplotlib.image

src = np.zeros((200,200))
print src.shape

rgb = np.zeros((200,200,3))
print rgb.shape

src[10,10] = 1
src[20,20] = 2
src[30,30] = 3

for i in range(src.shape[0]):
  for j in range(src.shape[1]):
    rgb[i,j,0] = 255 if src[i,j]==1 else 0 # R
    rgb[i,j,1] = 255 if src[i,j]==2 else 0 # G
    rgb[i,j,2] = 255 if src[i,j]==3 else 0 # B

matplotlib.image.imsave('test.png', rgb.astype(np.uint8))

诀窍是将其转换为形状为(x, y, 3) 的RGB 数组。您可以使用任何您想要生成每像素 RGB 值的公式。

另外，请注意将其转换为 uint8 数组。