in easycv/core/utils/array_converter.py [0:0]
def array_converter(to_torch=True,
apply_to=tuple(),
template_arg_name_=None,
recover=True):
"""Wrapper function for data-type agnostic processing.
First converts input arrays to PyTorch tensors or NumPy ndarrays
for middle calculation, then convert output to original data-type if
`recover=True`.
Args:
to_torch (Bool, optional): Whether convert to PyTorch tensors
for middle calculation. Defaults to True.
apply_to (tuple[str], optional): The arguments to which we apply
data-type conversion. Defaults to an empty tuple.
template_arg_name_ (str, optional): Argument serving as the template (
return arrays should have the same dtype and device
as the template). Defaults to None. If None, we will use the
first argument in `apply_to` as the template argument.
recover (Bool, optional): Whether or not recover the wrapped function
outputs to the `template_arg_name_` type. Defaults to True.
Raises:
ValueError: When template_arg_name_ is not among all args, or
when apply_to contains an arg which is not among all args,
a ValueError will be raised. When the template argument or
an argument to convert is a list or tuple, and cannot be
converted to a NumPy array, a ValueError will be raised.
TypeError: When the type of the template argument or
an argument to convert does not belong to the above range,
or the contents of such an list-or-tuple-type argument
do not share the same data type, a TypeError is raised.
Returns:
(function): wrapped function.
Example:
>>> import torch
>>> import numpy as np
>>>
>>> # Use torch addition for a + b,
>>> # and convert return values to the type of a
>>> @array_converter(apply_to=('a', 'b'))
>>> def simple_add(a, b):
>>> return a + b
>>>
>>> a = np.array([1.1])
>>> b = np.array([2.2])
>>> simple_add(a, b)
>>>
>>> # Use numpy addition for a + b,
>>> # and convert return values to the type of b
>>> @array_converter(to_torch=False, apply_to=('a', 'b'),
>>> template_arg_name_='b')
>>> def simple_add(a, b):
>>> return a + b
>>>
>>> simple_add()
>>>
>>> # Use torch funcs for floor(a) if flag=True else ceil(a),
>>> # and return the torch tensor
>>> @array_converter(apply_to=('a',), recover=False)
>>> def floor_or_ceil(a, flag=True):
>>> return torch.floor(a) if flag else torch.ceil(a)
>>>
>>> floor_or_ceil(a, flag=False)
"""
def array_converter_wrapper(func):
"""Outer wrapper for the function."""
@functools.wraps(func)
def new_func(*args, **kwargs):
"""Inner wrapper for the arguments."""
if len(apply_to) == 0:
return func(*args, **kwargs)
func_name = func.__name__
arg_spec = getfullargspec(func)
arg_names = arg_spec.args
arg_num = len(arg_names)
default_arg_values = arg_spec.defaults
if default_arg_values is None:
default_arg_values = []
no_default_arg_num = len(arg_names) - len(default_arg_values)
kwonly_arg_names = arg_spec.kwonlyargs
kwonly_default_arg_values = arg_spec.kwonlydefaults
if kwonly_default_arg_values is None:
kwonly_default_arg_values = {}
all_arg_names = arg_names + kwonly_arg_names
# in case there are args in the form of *args
if len(args) > arg_num:
named_args = args[:arg_num]
nameless_args = args[arg_num:]
else:
named_args = args
nameless_args = []
# template argument data type is used for all array-like arguments
if template_arg_name_ is None:
template_arg_name = apply_to[0]
else:
template_arg_name = template_arg_name_
if template_arg_name not in all_arg_names:
raise ValueError(f'{template_arg_name} is not among the '
f'argument list of function {func_name}')
# inspect apply_to
for arg_to_apply in apply_to:
if arg_to_apply not in all_arg_names:
raise ValueError(f'{arg_to_apply} is not '
f'an argument of {func_name}')
new_args = []
new_kwargs = {}
converter = ArrayConverter()
target_type = torch.Tensor if to_torch else np.ndarray
# non-keyword arguments
for i, arg_value in enumerate(named_args):
if arg_names[i] in apply_to:
new_args.append(
converter.convert(
input_array=arg_value, target_type=target_type))
else:
new_args.append(arg_value)
if arg_names[i] == template_arg_name:
template_arg_value = arg_value
kwonly_default_arg_values.update(kwargs)
kwargs = kwonly_default_arg_values
# keyword arguments and non-keyword arguments using default value
for i in range(len(named_args), len(all_arg_names)):
arg_name = all_arg_names[i]
if arg_name in kwargs:
if arg_name in apply_to:
new_kwargs[arg_name] = converter.convert(
input_array=kwargs[arg_name],
target_type=target_type)
else:
new_kwargs[arg_name] = kwargs[arg_name]
else:
default_value = default_arg_values[i - no_default_arg_num]
if arg_name in apply_to:
new_kwargs[arg_name] = converter.convert(
input_array=default_value, target_type=target_type)
else:
new_kwargs[arg_name] = default_value
if arg_name == template_arg_name:
template_arg_value = kwargs[arg_name]
# add nameless args provided by *args (if exists)
new_args += nameless_args
return_values = func(*new_args, **new_kwargs)
converter.set_template(template_arg_value)
def recursive_recover(input_data):
if isinstance(input_data, (tuple, list)):
new_data = []
for item in input_data:
new_data.append(recursive_recover(item))
return tuple(new_data) if isinstance(input_data,
tuple) else new_data
elif isinstance(input_data, dict):
new_data = {}
for k, v in input_data.items():
new_data[k] = recursive_recover(v)
return new_data
elif isinstance(input_data, (torch.Tensor, np.ndarray)):
return converter.recover(input_data)
else:
return input_data
if recover:
return recursive_recover(return_values)
else:
return return_values
return new_func
return array_converter_wrapper