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
if drop:
n_blocks_h, n_blocks_w = floor(height/stride_h), floor(width/stride_w)
else:
n_blocks_h, n_blocks_w = ceil(height/stride_h), 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
add patches_to_image
Be the first to comment
You can use [html][/html], [css][/css], [php][/php] and more to embed the code. Urls are automatically hyperlinked. Line breaks and paragraphs are automatically generated.