向量化逻辑索引

Question

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

是否可以矢量化该提取过程，即通道上的循环？

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

到目前为止，我的预感是通过掩蔽是不可能的，因为先验不清楚每个掩码具有相同数量的

True

值。

import numpy as np
np.random.seed(0)
# set up image
img = np.zeros((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

Answer 1

您可以通过重塑图像和掩模阵列，然后重塑结果来做到这一点：

import numpy as np
rng = np.random.default_rng()
# set up image
# img = np.zeros((3, 100, 111))
img = rng.random((3, 100, 111)) # changed so that img is random not just zeros

# set up some mask with same number of "True" pixels per channel
p = 0.3
mask_array = np.stack([rng.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)

vector_output = img[mask_array].reshape((3, -1))

assert np.allclose(output, vector_output) # True