索引numpy數組的最后一個維度

Question

我有兩個圖像，分別為image和mask，具有以下形狀：

image shape: (876, 1020, 3)
mask shape: (876, 1020)

我想做的是將圖像的三個元素中的每一個都除以遮罩，其中遮罩為非零值。 因此，我這樣做：

import numpy as np

index = np.nonzero(mask)
image[index, :] = image[index, :]/mask[index]

但是，這會使程序花費很長時間，然后失敗。

ValueError：操作數不能與形狀（2,302793,1020,3）一起廣播（302793，）

Answer 1

如果創建掩碼副本沒有內存問題，則可以將零元素設置為1並直接除法：

mask2 = mask.copy()
mask2[mask2==0]=1
image /= mask2[...,np.newaxis]