Aggregate values for merging. Args: values: a tensor of shape BCHW[D], representing the values of inference output. location: a tuple/list giving the top left location of the patch in the original image.
(self, values: torch.Tensor, location: Sequence[int])
| 428 | self.lock = nullcontext() |
| 429 | |
| 430 | def aggregate(self, values: torch.Tensor, location: Sequence[int]) -> None: |
| 431 | """ |
| 432 | Aggregate values for merging. |
| 433 | |
| 434 | Args: |
| 435 | values: a tensor of shape BCHW[D], representing the values of inference output. |
| 436 | location: a tuple/list giving the top left location of the patch in the original image. |
| 437 | """ |
| 438 | if self.is_finalized: |
| 439 | raise ValueError("`ZarrAvgMerger` is already finalized. Please instantiate a new object to aggregate.") |
| 440 | patch_size = values.shape[2:] |
| 441 | map_slice = tuple(slice(loc, loc + size) for loc, size in zip(location, patch_size)) |
| 442 | map_slice = ensure_tuple_size(map_slice, values.ndim, pad_val=slice(None), pad_from_start=True) |
| 443 | with self.lock: |
| 444 | self.values[map_slice] += values.numpy() |
| 445 | self.counts[map_slice] += 1 # type: ignore[operator] |
| 446 | |
| 447 | def finalize(self) -> zarr.Array: |
| 448 | """ |