Args: item (str, int, list, :obj:`slice`, :obj:`numpy.ndarray`, :obj:`torch.LongTensor`, :obj:`torch.BoolTensor`): Get the corresponding values according to item. Returns: :obj:`PointData`: Corresponding values.
(self, item: IndexType)
| 96 | __setitem__ = __setattr__ |
| 97 | |
| 98 | def __getitem__(self, item: IndexType) -> 'PointData': |
| 99 | """ |
| 100 | Args: |
| 101 | item (str, int, list, :obj:`slice`, :obj:`numpy.ndarray`, |
| 102 | :obj:`torch.LongTensor`, :obj:`torch.BoolTensor`): |
| 103 | Get the corresponding values according to item. |
| 104 | |
| 105 | Returns: |
| 106 | :obj:`PointData`: Corresponding values. |
| 107 | """ |
| 108 | if isinstance(item, list): |
| 109 | item = np.array(item) |
| 110 | if isinstance(item, np.ndarray): |
| 111 | # The default int type of numpy is platform dependent, int32 for |
| 112 | # windows and int64 for linux. `torch.Tensor` requires the index |
| 113 | # should be int64, therefore we simply convert it to int64 here. |
| 114 | # Mode details in https://github.com/numpy/numpy/issues/9464 |
| 115 | item = item.astype(np.int64) if item.dtype == np.int32 else item |
| 116 | item = torch.from_numpy(item) |
| 117 | assert isinstance( |
| 118 | item, (str, slice, int, torch.LongTensor, torch.cuda.LongTensor, |
| 119 | torch.BoolTensor, torch.cuda.BoolTensor)) |
| 120 | |
| 121 | if isinstance(item, str): |
| 122 | return getattr(self, item) |
| 123 | |
| 124 | if isinstance(item, int): |
| 125 | if item >= len(self) or item < -len(self): # type: ignore |
| 126 | raise IndexError(f'Index {item} out of range!') |
| 127 | else: |
| 128 | # keep the dimension |
| 129 | item = slice(item, None, len(self)) |
| 130 | |
| 131 | new_data = self.__class__(metainfo=self.metainfo) |
| 132 | if isinstance(item, torch.Tensor): |
| 133 | assert item.dim() == 1, 'Only support to get the' \ |
| 134 | ' values along the first dimension.' |
| 135 | if isinstance(item, (torch.BoolTensor, torch.cuda.BoolTensor)): |
| 136 | assert len(item) == len(self), 'The shape of the ' \ |
| 137 | 'input(BoolTensor) ' \ |
| 138 | f'{len(item)} ' \ |
| 139 | 'does not match the shape ' \ |
| 140 | 'of the indexed tensor ' \ |
| 141 | 'in results_field ' \ |
| 142 | f'{len(self)} at ' \ |
| 143 | 'first dimension.' |
| 144 | |
| 145 | for k, v in self.items(): |
| 146 | if isinstance(v, torch.Tensor): |
| 147 | new_data[k] = v[item] |
| 148 | elif isinstance(v, np.ndarray): |
| 149 | new_data[k] = v[item.cpu().numpy()] |
| 150 | elif isinstance( |
| 151 | v, (str, list, tuple)) or (hasattr(v, '__getitem__') |
| 152 | and hasattr(v, 'cat')): |
| 153 | # convert to indexes from BoolTensor |
| 154 | if isinstance(item, |
| 155 | (torch.BoolTensor, torch.cuda.BoolTensor)): |