| from collections.abc import Sequence |
|
|
| import annotator.mmpkg.mmcv as mmcv |
| import numpy as np |
| import torch |
| from annotator.mmpkg.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`. |
| |
| 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 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.') |
|
|
|
|
| @PIPELINES.register_module() |
| class ToTensor(object): |
| """Convert some results to :obj:`torch.Tensor` by given keys. |
| |
| Args: |
| keys (Sequence[str]): Keys that need to be converted to Tensor. |
| """ |
|
|
| def __init__(self, keys): |
| self.keys = keys |
|
|
| def __call__(self, results): |
| """Call function to convert data in results to :obj:`torch.Tensor`. |
| |
| Args: |
| results (dict): Result dict contains the data to convert. |
| |
| Returns: |
| dict: The result dict contains the data converted |
| to :obj:`torch.Tensor`. |
| """ |
|
|
| 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 ImageToTensor(object): |
| """Convert image to :obj:`torch.Tensor` by given keys. |
| |
| The dimension order of input image is (H, W, C). The pipeline will convert |
| it to (C, H, W). If only 2 dimension (H, W) is given, the output would be |
| (1, H, W). |
| |
| Args: |
| keys (Sequence[str]): Key of images to be converted to Tensor. |
| """ |
|
|
| def __init__(self, keys): |
| self.keys = keys |
|
|
| def __call__(self, results): |
| """Call function to convert image in results to :obj:`torch.Tensor` and |
| transpose the channel order. |
| |
| Args: |
| results (dict): Result dict contains the image data to convert. |
| |
| Returns: |
| dict: The result dict contains the image converted |
| to :obj:`torch.Tensor` and transposed to (C, H, W) order. |
| """ |
|
|
| for key in self.keys: |
| img = results[key] |
| if len(img.shape) < 3: |
| img = np.expand_dims(img, -1) |
| results[key] = to_tensor(img.transpose(2, 0, 1)) |
| return results |
|
|
| def __repr__(self): |
| return self.__class__.__name__ + f'(keys={self.keys})' |
|
|
|
|
| @PIPELINES.register_module() |
| class Transpose(object): |
| """Transpose some results by given keys. |
| |
| Args: |
| keys (Sequence[str]): Keys of results to be transposed. |
| order (Sequence[int]): Order of transpose. |
| """ |
|
|
| def __init__(self, keys, order): |
| self.keys = keys |
| self.order = order |
|
|
| def __call__(self, results): |
| """Call function to convert image in results to :obj:`torch.Tensor` and |
| transpose the channel order. |
| |
| Args: |
| results (dict): Result dict contains the image data to convert. |
| |
| Returns: |
| dict: The result dict contains the image converted |
| to :obj:`torch.Tensor` and transposed to (C, H, W) order. |
| """ |
|
|
| 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 ToDataContainer(object): |
| """Convert results to :obj:`mmcv.DataContainer` by given fields. |
| |
| Args: |
| fields (Sequence[dict]): Each field is a dict like |
| ``dict(key='xxx', **kwargs)``. The ``key`` in result will |
| be converted to :obj:`mmcv.DataContainer` with ``**kwargs``. |
| Default: ``(dict(key='img', stack=True), |
| dict(key='gt_semantic_seg'))``. |
| """ |
|
|
| def __init__(self, |
| fields=(dict(key='img', |
| stack=True), dict(key='gt_semantic_seg'))): |
| self.fields = fields |
|
|
| def __call__(self, results): |
| """Call function to convert data in results to |
| :obj:`mmcv.DataContainer`. |
| |
| Args: |
| results (dict): Result dict contains the data to convert. |
| |
| Returns: |
| dict: The result dict contains the data converted to |
| :obj:`mmcv.DataContainer`. |
| """ |
|
|
| for field in self.fields: |
| field = field.copy() |
| key = field.pop('key') |
| results[key] = DC(results[key], **field) |
| return results |
|
|
| def __repr__(self): |
| return self.__class__.__name__ + f'(fields={self.fields})' |
|
|
|
|
| @PIPELINES.register_module() |
| class DefaultFormatBundle(object): |
| """Default formatting bundle. |
| |
| It simplifies the pipeline of formatting common fields, including "img" |
| and "gt_semantic_seg". These fields are formatted as follows. |
| |
| - img: (1)transpose, (2)to tensor, (3)to DataContainer (stack=True) |
| - gt_semantic_seg: (1)unsqueeze dim-0 (2)to tensor, |
| (3)to DataContainer (stack=True) |
| """ |
|
|
| def __call__(self, results): |
| """Call function to transform and format common fields in results. |
| |
| Args: |
| results (dict): Result dict contains the data to convert. |
| |
| Returns: |
| dict: The result dict contains the data that is formatted with |
| default bundle. |
| """ |
|
|
| if 'img' in results: |
| img = results['img'] |
| if len(img.shape) < 3: |
| img = np.expand_dims(img, -1) |
| img = np.ascontiguousarray(img.transpose(2, 0, 1)) |
| results['img'] = DC(to_tensor(img), stack=True) |
| if 'gt_semantic_seg' in results: |
| |
| results['gt_semantic_seg'] = DC( |
| to_tensor(results['gt_semantic_seg'][None, |
| ...].astype(np.int64)), |
| stack=True) |
| return results |
|
|
| def __repr__(self): |
| return self.__class__.__name__ |
|
|
|
|
| @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. Typically keys |
| is set to some subset of "img", "gt_semantic_seg". |
| |
| The "img_meta" item is always populated. The contents of the "img_meta" |
| dictionary depends on "meta_keys". By default this includes: |
| |
| - "img_shape": shape of the image input to the network as a tuple |
| (h, w, c). Note that images may be zero padded on the bottom/right |
| if the batch tensor is larger than this shape. |
| |
| - "scale_factor": a float indicating the preprocessing scale |
| |
| - "flip": a boolean indicating if image flip transform was used |
| |
| - "filename": path to the image file |
| |
| - "ori_shape": original shape of the image as a tuple (h, w, c) |
| |
| - "pad_shape": image shape after padding |
| |
| - "img_norm_cfg": a dict of normalization information: |
| - mean - per channel mean subtraction |
| - std - per channel std divisor |
| - to_rgb - bool indicating if bgr was converted to rgb |
| |
| 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[img_metas]``. |
| Default: ``('filename', 'ori_filename', 'ori_shape', 'img_shape', |
| 'pad_shape', 'scale_factor', 'flip', 'flip_direction', |
| 'img_norm_cfg')`` |
| """ |
|
|
| def __init__(self, |
| keys, |
| meta_keys=('filename', 'ori_filename', 'ori_shape', |
| 'img_shape', 'pad_shape', 'scale_factor', 'flip', |
| 'flip_direction', 'img_norm_cfg')): |
| self.keys = keys |
| self.meta_keys = meta_keys |
|
|
| def __call__(self, results): |
| """Call function to collect keys in results. The keys in ``meta_keys`` |
| will be converted to :obj:mmcv.DataContainer. |
| |
| Args: |
| results (dict): Result dict contains the data to collect. |
| |
| Returns: |
| dict: The result dict contains the following keys |
| - keys in``self.keys`` |
| - ``img_metas`` |
| """ |
|
|
| data = {} |
| img_meta = {} |
| for key in self.meta_keys: |
| img_meta[key] = results[key] |
| data['img_metas'] = DC(img_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})' |
|
|
