矢量化逻辑索引

问：

我有一个有三个通道的图像。我还有一个掩码，可以从每个通道中随机提取相同数量的像素（但不一定在各个通道之间相同）。

是否可以对提取过程（即通道上的循环）进行矢量化？

这是我试图以非矢量化的方式做的事情：

到目前为止，我的直觉是，通过掩码是不可能的，因为先验地不清楚每个掩码是否具有相同数量的值。True

import numpy as np
np.random.seed(0)
# set up image
img = np.random.random((3, 100, 111))

# set up some mask with same number of "True" pixels per channel
p = 0.3
mask_array = np.stack([np.random.permutation(np.prod(img.shape[1:])).reshape(img.shape[1:]) > p for _ in range(img.shape[0])], axis=0)
print(mask_array.shape) # same as img.shape

# can we vectorize the loop over k away?
output = np.stack([img[k, mask_array[k, ...]] for k in range(img.shape[0])], axis=0)

print(output.shape)  # img.shape[0] x N