How to crop same size image patches with different locations from a stack of images?

image-processing, multidimensional-array, numpy, python

Solution

Proper slicing is out of the question, because you need to access the underlying data on irregular intervals. You could get the crops with a single "fancy indexing" operation, but you'll need a (very) large indexing array. Therefor I think using a loop is easier and faster.

Compare the following two functions:

def fancy_indexing(imgs, centers, pw, ph):
    n = imgs.shape[0]
    img_i, RGB, x, y = np.ogrid[:n, :3, :pw, :ph]
    corners = centers - [pw//2, ph//2]
    x_i = x + corners[:,0,None,None,None]
    y_i = y + corners[:,1,None,None,None]
    return imgs[img_i, RGB, x_i, y_i]

def just_a_loop(imgs, centers, pw, ph):
    crops = np.empty(imgs.shape[:2]+(pw,ph), imgs.dtype)
    for i, (x,y) in enumerate(centers):
        crops[i] = imgs[i,:,x-pw//2:x+pw//2,y-ph//2:y+ph//2]
    return crops

Problem

Suppose I have an ndarray `imgs` of shape `( num_images, 3, width, height )` that stores a stack of `num_images` RGB images all of the same size. I would like to slice/crop from each image a patch of shape `( 3, pw, ph )` but the center location of the patch is different for each image and is given in `centers` array of shape `(num_images, 2)`. Is there a nice/pythonic way of slicing `imgs` to get `patches` (of shape `(num_images,3,pw,ph)`) each patch is centered around its corresponding `centers`? for simplicity it is safe to assume all patches fall within image boundaries.

Original source