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
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.