[FIXED] Vectorize logical indexing

Issue

I have an image with three channels. I also have a mask that randomly extracts the same number of pixels from each channel (but not necessarily the same across channels).

Is it possible to vectorize that extracting process, that is, the loop over the channels?

Here is what I'm trying to do in an unvectorized way:

My hunch so far is that it is not possible via the masking, as it is a priori not clear that each of the mask has the same number of True values.

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

Solution

Apply the mask to your image will result in a flattened array. You can then reshape the result:

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

assert np.allclose(output, vector_output) # Assertion succeeds

Answered By - Nin17