如何用NumPy将二维0-1数组转换为对应RGB像素数组？

这里有几种高效的NumPy实现方式，能快速将二维0-1掩码数组转换为对应的RGB数组，满足你传入PIL保存PNG的需求：

方法1：利用数组广播直接计算 #

这是最简洁高效的方式，通过扩展维度并结合标量乘法实现：

import numpy as np

# 将输入列表转为NumPy数组
mask = np.array([[0, 0, 1], [1, 0, 0]])
# 扩展维度为(高度, 宽度, 1)，再与255广播得到3通道RGB数组
rgb_mask = mask[..., np.newaxis] * 255
# 转换为PIL要求的uint8类型
rgb_mask = rgb_mask.astype(np.uint8)

mask[..., np.newaxis]会给原二维数组新增一个维度（形状变为(2,3,1)），之后与标量255相乘时，NumPy会自动广播到3个通道，最终得到(2,3,3)的RGB数组。

方法2：正确使用numpy.where #

你之前遇到的广播问题，是因为没给掩码数组扩展维度。只需先扩展掩码的维度，就能让np.where正确广播：

import numpy as np

mask = np.array([[0, 0, 1], [1, 0, 0]])
black = np.array([0, 0, 0])
white = np.array([255, 255, 255])

# 扩展掩码维度为3D，让where可以在每个像素位置选择对应颜色
rgb_mask = np.where(mask[..., np.newaxis], white, black).astype(np.uint8)

方法3：使用颜色映射索引 #

通过创建颜色映射数组，直接用掩码值作为索引获取对应RGB像素，可读性高且高效：

import numpy as np

mask = np.array([[0, 0, 1], [1, 0, 0]])
# 创建颜色映射：索引0对应黑色，索引1对应白色
color_map = np.array([[0, 0, 0], [255, 255, 255]], dtype=np.uint8)
# 直接用掩码索引映射数组，得到RGB结果
rgb_mask = color_map[mask]

验证结果 #

以上三种方法得到的rgb_mask都与你预期的结果一致，打印输出如下：

print(rgb_mask)
# 输出：
# [[[  0   0   0]
#   [  0   0   0]
#   [255 255 255]]
# 
#  [[255 255 255]
#   [  0   0   0]
#   [  0   0   0]]]

最后将rgb_mask传入PIL保存为PNG即可：

from PIL import Image
Image.fromarray(rgb_mask).save("mask.png")

内容的提问来源于stack exchange，提问作者TheMaster