Generate batch of patches and locations Args: patches: a tensor or list of tensors Yields: A batch of patches (torch.Tensor or MetaTensor), a sequence of location tuples, and the batch size
(
self, patches: Iterable[tuple[torch.Tensor, Sequence[int]]] | MetaTensor
)
| 191 | self.buffer_size = buffer_size |
| 192 | |
| 193 | def _batch_sampler( |
| 194 | self, patches: Iterable[tuple[torch.Tensor, Sequence[int]]] | MetaTensor |
| 195 | ) -> Iterator[tuple[torch.Tensor, Sequence, int]]: |
| 196 | """Generate batch of patches and locations |
| 197 | |
| 198 | Args: |
| 199 | patches: a tensor or list of tensors |
| 200 | |
| 201 | Yields: |
| 202 | A batch of patches (torch.Tensor or MetaTensor), a sequence of location tuples, and the batch size |
| 203 | """ |
| 204 | if isinstance(patches, MetaTensor): |
| 205 | total_size = len(patches) |
| 206 | for i in range(0, total_size, self.batch_size): |
| 207 | batch_size = min(self.batch_size, total_size - i) |
| 208 | yield patches[i : i + batch_size], patches[i : i + batch_size].meta[PatchKeys.LOCATION], batch_size # type: ignore |
| 209 | else: |
| 210 | buffer: Iterable | ThreadBuffer |
| 211 | if self.buffer_size > 0: |
| 212 | # Use multi-threading to sample patches with a buffer |
| 213 | buffer = ThreadBuffer(patches, buffer_size=self.buffer_size, timeout=0.1) |
| 214 | else: |
| 215 | buffer = patches |
| 216 | patch_batch: list[Any] = [None] * self.batch_size |
| 217 | location_batch: list[Any] = [None] * self.batch_size |
| 218 | idx_in_batch = 0 |
| 219 | for sample in buffer: |
| 220 | patch_batch[idx_in_batch] = sample[0] |
| 221 | location_batch[idx_in_batch] = sample[1] |
| 222 | idx_in_batch += 1 |
| 223 | if idx_in_batch == self.batch_size: |
| 224 | # concatenate batch of patches to create a tensor |
| 225 | yield torch.cat(patch_batch), location_batch, idx_in_batch |
| 226 | patch_batch = [None] * self.batch_size |
| 227 | location_batch = [None] * self.batch_size |
| 228 | idx_in_batch = 0 |
| 229 | if idx_in_batch > 0: |
| 230 | # concatenate batch of patches to create a tensor |
| 231 | yield torch.cat(patch_batch[:idx_in_batch]), location_batch, idx_in_batch |
| 232 | |
| 233 | def _ensure_tuple_outputs(self, outputs: Any) -> tuple: |
| 234 | if isinstance(outputs, dict): |
no test coverage detected