232 lines
6.9 KiB
Python
232 lines
6.9 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
||
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||
|
||
import time
|
||
from collections.abc import Sequence
|
||
from typing import Mapping
|
||
|
||
import numpy as np
|
||
import torch
|
||
|
||
from modelscope.utils.registry import default_group
|
||
from .builder import PREPROCESSORS, build_preprocessor
|
||
|
||
|
||
@PREPROCESSORS.register_module()
|
||
class Compose(object):
|
||
"""Compose a data pipeline with a sequence of transforms.
|
||
Args:
|
||
transforms (list[dict | callable]):
|
||
Either config dicts of transforms or transform objects.
|
||
profiling (bool, optional): If set True, will profile and
|
||
print preprocess time for each step.
|
||
"""
|
||
|
||
def __init__(self, transforms, field_name=None, profiling=False):
|
||
assert isinstance(transforms, Sequence)
|
||
self.profiling = profiling
|
||
self.transforms = []
|
||
self.field_name = field_name
|
||
for transform in transforms:
|
||
if isinstance(transform, dict):
|
||
if self.field_name is None:
|
||
transform = build_preprocessor(transform, default_group)
|
||
else:
|
||
# if not found key in field_name, try field_name=None(default_group)
|
||
try:
|
||
transform = build_preprocessor(transform, field_name)
|
||
except KeyError:
|
||
transform = build_preprocessor(transform,
|
||
default_group)
|
||
elif callable(transform):
|
||
pass
|
||
else:
|
||
raise TypeError('transform must be callable or a dict, but got'
|
||
f' {type(transform)}')
|
||
self.transforms.append(transform)
|
||
|
||
def __call__(self, data):
|
||
for t in self.transforms:
|
||
if self.profiling:
|
||
start = time.time()
|
||
|
||
data = t(data)
|
||
|
||
if self.profiling:
|
||
print(f'{t} time {time.time()-start}')
|
||
|
||
if data is None:
|
||
return None
|
||
return data
|
||
|
||
def __repr__(self):
|
||
format_string = self.__class__.__name__ + '('
|
||
for t in self.transforms:
|
||
format_string += f'\n {t}'
|
||
format_string += '\n)'
|
||
return format_string
|
||
|
||
|
||
def to_tensor(data):
|
||
"""Convert objects of various python types to :obj:`torch.Tensor`.
|
||
|
||
Supported types are: :class:`numpy.ndarray`, :class:`torch.Tensor`,
|
||
:class:`Sequence`, :class:`int` and :class:`float`.
|
||
|
||
Args:
|
||
data (torch.Tensor | numpy.ndarray | Sequence | int | float): Data to
|
||
be converted.
|
||
"""
|
||
|
||
if isinstance(data, torch.Tensor):
|
||
return data
|
||
elif isinstance(data, np.ndarray):
|
||
return torch.from_numpy(data)
|
||
elif isinstance(data, Sequence) and not isinstance(data, str):
|
||
return torch.tensor(data)
|
||
elif isinstance(data, int):
|
||
return torch.LongTensor([data])
|
||
elif isinstance(data, float):
|
||
return torch.FloatTensor([data])
|
||
else:
|
||
raise TypeError(f'type {type(data)} cannot be converted to tensor.')
|
||
|
||
|
||
@PREPROCESSORS.register_module()
|
||
class ToTensor(object):
|
||
"""Convert target object to tensor.
|
||
|
||
Args:
|
||
keys (Sequence[str]): Key of data to be converted to Tensor.
|
||
Only valid when data is type of `Mapping`. If `keys` is None,
|
||
all values of keys will be converted to tensor by default.
|
||
"""
|
||
|
||
def __init__(self, keys=None):
|
||
self.keys = keys
|
||
|
||
def __call__(self, data):
|
||
if isinstance(data, Mapping):
|
||
if self.keys is None:
|
||
self.keys = list(data.keys())
|
||
|
||
for key in self.keys:
|
||
if key in data:
|
||
data[key] = to_tensor(data[key])
|
||
else:
|
||
data = to_tensor(data)
|
||
|
||
return data
|
||
|
||
def __repr__(self):
|
||
return self.__class__.__name__ + f'(keys={self.keys})'
|
||
|
||
|
||
@PREPROCESSORS.register_module()
|
||
class Filter(object):
|
||
"""This is usually the last stage of the dataloader transform.
|
||
Only data of reserved keys will be kept and passed directly to the model, others will be removed.
|
||
|
||
Args:
|
||
keys (Sequence[str]): Keys of data to be reserved, others will be removed.
|
||
"""
|
||
|
||
def __init__(self, reserved_keys):
|
||
self.reserved_keys = reserved_keys
|
||
|
||
def __call__(self, data):
|
||
assert isinstance(data, Mapping)
|
||
|
||
reserved_data = {}
|
||
for key in self.reserved_keys:
|
||
if key in data:
|
||
reserved_data[key] = data[key]
|
||
|
||
return reserved_data
|
||
|
||
def __repr__(self):
|
||
return self.__class__.__name__ + f'(keys={self.reserved_keys})'
|
||
|
||
|
||
def to_numpy(data):
|
||
"""Convert objects of various python types to `numpy.ndarray`.
|
||
|
||
Args:
|
||
data (torch.Tensor | numpy.ndarray | Sequence | int | float): Data to
|
||
be converted.
|
||
"""
|
||
|
||
if isinstance(data, torch.Tensor):
|
||
return data.numpy()
|
||
elif isinstance(data, np.ndarray):
|
||
return data
|
||
elif isinstance(data, Sequence) and not isinstance(data, str):
|
||
return np.asarray(data)
|
||
elif isinstance(data, int):
|
||
return np.asarray(data, dtype=np.int64)
|
||
elif isinstance(data, float):
|
||
return np.asarray(data, dtype=np.float64)
|
||
else:
|
||
raise TypeError(f'type {type(data)} cannot be converted to tensor.')
|
||
|
||
|
||
@PREPROCESSORS.register_module()
|
||
class ToNumpy(object):
|
||
"""Convert target object to numpy.ndarray.
|
||
|
||
Args:
|
||
keys (Sequence[str]): Key of data to be converted to numpy.ndarray.
|
||
Only valid when data is type of `Mapping`. If `keys` is None,
|
||
all values of keys will be converted to numpy.ndarray by default.
|
||
"""
|
||
|
||
def __init__(self, keys=None):
|
||
self.keys = keys
|
||
|
||
def __call__(self, data):
|
||
if isinstance(data, Mapping):
|
||
if self.keys is None:
|
||
self.keys = list(data.keys())
|
||
|
||
for key in self.keys:
|
||
if key in data:
|
||
data[key] = to_numpy(data[key])
|
||
else:
|
||
data = to_numpy(data)
|
||
|
||
return data
|
||
|
||
def __repr__(self):
|
||
return self.__class__.__name__ + f'(keys={self.keys})'
|
||
|
||
|
||
@PREPROCESSORS.register_module()
|
||
class Rename(object):
|
||
"""Change the name of the input keys to output keys, respectively.
|
||
"""
|
||
|
||
def __init__(self, input_keys=[], output_keys=[]):
|
||
self.input_keys = input_keys
|
||
self.output_keys = output_keys
|
||
|
||
def __call__(self, data):
|
||
if isinstance(data, Mapping):
|
||
for in_key, out_key in zip(self.input_keys, self.output_keys):
|
||
if in_key in data and out_key not in data:
|
||
data[out_key] = data[in_key]
|
||
data.pop(in_key)
|
||
return data
|
||
|
||
def __repr__(self):
|
||
return self.__class__.__name__ + f'(keys={self.keys})'
|
||
|
||
|
||
@PREPROCESSORS.register_module()
|
||
class Identity(object):
|
||
|
||
def __init__(self):
|
||
pass
|
||
|
||
def __call__(self, item):
|
||
return item
|