如何用布尔索引为NumPy数组设置RGB像素值？

我有一个形状为(100, 80, 3)的RGB图像NumPy数组，以及一个形状为(100, 80)的布尔掩码。需要将掩码为True的所有像素设置为pix_val = np.array([0.1, 0.2, 0.3])。

目前已经实现了循环版本的代码：

cols = 100
rows = 80
img = np.random.rand(rows, cols, 3)
mask = np.random.randint(2, size=(rows, cols), dtype=np.bool_)
px = np.array([0.1, 0.2, 0.3])

for ch in range(3):
  img[:, :, ch][mask] = px[ch]

原本尝试用广播方式：

img[mask[:, :, None]] = px

但实际无效，现寻求向量化的高效实现方法。

高效向量化实现方法 #

方法1：直接布尔掩码索引赋值 #

这是最简洁的实现方式，利用NumPy花式索引的特性，二维布尔掩码会自动提取三维数组中掩码为True位置的完整像素（三个通道），目标像素值会被自动广播匹配：

img[mask] = px

原理：img[mask]返回形状为(N, 3)的数组（N为掩码中True的数量），px是形状为(3,)的数组，NumPy会自动将px扩展为(N, 3)的形状，完成逐像素赋值。

方法2：np.where条件赋值 #

适合需要同时保留非掩码区域像素的场景，通过条件判断选择保留原像素或替换为目标值：

img = np.where(mask[..., None], px, img)

原理：mask[..., None]将二维掩码扩展为(100, 80, 1)的三维掩码，使其能与img（(100,80,3)）和px（(3,)）广播匹配。np.where会遍历每个位置，掩码为True时取px对应通道值，否则保留原图像素。

方法3：扩展掩码维度后赋值 #

修正你之前的尝试逻辑，先将掩码扩展为三维再赋值，效果与方法1一致：

mask_3d = mask[:, :, np.newaxis]  # 形状变为(100,80,1)
img[mask_3d] = px

原理：三维掩码索引img后得到(N,3)的像素数组，px自动广播完成赋值，逻辑更直观。

效率对比 #

三种向量化方法的效率远高于循环版本，NumPy底层C实现的运算速度比Python循环快几个数量级。以测试代码为例：

• 循环版本：约几十毫秒（随数组大小波动）
• 向量化版本：约几微秒到几十微秒

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