| from collections.abc import Sequence |
|
|
| import mmcv |
| import numpy as np |
| import torch |
| from mmcv.parallel import DataContainer as DC |
|
|
| from ..builder import PIPELINES |
|
|
|
|
| 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`. |
| """ |
| if isinstance(data, torch.Tensor): |
| return data |
| elif isinstance(data, np.ndarray): |
| return torch.from_numpy(data) |
| elif isinstance(data, Sequence) and not mmcv.is_str(data): |
| 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.' |
| 'Supported types are: `numpy.ndarray`, `torch.Tensor`, ' |
| '`Sequence`, `int` and `float`') |
|
|
|
|
| @PIPELINES.register_module() |
| class ToTensor(object): |
|
|
| def __init__(self, keys): |
| self.keys = keys |
|
|
| def __call__(self, results): |
| for key in self.keys: |
| results[key] = to_tensor(results[key]) |
| return results |
|
|
| def __repr__(self): |
| return self.__class__.__name__ + f'(keys={self.keys})' |
|
|
|
|
| @PIPELINES.register_module() |
| class Transpose(object): |
|
|
| def __init__(self, keys, order): |
| self.keys = keys |
| self.order = order |
|
|
| def __call__(self, results): |
| for key in self.keys: |
| results[key] = results[key].transpose(self.order) |
| return results |
|
|
| def __repr__(self): |
| return self.__class__.__name__ + \ |
| f'(keys={self.keys}, order={self.order})' |
|
|
|
|
| @PIPELINES.register_module() |
| class Collect(object): |
| """Collect data from the loader relevant to the specific task. |
| |
| This is usually the last stage of the data loader pipeline. |
| |
| Args: |
| keys (Sequence[str]): Keys of results to be collected in ``data``. |
| meta_keys (Sequence[str], optional): Meta keys to be converted to |
| ``mmcv.DataContainer`` and collected in ``data[motion_metas]``. |
| Default: ``('filename', 'ori_filename', |
| 'ori_shape', 'motion_shape', 'motion_mask')`` |
| |
| Returns: |
| dict: The result dict contains the following keys |
| - keys in``self.keys`` |
| - ``motion_metas`` if available |
| """ |
|
|
| def __init__(self, |
| keys, |
| meta_keys=('filename', 'ori_filename', 'ori_shape', |
| 'motion_shape', 'motion_mask')): |
| self.keys = keys |
| self.meta_keys = meta_keys |
|
|
| def __call__(self, results): |
| data = {} |
| motion_meta = {} |
| for key in self.meta_keys: |
| if key in results: |
| motion_meta[key] = results[key] |
| data['motion_metas'] = DC(motion_meta, cpu_only=True) |
| for key in self.keys: |
| data[key] = results[key] |
| return data |
|
|
| def __repr__(self): |
| return self.__class__.__name__ + \ |
| f'(keys={self.keys}, meta_keys={self.meta_keys})' |
|
|
|
|
| @PIPELINES.register_module() |
| class WrapFieldsToLists(object): |
| """Wrap fields of the data dictionary into lists for evaluation. |
| |
| This class can be used as a last step of a test or validation |
| pipeline for single image evaluation or inference. |
| |
| Example: |
| >>> test_pipeline = [ |
| >>> dict(type='LoadImageFromFile'), |
| >>> dict(type='Normalize', |
| mean=[123.675, 116.28, 103.53], |
| std=[58.395, 57.12, 57.375], |
| to_rgb=True), |
| >>> dict(type='ImageToTensor', keys=['img']), |
| >>> dict(type='Collect', keys=['img']), |
| >>> dict(type='WrapIntoLists') |
| >>> ] |
| """ |
|
|
| def __call__(self, results): |
| |
| for key, val in results.items(): |
| results[key] = [val] |
| return results |
|
|
| def __repr__(self): |
| return f'{self.__class__.__name__}()' |
|
|
