Extracting 2d patches from 3d arrays

Question

scikit-learn's extract_patches_2d can be used to reshape a 2D image into a collection of patches. extract_patches is the generalized form that uses NumPy's as_strided .

import numpy as np
from sklearn.feature_extraction import image

ex = np.arange(3 * 3).reshape(3, 3)
image.extract_patches_2d(ex, patch_size=(2, 2))
[[[0 1]
  [3 4]]

 [[1 2]
  [4 5]]

 [[3 4]
  [6 7]]

 [[4 5]
  [7 8]]]

I have a three-dimensional array a and would like to extract 2d patches from each "innermost" 2d array, then find the (axis-agnostic) mean of each 2d patch.

a = np.arange(2 * 3 * 3).reshape(2, 3, 3)

In this case, I effectively want to first call extract_patches_2d on each (3, 3) inner array.

patches = np.array([image.extract_patches_2d(i, patch_size=(2, 2)) for i in a])

and then find the mean of each innermost 2d array (each patch):

means = patches.reshape(*patches.shape[:-2], -1).mean(axis=-1)
print(means)
[[  2.   3.   5.   6.]
 [ 11.  12.  14.  15.]]

How can I vectorize this and get rid of the for-loop above? The important thing here is that the size of the first dimension of means is equal to the size of a 's first dimension.

Answer 1

You can use scikit-image as view_as_windows to get those patches as a view into input array -

from skimage.util.shape import view_as_windows

size = 2 # patch size
patches = view_as_windows(a, (1,size,size))[...,0,:,:]

This gives us a 5D array as patches , on which we can use mean reduction along the last two axes for a 3D output -

out = patches.mean((-2,-1))

If the final output is needed as a 2D one, reshape to merge last two axes -

out.reshape(a.shape[0],-1)

This can also utilize sklearn 's extract_patches :

def inner_means(arr_3d, patch_size):
    """Axis-agnostic mean of each 2d patch.

    Maintains the first dimension of `arr_3d`.

    patch_size: tuple
        Same syntax as the parameter passed to extract_patches_2d
    """
    shape = (1,) + patch_size
    patches = image.extract_patches(arr_3d, shape)[..., 0, :, :].mean((-2, -1))
    return patches.reshape(*patches.shape[:-2], -1)


a = np.arange(2 * 3 * 3).reshape(2, 3, 3)
    print(inner_means(a, patch_size=(2, 2)))

[[  2.   3.   5.   6.]
 [ 11.  12.  14.  15.]]

---

Alternatively, to directly get to the blocky average values, we can use one of those convolution tools from Scipy. So with fftconvolve -

from scipy.signal import fftconvolve

out = fftconvolve(a, np.ones((1,size,size)),mode='valid')/size**2

Or use scipy.signal.convolve or scipy.ndimage.filters.uniform_filter there without the division.