使用布尔索引为 numpy 数组设置 RGB 值

Question

我有一个形状为

(100, 80, 3)

的数组，它是一个 RGB 图像。
我有一个形状为

(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

会起作用的。但事实并非如此。

我正在寻找一种矢量化（有效）的方法来实现它。

Answer 1

您可以使用：

img[mask, :] = px