Args: batch: batch of data to pad-collate
(self, batch: Any)
| 70 | self.kwargs = kwargs |
| 71 | |
| 72 | def __call__(self, batch: Any): |
| 73 | """ |
| 74 | Args: |
| 75 | batch: batch of data to pad-collate |
| 76 | """ |
| 77 | # data is either list of dicts or list of lists |
| 78 | is_list_of_dicts = isinstance(batch[0], dict) |
| 79 | # loop over items inside of each element in a batch |
| 80 | batch_item = tuple(batch[0].keys()) if is_list_of_dicts else range(len(batch[0])) |
| 81 | for key_or_idx in batch_item: |
| 82 | # calculate max size of each dimension |
| 83 | max_shapes = [] |
| 84 | for elem in batch: |
| 85 | if not isinstance(elem[key_or_idx], (torch.Tensor, np.ndarray)): |
| 86 | break |
| 87 | max_shapes.append(elem[key_or_idx].shape[1:]) |
| 88 | # len > 0 if objects were arrays, else skip as no padding to be done |
| 89 | if not max_shapes: |
| 90 | continue |
| 91 | max_shape = np.array(max_shapes).max(axis=0) |
| 92 | # If all same size, skip |
| 93 | if np.all(np.array(max_shapes).min(axis=0) == max_shape): |
| 94 | continue |
| 95 | |
| 96 | # Use `SpatialPad` to match sizes, Default params are central padding, padding with 0's |
| 97 | padder = SpatialPad(spatial_size=max_shape, method=self.method, mode=self.mode, **self.kwargs) |
| 98 | for idx, batch_i in enumerate(batch): |
| 99 | orig_size = batch_i[key_or_idx].shape[1:] |
| 100 | padded = padder(batch_i[key_or_idx]) |
| 101 | batch = replace_element(padded, batch, idx, key_or_idx) |
| 102 | |
| 103 | # If we have a dictionary of data, append to list |
| 104 | # padder transform info is re-added with self.push_transform to ensure one info dict per transform. |
| 105 | if is_list_of_dicts: |
| 106 | self.push_transform( |
| 107 | batch[idx], |
| 108 | key_or_idx, |
| 109 | orig_size=orig_size, |
| 110 | extra_info=self.pop_transform(batch[idx], key_or_idx, check=False), |
| 111 | ) |
| 112 | |
| 113 | # After padding, use default list collator |
| 114 | return list_data_collate(batch) |
| 115 | |
| 116 | @staticmethod |
| 117 | def inverse(data: dict) -> dict[Hashable, np.ndarray]: |
nothing calls this directly
no test coverage detected