通过NumPy 3D数组进行循环的矢量化代码

Question

我当前正在编写一个处理图像并返回相同图像的程序，除了一些被识别为某个特征（通过truth_table找到的像素）被涂黑之外。它可以工作，但是由于图像大约为5000 x 5000像素，因此只需花费整整1分钟的时间：

image_DN_no_12s = np.copy(image_DN)
for y in range(y_size):
    for x in range(x_size):
        if truth_table[y, x] == 12:
            image_DN_no_12s[y, x, :] = np.zeros((1, 1, 4))

（np.zeros()中的第3维为4，因为我正在读取红色，绿色，蓝色和红外通道。）

我想将这些代码向量化，以使其运行更快。我曾尝试使用np.where()和布尔索引，但似乎找不到找到使它运行更快的方法。

预先感谢您的帮助！

Answer 1

如果img和truth_table是其中truth_table具有可广播形状的数组：

img[truth_table == 12] = 0

由于广播会为您处理，因此您无需显式创建4个频道的另一个轴。