from numpy import min, max, all, array, subtract, pad
__all__ = ['pad_to_square']
[docs]def pad_to_square(arr: array, new_shape: tuple = None, **kwargs) -> array:
"""
Given a numpy array of unequal dimensions,
pad with zeros until it all dimensions have equal length
:param arr: numpy array
:param new_shape: desired shape
:param kwargs: are passed on to np.pad
:return:
"""
old_shape = arr.shape
if max(old_shape) == min(old_shape) and new_shape is None:
return arr
if new_shape is None:
new_shape = tuple([max(old_shape) for dim in old_shape])
elif type(new_shape) == int:
new_shape = tuple(new_shape for dim in old_shape)
assert all(array(new_shape) >= array(old_shape)), 'New dimensions must be larger than old dimensions'
padding = subtract(new_shape, old_shape)
left_pad = padding // 2
right_pad = padding - left_pad
cv = kwargs.pop('constant_values', 0)
padmode = kwargs.pop('mode', 'constant')
return pad(arr, tuple([(left_pad[i], right_pad[i]) for i in range(len(left_pad))]),
mode=padmode, constant_values=cv, **kwargs)