import numpy as np

def image_cut(image, size0, size1):
    """
    Takes an image and cuts it into smaller images of specified dimensions.
    Images must be cut before axis rolling.
    """
    dims = image.shape
    assert dims[0] >= size0
    assert dims[1] >= size1
    return np.array([image[size0 * i:size0 * (i+1), size1 * j:size1 * (j+1), :] \
        for i in range(dims[0] // size0) for j in range(dims[1] // size1)] + \
        [image[size0 * i:size0 * (i+1), dims[1]-size1:] \
            for i in range(dims[0] // size0) if dims[1] % size1 != 0] + \
        [image[dims[0]-size0:, size1 * j:size1 * (j+1)] \
            for j in range(dims[1] // size1) if dims[0] % size0 != 0] \
        )


if __name__ == '__main__':
    # test = np.random.rand(5,4,3)
    test = np.array([[
        [k + 4*j, k + 4*j] for k in range(4)
        ] for j in range(5)])
    print(test)
    print(image_cut(test, 2, 2))