| 139 | self.patch_iter = PatchIter(patch_size=patch_size, start_pos=start_pos, mode=mode, **pad_opts) |
| 140 | |
| 141 | def __call__( |
| 142 | self, data: Mapping[Hashable, NdarrayTensor] |
| 143 | ) -> Generator[tuple[Mapping[Hashable, NdarrayTensor], np.ndarray], None, None]: |
| 144 | d = dict(data) |
| 145 | original_spatial_shape = d[first(self.keys)].shape[1:] |
| 146 | |
| 147 | for patch in zip(*[self.patch_iter(d[key]) for key in self.keys]): |
| 148 | coords = patch[0][1] # use the coordinate of the first item |
| 149 | ret = {k: v[0] for k, v in zip(self.keys, patch)} |
| 150 | # fill in the extra keys with unmodified data |
| 151 | for k in set(d.keys()).difference(set(self.keys)): |
| 152 | ret[k] = deepcopy(d[k]) |
| 153 | # also store the `coordinate`, `spatial shape of original image`, `start position` in the dictionary |
| 154 | ret[self.coords_key] = coords |
| 155 | ret[self.original_spatial_shape_key] = original_spatial_shape |
| 156 | ret[self.start_pos_key] = self.patch_iter.start_pos |
| 157 | yield ret, coords |
| 158 | |
| 159 | |
| 160 | class GridPatchDataset(IterableDataset): |