|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
|
|
|
|
|
|
|
|
|
