Efficiently extract a patch from image and lable

Question

I have a segmentation project. I have images and labels, which holds ground truth for the segmentation. The images are large, and contains a lot of "empty" areas. I want to cut patches from image and label, so that the patch will have non-zero labeling in it.

I need it to be efficient as possible.

I wrote the following code, but its very slow. any improvement will be highly appreciated.

import numpy as np
import matplotlib.pyplot as plt

Lets create dummy data

img = np.random.rand(300,200,3)
img[240:250,120:200]=0

mask = np.zeros((300,200))
mask[220:260,120:300]=0.7
mask[250:270,140:170]=0.3

f, axarr = plt.subplots(1,2, figsize = (10, 5))
axarr[0].imshow(img)
axarr[1].imshow(mask)[![enter image description here][1]][1]
plt.show()

My inefficient code:

IM_SIZE = 60     # Patch size

x_min, y_min = 0,0
x_max = img.shape[0] - IM_SIZE
y_max = img.shape[1] - IM_SIZE
xd, yd, x, y = 0,0,0,0

if (mask.max() > 0):
    xd, yd = np.where(mask>0)

    x_min = xd.min()
    y_min = yd.min()
    x_max = min(xd.max()- IM_SIZE-1, img.shape[0] - IM_SIZE-1)
    y_max = min(yd.max()- IM_SIZE-1, img.shape[1] - IM_SIZE-1)

    if (y_min >= y_max):

        y = y_max
        if (y + IM_SIZE >= img.shape[1] ): 
            print('Error')

    else:
        y = np.random.randint(y_min,y_max)

    if (x_min>=x_max):

        x = x_max
        if (x+IM_SIZE >= img.shape[0] ):
            print('Error')

    else:
        x = np.random.randint(x_min,x_max )
print(x,y)    
img = img[x:x+IM_SIZE, y:y+IM_SIZE,:]
mask = mask[x:x+IM_SIZE, y:y+IM_SIZE]

f, axarr = plt.subplots(1,2, figsize = (10, 5))
axarr[0].imshow(img)
axarr[1].imshow(mask)
plt.show()

Answer 1

A snapshot of the result given by line profiler is as follows :

Most of the time is used by mask.max() (which can be changed to np.max(mask) for some speedup) and np.where(mask>0).

If you need to use a where function on a different mask each time then take a look at numexpr. Or you could use joblib to store the results for x/y_min/max for a given mask by running many such cases in parallel.

Rearranging the function by using numba.jit gives me better results :

@jit
def temp(mask):
    xd, yd = np.where(mask>0)

    x_min = np.min(xd)
    y_min = np.min(yd)
    x_max = min(np.max(xd)- IM_SIZE-1, img.shape[0] - IM_SIZE-1)
    y_max = min(np.max(yd)- IM_SIZE-1, img.shape[1] - IM_SIZE-1)
    return x_min,x_max,y_min,y_max

def solver_new(img):
    IM_SIZE = 60     # Patch size

    x_min, y_min = 0,0
    x_max = img.shape[0] - IM_SIZE
    y_max = img.shape[1] - IM_SIZE
    xd, yd, x, y = 0,0,0,0

    if (np.max(mask) > 0):
        x_min,x_max,y_min,y_max = temp(mask)
        if (y_min >= y_max):

            y = y_max
            if (y + IM_SIZE >= img.shape[1] ): 
                print('Error')

        else:
            y = np.random.randint(y_min,y_max)

        if (x_min>=x_max):

            x = x_max
            if (x+IM_SIZE >= img.shape[0] ):
                print('Error')

        else:
            x = np.random.randint(x_min,x_max )
    return x,y

Since the image and patch sizes are small, the result's aren't too meaningful as caching has a big effect on the timing. I'm getting roughly ~200us for the implementation posted in the question and ~90us for the one posted here.