python - Numpy 使用二进制掩码迭代颜色 channel 的更好方法

Question

有没有更好的方法将二元掩码应用于 numpy 中的颜色 channel ？我最终不得不一直这样做，感觉应该这样做。

for c in range(3):
    a_image[mask, c] = b_image[mask, c]

形状是 (x, y, c) a_image 和 b_image，以及 (x, y) 用于掩码。

Answer 1

您可以简单地在 3-D 数组上使用 2-D 掩码，而无需循环。 Numpy 将为您将其广播到三维空间。

a_image[mask] = b_image[mask]

简单的例子:

a_image = np.arange(6).reshape(1,2,3)
#[[[0 1 2]
#  [3 4 5]]]

b_image = np.ones((1,2,3))
#[[[1. 1. 1.]
#  [1. 1. 1.]]]

mask = np.array([[False,True]])
#[[False  True]]

输出:

[[[0 1 2]
  [1 1 1]]]