Divide image in blocks

import numpy as np 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 = (height - block_h + stride_h) / stride_h n_blocks_w = (width - block_w + stride_w) / stride_w n_blocks_h = floor(n_blocks_h) if drop else ceil(n_blocks_h) n_blocks_w = floor(n_blocks_w) if drop else ceil(n_blocks_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 def patches_to_image(patches, internal_size, image_size): from math import ceil n_patches, patch_height, patch_width = patches.shape img_height, img_width = image_size stride_h, stride_w = internal_size n_patches_h = ceil(img_height / stride_h) n_patches_w = ceil(img_width / stride_w) assert(n_patches == n_patches_h * n_patches_w) center_x, center_y = patch_width // 2, patch_height // 2 top, left = center_y - (stride_h // 2), center_x - (stride_w // 2) bottom, right = top + stride_h, left + stride_w image = np.zeros((n_patches_h*stride_h, n_patches_w*stride_w)) for patch_index, p in enumerate(patches): row = (patch_index // n_patches_w) * stride_h col = (patch_index % n_patches_w) * stride_w image[row:row + stride_h, col:col + stride_w] = p[top:bottom, left:right] return image[0:img_height, 0:img_width]