基于布尔掩码为RGB图像批量赋值NaN：Python2.7高效实现

在Python 2.7里，我们可以借助NumPy的向量化操作+广播机制来实现和Matlab类似的高效处理，完全不需要双重循环，哪怕是10000x10000的超大图像也能快速完成。

核心思路 #

和Matlab的repmat思路类似，我们需要让掩码数组的维度和RGB图像匹配，然后利用布尔索引直接批量赋值。NumPy的广播机制会帮我们自动完成维度扩展，不需要手动复制通道。

具体步骤 #

1. 确保数据类型正确：首先要把RGB图像数组转换成浮点型（比如float64），因为整数数组无法存储NaN值。
2. 维度匹配与赋值：通过扩展掩码的维度，让它和RGB图像的通道维度对齐，然后用布尔索引选中需要设为NaN的位置，一次性赋值。

代码示例 #

情况1：掩码是布尔数组（值为True/False，对应原需求的1/0）

import numpy as np

# 假设img_rgb是形状为(H, W, 3)的NumPy数组，mask是(H, W)的布尔数组
# 先转换为浮点型（如果原数组是整数类型）
img_rgb = img_rgb.astype(np.float64)

# 扩展掩码维度并赋值NaN
# ~mask表示取反，选中mask为False的位置
img_rgb[~mask[..., np.newaxis]] = np.nan

情况2：掩码是整数数组（值为1/0）

如果你的掩码是整数类型（比如0和1），只需要稍作修改：

import numpy as np

img_rgb = img_rgb.astype(np.float64)
# 选中mask为0的位置
img_rgb[mask[..., np.newaxis] == 0] = np.nan

验证示例 #

用你给出的小测试数据验证：

# 构造示例RGB图像
img_rgb = np.array([[[125, 125, 125], [160, 160, 160]],
                    [[130, 130, 130], [125, 125, 125]]], dtype=np.float64)
# 构造掩码
mask = np.array([[1, 0], [0, 1]])

# 执行处理
img_rgb[mask[..., np.newaxis] == 0] = np.nan

# 查看第一个通道的结果
print(img_rgb[:, :, 0])

输出结果和你预期的一致：

[[125. nan]
 [nan 125.]]

为什么高效？ #

这种方式完全基于NumPy的底层向量化操作，所有循环都是在C语言层面完成的，比Python层面的双重循环效率高几个数量级，非常适合处理超大尺寸的图像。

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