fastest way to reshape 2D numpy array (gray image) into a 3D stacked array

I have a 2D image with the shape (M, N), and would like to reshape it into (M//m * N//n, m, n). That is, to stack small patches of images into a 3D array.

Currently, I used two for-loop to achieve that

import numpy as np
a = np.arange(16).reshape(4,4)
b = np.zeros( (4,2,2))

# a = array([[ 0,  1,  2,  3],
#       [ 4,  5,  6,  7],
#       [ 8,  9, 10, 11],
#       [12, 13, 14, 15]])


row_step = 2
col_step = 2
row_indices = np.arange(4, step = 2)
col_indices = np.arange(4, step = 2)

for i, row_idx in enumerate(row_indices):
    for j, col_idx in enumerate(col_indices):
        b[i*2+j, :,:] = a[row_idx:row_idx+row_step, col_idx:col_idx+col_step]

# b = array([[[ 0.,  1.],
#        [ 4.,  5.]],
#
#       [[ 2.,  3.],
#        [ 6.,  7.]],
#
#       [[ 8.,  9.],
#        [12., 13.]],
#
#       [[10., 11.],
#        [14., 15.]]])

Is there any other faster way to do this? Thanks a lot!

Answer

Use skimage.util.view_as_blocks:

from skimage.util.shape import view_as_blocks

out = view_as_blocks(a, (2,2))#.copy()

Output:

array([[[[ 0,  1],
         [ 4,  5]],

        [[ 2,  3],
         [ 6,  7]]],


       [[[ 8,  9],
         [12, 13]],

        [[10, 11],
         [14, 15]]]])

NB. Be aware that you have a view of the original object. If you want a copy, use view_as_blocks(a, (2,2)).copy()