Crop black border of image using NumPy

Objective: Crop the image so only the number stays in the image

Problem: Slow Performance

I have code that crops an image.
The image pixels are 0 or 255. There are no values between.

The background is 0 (black) and the letter/number is between 0 (not-inclusive) - 255 (white).

This code is being used to crop a Mnist Digit of the Mnist Dataset.

Image Source

The code does it, however, it does with 2 fors and it takes a long time! How can I optimize it?

def crop_image(data_to_crop):
    cropped_data = []

    for z in xrange(data_to_crop.shape[0]):

        img = data_to_crop[z]
        img = img.reshape(data_to_crop.shape[0], data_to_crop.shape[1])

        rx = -1
        upy = -1

        lx = -1
        by = -1

        for x in xrange(data_to_crop.shape[0]):
            for y in xrange(data_to_crop.shape[1]):

                px = img[x, y]

                if px > 0:

                    if rx == -1 or x > rx:
                        rx = x

                    if lx == -1 or x  upy:
                        upy = y

                    if by == -1 or y < by:
                        by = y

        img = img[lx:rx, by:upy]

        cropped_data.append(img)

    return cropped_data

Vectorization with NumPy

When read with cv2.imread or skimage.io.imread or scipy.misc.imread, you would already have the image data as a NumPy array. Now, NumPy supports various vectorization capabilities, which we can use to speed up things quite a bit.
I. Crop to remove all black rows and columns across entire image

To solve our case, one method would be to look for rows and columns that have at least one pixel along rows and columns that is greater than some lower limit or threshold as a pixel value. So, if you are sure that the black areas are absolutely zeros, you can set that threshold as 0. Thus, if img represents the image data, you would have correspondingly two boolean arrays : (img>tol).any(1) and (img>tol).any(0).

Next up, you can use these boolean arrays to index into the image data for extraction of valid bounding box using broadcasted indexing with np._ix -

np.ix_((img>tol).any(1),(img>tol).any(0))

Finally, we index into image data with it for the final extracted data, which is the required bounding box data.

To sum up, the final implementation would be -

def crop_image(img,tol=0):
    # img is 2D image data
    # tol  is tolerance
    mask = img>tol
    return img[np.ix_(mask.any(1),mask.any(0))]

Sample results

1] With tolerance = 0 :

2] With tolerance = 80 (tighter) :

II. Crop while keeping the inner all black rows or columns

To crop the image while keeping the inner all black rows or columns, the implementation would be close to the previous method. The basic idea here would be getting the start, stop indices along rows and columns that decide the bounding box. We will start off with the same mask of ANY match along rows and columns as used in the previous one. Then, argmax would be useful to get start and stop indices of those matches along rows and cols. We will use these indices to slice the 2D input array, which is the desired cropped image output. The implementation would look something like this -

def crop_image_only_outside(img,tol=0):
    # img is 2D image data
    # tol  is tolerance
    mask = img>tol
    m,n = img.shape
    mask0,mask1 = mask.any(0),mask.any(1)
    col_start,col_end = mask0.argmax(),n-mask0[::-1].argmax()
    row_start,row_end = mask1.argmax(),m-mask1[::-1].argmax()
    return img[row_start:row_end,col_start:col_end]

Sample run -

Input :

Output :

Benchmarking

Since we are talking about performance in this Q&A, let's test out how this method works in comparison to others.

@Gareth Rees's solution is another vectorized one that finds all indices and then slices the input array. Finding all indices could become costly. We will try to time and see how much would that affect the performance. we will use the sample data used in the earlier section -

# @Gareth Rees's solution
def crop_with_argwhere(image):
    # Mask of non-black pixels (assuming image has a single channel).
    mask = image > 0
    
    # Coordinates of non-black pixels.
    coords = np.argwhere(mask)
    
    # Bounding box of non-black pixels.
    x0, y0 = coords.min(axis=0)
    x1, y1 = coords.max(axis=0) + 1   # slices are exclusive at the top
    
    # Get the contents of the bounding box.
    cropped = image[x0:x1, y0:y1]
    return cropped

Timings -

# Import the "4" digit image from previous section
In [50]: from skimage import io
    ...: im = io.imread('https://i.stack.imgur.com/tmdiH.png')

# @Gareth Rees's solution
In [51]: %timeit crop_with_argwhere(im)
1000 loops, best of 3: 1.4 ms per loop

In [52]: %timeit crop_image_only_outside(im,tol=0)
10000 loops, best of 3: 81.8 µs per loop

The memory efficiency with crop_image_only_outside is noticeable on performance.

Extend to generic 2D or 3D image data cases

Assuming we are looking to check for ALL matches across all channels along the last dimension/axis, the extension would be simply performing numpy.all reduction along the last axis. Hence, we would have generic solutions to handle both 2D and 3D image data cases like so -

def crop_image(img,tol=0):
    # img is 2D or 3D image data
    # tol  is tolerance
    mask = img>tol
    if img.ndim==3:
        mask = mask.all(2)
    mask0,mask1 = mask.any(0),mask.any(1)
    return img[np.ix_(mask0,mask1)]

def crop_image_only_outside(img,tol=0):
    # img is 2D or 3D image data
    # tol  is tolerance
    mask = img>tol
    if img.ndim==3:
        mask = mask.all(2)
    m,n = mask.shape
    mask0,mask1 = mask.any(0),mask.any(1)
    col_start,col_end = mask0.argmax(),n-mask0[::-1].argmax()
    row_start,row_end = mask1.argmax(),m-mask1[::-1].argmax()
    return img[row_start:row_end,col_start:col_end]

StackExchange Code Review Q#132914, answer score: 29