Divide image in blocks

def split_in_blocks(array_2d, block_size, stride, pad=False, drop=False): ''' Split a 2d-array in blocks given a block_size and stride. The array is padded with zeros if block_size and stride do not fit array size. Keyword arguments: array_2d -- 2d numpy array-like block_size -- tuple (block_height, block_width) stride -- tuple (stride_height, stride_width) pad -- optinal pad [top, bottom, left, right] with zeros can be applied to original array before spliting drop -- drop blocks not fully contained inside array_2d Returns: 3d numpy array with all blocks (n_blocks, block_height, block_width) Authors: Arnaldo Gualberto and Arthur Araújo ''' from math import ceil, floor height, width = array_2d.shape block_h, block_w = block_size stride_h, stride_w = stride pad_array = None if isinstance(pad, list): top, bottom, left, right = pad pad_array = np.zeros((height + top + bottom, width + left + right), array_2d.dtype) pad_array[top:(top+height), left:(left+width)] = array_2d else: pad_array = array_2d height, width = pad_array.shape new_pad = None n_blocks_h = floor(height/stride_h) if drop else ceil(height/stride_h) n_blocks_w = floor(width/stride_w) if drop else ceil(width/stride_w) overlap_h = max(0, stride_h*(n_blocks_h - 1) + block_h - height) overlap_w = max(0, stride_w*(n_blocks_w - 1) + block_w - width) new_pad = np.zeros((height + overlap_h, width + overlap_w), pad_array.dtype) new_pad[:height, :width] = pad_array height, width = new_pad.shape block_index = 0 blocks = np.zeros((n_blocks_h*n_blocks_w, block_h, block_w), new_pad.dtype) for i in range(0, height-block_h+1, stride_h): for j in range(0, width-block_w+1, stride_w): blocks[block_index] = new_pad[i:i+block_h, j:j+block_w] block_index += 1 return blocks