[Python] What is the most efficient way to randomly pick one positive location within a large...

I am writing a custom image data loading function to randomly crop part of a large image according to its binary mask. The function will be used in PyTorch dataloader so I want it to be as fast and memory-efficient as possible. The image and the mask are quite large, as both the width and height are on the order of 10k~20k pixels.

I want every crop of the image to contain at last one positive point in the binary mask. My current sollution is to first randomly sample one positive point from the mask image, and then generate a crop box around it. The implementation contains a section of code as follows:

import PIL
import numpy as np

... # Some preprocessing to find all mask and image files

mask = PIL.Image.open(mask_file) # both the width and height have 10k~20k pixels.
# fast_pil_to_numpy: https://uploadcare.com/blog/fast-import-of-pillow-images-to-numpy-opencv-arrays/
mask_np = fast_pil_to_numpy(mask).astype(bool) # dim: [height, width]
mask_loc = np.where(mask_np) # get (loc_y, loc_x) of all positive indices
idx = np.random.randint(low=0, high=len(mask_loc[0]))

... # Generate a crop box around (mask_loc[1][idx], mask_loc[0][idx])



After profiling the entire function with line_profiler, I find that the line mask_loc = np.where(mask_np) is one of the performance bottle-necks. How can I optimize this part? Is there another more efficient way to randomly sample one positive point from a binary image?

Continue reading...