| |
| import warnings |
| from typing import List, Sequence, Union |
|
|
| import numpy as np |
| import torch |
|
|
| from .base_data_element import BaseDataElement |
|
|
|
|
| class PixelData(BaseDataElement): |
| """Data structure for pixel-level annotations or predictions. |
| |
| All data items in ``data_fields`` of ``PixelData`` meet the following |
| requirements: |
| |
| - They all have 3 dimensions in orders of channel, height, and width. |
| - They should have the same height and width. |
| |
| Examples: |
| >>> metainfo = dict( |
| ... img_id=random.randint(0, 100), |
| ... img_shape=(random.randint(400, 600), random.randint(400, 600))) |
| >>> image = np.random.randint(0, 255, (4, 20, 40)) |
| >>> featmap = torch.randint(0, 255, (10, 20, 40)) |
| >>> pixel_data = PixelData(metainfo=metainfo, |
| ... image=image, |
| ... featmap=featmap) |
| >>> print(pixel_data.shape) |
| (20, 40) |
| |
| >>> # slice |
| >>> slice_data = pixel_data[10:20, 20:40] |
| >>> assert slice_data.shape == (10, 20) |
| >>> slice_data = pixel_data[10, 20] |
| >>> assert slice_data.shape == (1, 1) |
| |
| >>> # set |
| >>> pixel_data.map3 = torch.randint(0, 255, (20, 40)) |
| >>> assert tuple(pixel_data.map3.shape) == (1, 20, 40) |
| >>> with self.assertRaises(AssertionError): |
| ... # The dimension must be 3 or 2 |
| ... pixel_data.map2 = torch.randint(0, 255, (1, 3, 20, 40)) |
| """ |
|
|
| def __setattr__(self, name: str, value: Union[torch.Tensor, np.ndarray]): |
| """Set attributes of ``PixelData``. |
| |
| If the dimension of value is 2 and its shape meet the demand, it |
| will automatically expand its channel-dimension. |
| |
| Args: |
| name (str): The key to access the value, stored in `PixelData`. |
| value (Union[torch.Tensor, np.ndarray]): The value to store in. |
| The type of value must be `torch.Tensor` or `np.ndarray`, |
| and its shape must meet the requirements of `PixelData`. |
| """ |
| if name in ('_metainfo_fields', '_data_fields'): |
| if not hasattr(self, name): |
| super().__setattr__(name, value) |
| else: |
| raise AttributeError(f'{name} has been used as a ' |
| 'private attribute, which is immutable.') |
|
|
| else: |
| assert isinstance(value, (torch.Tensor, np.ndarray)), \ |
| f'Can not set {type(value)}, only support' \ |
| f' {(torch.Tensor, np.ndarray)}' |
|
|
| if self.shape: |
| assert tuple(value.shape[-2:]) == self.shape, ( |
| 'The height and width of ' |
| f'values {tuple(value.shape[-2:])} is ' |
| 'not consistent with ' |
| 'the shape of this ' |
| ':obj:`PixelData` ' |
| f'{self.shape}') |
| assert value.ndim in [ |
| 2, 3 |
| ], f'The dim of value must be 2 or 3, but got {value.ndim}' |
| if value.ndim == 2: |
| value = value[None] |
| warnings.warn('The shape of value will convert from ' |
| f'{value.shape[-2:]} to {value.shape}') |
| super().__setattr__(name, value) |
|
|
| |
| def __getitem__(self, item: Sequence[Union[int, slice]]) -> 'PixelData': |
| """ |
| Args: |
| item (Sequence[Union[int, slice]]): Get the corresponding values |
| according to item. |
| |
| Returns: |
| :obj:`PixelData`: Corresponding values. |
| """ |
|
|
| new_data = self.__class__(metainfo=self.metainfo) |
| if isinstance(item, tuple): |
|
|
| assert len(item) == 2, 'Only support to slice height and width' |
| tmp_item: List[slice] = list() |
| for index, single_item in enumerate(item[::-1]): |
| if isinstance(single_item, int): |
| tmp_item.insert( |
| 0, slice(single_item, None, self.shape[-index - 1])) |
| elif isinstance(single_item, slice): |
| tmp_item.insert(0, single_item) |
| else: |
| raise TypeError( |
| 'The type of element in input must be int or slice, ' |
| f'but got {type(single_item)}') |
| tmp_item.insert(0, slice(None, None, None)) |
| item = tuple(tmp_item) |
| for k, v in self.items(): |
| setattr(new_data, k, v[item]) |
| else: |
| raise TypeError( |
| f'Unsupported type {type(item)} for slicing PixelData') |
| return new_data |
|
|
| @property |
| def shape(self): |
| """The shape of pixel data.""" |
| if len(self._data_fields) > 0: |
| return tuple(self.values()[0].shape[-2:]) |
| else: |
| return None |
|
|
| |
|
|
