A structure for storing masks and their related data in batched format. Implements basic filtering and concatenation.
| 14 | |
| 15 | |
| 16 | class MaskData: |
| 17 | """ |
| 18 | A structure for storing masks and their related data in batched format. |
| 19 | Implements basic filtering and concatenation. |
| 20 | """ |
| 21 | |
| 22 | def __init__(self, **kwargs) -> None: |
| 23 | for v in kwargs.values(): |
| 24 | assert isinstance( |
| 25 | v, (list, np.ndarray, torch.Tensor) |
| 26 | ), "MaskData only supports list, numpy arrays, and torch tensors." |
| 27 | self._stats = dict(**kwargs) |
| 28 | |
| 29 | def __setitem__(self, key: str, item: Any) -> None: |
| 30 | assert isinstance( |
| 31 | item, (list, np.ndarray, torch.Tensor) |
| 32 | ), "MaskData only supports list, numpy arrays, and torch tensors." |
| 33 | self._stats[key] = item |
| 34 | |
| 35 | def __delitem__(self, key: str) -> None: |
| 36 | del self._stats[key] |
| 37 | |
| 38 | def __getitem__(self, key: str) -> Any: |
| 39 | return self._stats[key] |
| 40 | |
| 41 | def items(self) -> ItemsView[str, Any]: |
| 42 | return self._stats.items() |
| 43 | |
| 44 | def filter(self, keep: torch.Tensor) -> None: |
| 45 | for k, v in self._stats.items(): |
| 46 | if v is None: |
| 47 | self._stats[k] = None |
| 48 | elif isinstance(v, torch.Tensor): |
| 49 | self._stats[k] = v[torch.as_tensor(keep, device=v.device)] |
| 50 | elif isinstance(v, np.ndarray): |
| 51 | self._stats[k] = v[keep.detach().cpu().numpy()] |
| 52 | elif isinstance(v, list) and keep.dtype == torch.bool: |
| 53 | self._stats[k] = [a for i, a in enumerate(v) if keep[i]] |
| 54 | elif isinstance(v, list): |
| 55 | self._stats[k] = [v[i] for i in keep] |
| 56 | else: |
| 57 | raise TypeError(f"MaskData key {k} has an unsupported type {type(v)}.") |
| 58 | |
| 59 | def cat(self, new_stats: "MaskData") -> None: |
| 60 | for k, v in new_stats.items(): |
| 61 | if k not in self._stats or self._stats[k] is None: |
| 62 | self._stats[k] = deepcopy(v) |
| 63 | elif isinstance(v, torch.Tensor): |
| 64 | self._stats[k] = torch.cat([self._stats[k], v], dim=0) |
| 65 | elif isinstance(v, np.ndarray): |
| 66 | self._stats[k] = np.concatenate([self._stats[k], v], axis=0) |
| 67 | elif isinstance(v, list): |
| 68 | self._stats[k] = self._stats[k] + deepcopy(v) |
| 69 | else: |
| 70 | raise TypeError(f"MaskData key {k} has an unsupported type {type(v)}.") |
| 71 | |
| 72 | def to_numpy(self) -> None: |
| 73 | for k, v in self._stats.items(): |
no outgoing calls
no test coverage detected