python - 按 2D 数组索引 3D 数组

Question

我有一个 3D 彩色图像 im(形状 512 512 3)和一个 2D 数组掩模(512 512)。我想用蒙版注释这个彩色图像:

im = im[mask>threshold] + im[mask<threshold] * 0.2 + (255,0,0) * [mask<threshold].

如何用 Python 高效地编写此代码？

Answer 1

这有效:

mask3 = numpy.dstack(mask,mask,mask)
im = im * (mask3>threshold) + im * (mask3<threshold) * 0.2
im[:,:,0] += 255 * (mask<threshold)

它依赖于 true 的数值为 1，false 的数值为 0。

它可能不是最清晰或最有效的，但它仍然可能比 bool 数组索引快得多，例如 im[ mask3 < threshold ] *= 0.2 (除非索引的真实值数量非常少)。