| |
| from numbers import Number |
| from typing import Sequence, Union |
|
|
| import mmengine |
| import numpy as np |
| import torch |
| from mmengine.structures import BaseDataElement, LabelData |
|
|
|
|
| def format_label(value: Union[torch.Tensor, np.ndarray, Sequence, int], |
| num_classes: int = None) -> LabelData: |
| """Convert label of various python types to :obj:`mmengine.LabelData`. |
| |
| Supported types are: :class:`numpy.ndarray`, :class:`torch.Tensor`, |
| :class:`Sequence`, :class:`int`. |
| |
| Args: |
| value (torch.Tensor | numpy.ndarray | Sequence | int): Label value. |
| num_classes (int, optional): The number of classes. If not None, set |
| it to the metainfo. Defaults to None. |
| |
| Returns: |
| :obj:`mmengine.LabelData`: The foramtted label data. |
| """ |
|
|
| |
| if isinstance(value, (torch.Tensor, np.ndarray)) and value.ndim == 0: |
| value = int(value.item()) |
|
|
| if isinstance(value, np.ndarray): |
| value = torch.from_numpy(value) |
| elif isinstance(value, Sequence) and not mmengine.utils.is_str(value): |
| value = torch.tensor(value) |
| elif isinstance(value, int): |
| value = torch.LongTensor([value]) |
| elif not isinstance(value, torch.Tensor): |
| raise TypeError(f'Type {type(value)} is not an available label type.') |
|
|
| metainfo = {} |
| if num_classes is not None: |
| metainfo['num_classes'] = num_classes |
| if value.max() >= num_classes: |
| raise ValueError(f'The label data ({value}) should not ' |
| f'exceed num_classes ({num_classes}).') |
| label = LabelData(label=value, metainfo=metainfo) |
| return label |
|
|
|
|
| class ReIDDataSample(BaseDataElement): |
| """A data structure interface of ReID task. |
| |
| It's used as interfaces between different components. |
| |
| Meta field: |
| img_shape (Tuple): The shape of the corresponding input image. |
| Used for visualization. |
| ori_shape (Tuple): The original shape of the corresponding image. |
| Used for visualization. |
| num_classes (int): The number of all categories. |
| Used for label format conversion. |
| |
| Data field: |
| gt_label (LabelData): The ground truth label. |
| pred_label (LabelData): The predicted label. |
| scores (torch.Tensor): The outputs of model. |
| """ |
|
|
| @property |
| def gt_label(self): |
| return self._gt_label |
|
|
| @gt_label.setter |
| def gt_label(self, value: LabelData): |
| self.set_field(value, '_gt_label', dtype=LabelData) |
|
|
| @gt_label.deleter |
| def gt_label(self): |
| del self._gt_label |
|
|
| def set_gt_label( |
| self, value: Union[np.ndarray, torch.Tensor, Sequence[Number], Number] |
| ) -> 'ReIDDataSample': |
| """Set label of ``gt_label``.""" |
| label = format_label(value, self.get('num_classes')) |
| if 'gt_label' in self: |
| self.gt_label.label = label.label |
| else: |
| self.gt_label = label |
| return self |
|
|
| def set_gt_score(self, value: torch.Tensor) -> 'ReIDDataSample': |
| """Set score of ``gt_label``.""" |
| assert isinstance(value, torch.Tensor), \ |
| f'The value should be a torch.Tensor but got {type(value)}.' |
| assert value.ndim == 1, \ |
| f'The dims of value should be 1, but got {value.ndim}.' |
|
|
| if 'num_classes' in self: |
| assert value.size(0) == self.num_classes, \ |
| f"The length of value ({value.size(0)}) doesn't "\ |
| f'match the num_classes ({self.num_classes}).' |
| metainfo = {'num_classes': self.num_classes} |
| else: |
| metainfo = {'num_classes': value.size(0)} |
|
|
| if 'gt_label' in self: |
| self.gt_label.score = value |
| else: |
| self.gt_label = LabelData(score=value, metainfo=metainfo) |
| return self |
|
|
| @property |
| def pred_feature(self): |
| return self._pred_feature |
|
|
| @pred_feature.setter |
| def pred_feature(self, value: torch.Tensor): |
| self.set_field(value, '_pred_feature', dtype=torch.Tensor) |
|
|
| @pred_feature.deleter |
| def pred_feature(self): |
| del self._pred_feature |
|
|
