WSI sliding window splitter (#6107)

Fixes #5871 

### Description
This PR:
- implements a `WSISlidingWindowSplitter ` a splitter for whole slide
imaging,
- refactors `SlidingWindowSplitter` to make it inheritable for wsi
version (backward incompatible arguments),
- adds new feature to `AvgMerger` to crop the output (backward
compatible), and
- updates `PatchInference` to set cropping shape for new feature in
`AvgMerger` (backward compatible)

### Types of changes
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [x] Non-breaking change (fix or new feature that would not break
existing functionality).
- [x] New tests added to cover the changes.
- [x] Integration tests passed locally by running `./runtests.sh -f -u
--net --coverage`.
- [x] Quick tests passed locally by running `./runtests.sh --quick
--unittests --disttests`.
- [x] In-line docstrings updated.
- [x] Documentation updated, tested `make html` command in the `docs/`
folder.

---------

Signed-off-by: Behrooz <3968947+drbeh@users.noreply.github.com>
Signed-off-by: monai-bot <monai.miccai2019@gmail.com>
Signed-off-by: Dr. Behrooz Hashemian <3968947+drbeh@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: monai-bot <monai.miccai2019@gmail.com>
Co-authored-by: Eric Kerfoot <17726042+ericspod@users.noreply.github.com>

Author: 4 people

docs/source/inferers.rst

@@ -57,6 +57,13 @@ Splitters
     :members:
     :special-members: __call__
 
+`WSISlidingWindowSplitter`
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+.. autoclass:: WSISlidingWindowSplitter
+    :members:
+    :special-members: __call__
+
+
 Mergers
 -------
 .. currentmodule:: monai.inferers

monai/data/utils.py

@@ -209,7 +209,7 @@ def iter_patch_position(
     image_size: Sequence[int],
     patch_size: Sequence[int] | int | np.ndarray,
     start_pos: Sequence[int] = (),
-    overlap: Sequence[float] | float = 0.0,
+    overlap: Sequence[float] | float | Sequence[int] | int = 0.0,
     padded: bool = False,
 ):
     """
@@ -221,8 +221,10 @@ def iter_patch_position(
         image_size: dimensions of array to iterate over
         patch_size: size of patches to generate slices for, 0 or None selects whole dimension
         start_pos: starting position in the array, default is 0 for each dimension
-        overlap: the amount of overlap of neighboring patches in each dimension (a value between 0.0 and 1.0).
-            If only one float number is given, it will be applied to all dimensions. Defaults to 0.0.
+        overlap: the amount of overlap of neighboring patches in each dimension.
+            Either a float or list of floats between 0.0 and 1.0 to define relative overlap to patch size, or
+            an int or list of ints to define number of pixels for overlap.
+            If only one float/int number is given, it will be applied to all dimensions. Defaults to 0.0.
         padded: if the image is padded so the patches can go beyond the borders. Defaults to False.
 
     Yields:
@@ -236,7 +238,10 @@ def iter_patch_position(
     overlap = ensure_tuple_rep(overlap, ndim)
 
     # calculate steps, which depends on the amount of overlap
-    steps = tuple(round(p * (1.0 - o)) for p, o in zip(patch_size_, overlap))
+    if isinstance(overlap[0], float):
+        steps = tuple(round(p * (1.0 - o)) for p, o in zip(patch_size_, overlap))
+    else:
+        steps = tuple(p - o for p, o in zip(patch_size_, overlap))
 
     # calculate the last starting location (depending on the padding)
     end_pos = image_size if padded else tuple(s - round(p) + 1 for s, p in zip(image_size, patch_size_))

monai/inferers/__init__.py

@@ -21,5 +21,5 @@
     SlidingWindowInfererAdapt,
 )
 from .merger import AvgMerger, Merger
-from .splitter import SlidingWindowSplitter, Splitter
+from .splitter import SlidingWindowSplitter, Splitter, WSISlidingWindowSplitter
 from .utils import sliding_window_inference

monai/inferers/inferer.py

@@ -87,6 +87,8 @@ class PatchInferer(Inferer):
     Args:
         splitter: a `Splitter` object that split the inputs into patches. Defaults to None.
             If not provided or None, the inputs are considered to be already split into patches.
+            In this case, the output `merged_shape` and the optional `cropped_shape` cannot be inferred
+            and should be explicitly provided.
         merger_cls: a `Merger` subclass that can be instantiated to merges patch outputs.
             It can also be a string that matches the name of a class inherited from `Merger` class.
             Defaults to `AvgMerger`.
@@ -100,34 +102,29 @@ class PatchInferer(Inferer):
         output_keys: if the network output is a dictionary, this defines the keys of
             the output dictionary to be used for merging.
             Defaults to None, where all the keys are used.
+        match_spatial_shape: whether to crop the output to match the input shape. Defaults to True.
         merger_kwargs: arguments to be passed to `merger_cls` for instantiation.
-            `output_shape` is calculated automatically based on the input shape and
+            `merged_shape` is calculated automatically based on the input shape and
             the output patch shape unless it is passed here.
     """
 
     def __init__(
         self,
-        splitter: Splitter | Callable | None = None,
+        splitter: Splitter | None = None,
         merger_cls: type[Merger] | str = AvgMerger,
         batch_size: int = 1,
         preprocessing: Callable | None = None,
         postprocessing: Callable | None = None,
         output_keys: Sequence | None = None,
+        match_spatial_shape: bool = True,
         **merger_kwargs: Any,
     ) -> None:
         Inferer.__init__(self)
-
         # splitter
-        if splitter is not None and not isinstance(splitter, Splitter):
-            if callable(splitter):
-                warnings.warn(
-                    "`splitter` is a callable instead of `Splitter` object, please make sure that it returns "
-                    "the correct values. Either Iterable[tuple[torch.Tensor, Sequence[int]]], or "
-                    "a MetaTensor with defined `PatchKey.LOCATION` metadata."
-                )
-            else:
+        if not isinstance(splitter, (Splitter, type(None))):
+            if not isinstance(splitter, Splitter):
                 raise TypeError(
-                    f"'splitter' should be a `Splitter` object  (or a callable that returns "
+                    f"'splitter' should be a `Splitter` object that returns: "
                     "an iterable of pairs of (patch, location) or a MetaTensor that has `PatchKeys.LOCATION` metadata)."
                     f"{type(splitter)} is given."
                 )
@@ -165,6 +162,9 @@ def __init__(
         # model output keys
         self.output_keys = output_keys
 
+        # whether to crop the output to match the input shape
+        self.match_spatial_shape = match_spatial_shape
+
     def _batch_sampler(
         self, patches: Iterable[tuple[torch.Tensor, Sequence[int]]] | MetaTensor
     ) -> Iterator[tuple[torch.Tensor, Sequence, int]]:
@@ -226,14 +226,24 @@ def _initialize_mergers(self, inputs, outputs, patches, batch_size):
             out_patch = torch.chunk(out_patch_batch, batch_size)[0]
             # calculate the ratio of input and output patch sizes
             ratio = tuple(op / ip for ip, op in zip(in_patch.shape[2:], out_patch.shape[2:]))
-            ratios.append(ratio)
-            # calculate output_shape only if it is not provided and splitter is not None.
-            if self.splitter is not None and "output_shape" not in self.merger_kwargs:
-                output_shape = self._get_output_shape(inputs, out_patch, ratio)
-                merger = self.merger_cls(output_shape=output_shape, **self.merger_kwargs)
-            else:
-                merger = self.merger_cls(**self.merger_kwargs)
+
+            # calculate merged_shape and cropped_shape
+            merger_kwargs = self.merger_kwargs.copy()
+            cropped_shape, merged_shape = self._get_merged_shapes(inputs, out_patch, ratio)
+            if "merged_shape" not in merger_kwargs:
+                merger_kwargs["merged_shape"] = merged_shape
+                if merger_kwargs["merged_shape"] is None:
+                    raise ValueError("`merged_shape` cannot be `None`.")
+            if "cropped_shape" not in merger_kwargs:
+                merger_kwargs["cropped_shape"] = cropped_shape
+
+            # initialize the merger
+            merger = self.merger_cls(**merger_kwargs)
+
+            # store mergers and input/output ratios
             mergers.append(merger)
+            ratios.append(ratio)
+
         return mergers, ratios
 
     def _aggregate(self, outputs, locations, batch_size, mergers, ratios):
@@ -243,12 +253,27 @@ def _aggregate(self, outputs, locations, batch_size, mergers, ratios):
                 out_loc = [round(l * r) for l, r in zip(in_loc, ratio)]
                 merger.aggregate(out_patch, out_loc)
 
-    def _get_output_shape(self, inputs, out_patch, ratio):
-        """Define the shape of output merged tensors"""
-        in_spatial_shape = inputs.shape[2:]
-        out_spatial_shape = tuple(round(s * r) for s, r in zip(in_spatial_shape, ratio))
-        output_shape = out_patch.shape[:2] + out_spatial_shape
-        return output_shape
+    def _get_merged_shapes(self, inputs, out_patch, ratio):
+        """Define the shape of merged tensors (non-padded and padded)"""
+        if self.splitter is None:
+            return None, None
+
+        # input spatial shapes
+        original_spatial_shape = self.splitter.get_input_shape(inputs)
+        padded_spatial_shape = self.splitter.get_padded_shape(inputs)
+
+        # output spatial shapes
+        output_spatial_shape = tuple(round(s * r) for s, r in zip(original_spatial_shape, ratio))
+        padded_output_spatial_shape = tuple(round(s * r) for s, r in zip(padded_spatial_shape, ratio))
+
+        # output shapes
+        cropped_shape = out_patch.shape[:2] + output_spatial_shape
+        merged_shape = out_patch.shape[:2] + padded_output_spatial_shape
+
+        if not self.match_spatial_shape:
+            cropped_shape = merged_shape
+
+        return cropped_shape, merged_shape
 
     def __call__(
         self,
@@ -270,6 +295,7 @@ def __call__(
         """
         patches_locations: Iterable[tuple[torch.Tensor, Sequence[int]]] | MetaTensor
         if self.splitter is None:
+            # handle situations where the splitter is not provided
             if isinstance(inputs, torch.Tensor):
                 if isinstance(inputs, MetaTensor):
                     if PatchKeys.LOCATION not in inputs.meta:
@@ -288,6 +314,7 @@ def __call__(
                     )
             patches_locations = inputs
         else:
+            # apply splitter
             patches_locations = self.splitter(inputs)
 
         ratios: list[float] = []
@@ -302,7 +329,8 @@ def __call__(
             self._aggregate(outputs, locations, batch_size, mergers, ratios)
 
         # finalize the mergers and get the results
-        merged_outputs = tuple(merger.finalize() for merger in mergers)
+        merged_outputs = [merger.finalize() for merger in mergers]
+
         # return according to the model output
         if self.output_keys:
             return dict(zip(self.output_keys, merged_outputs))

monai/inferers/merger.py

@@ -32,12 +32,20 @@ class Merger(ABC):
         - finalize: perform any final process and return the merged output
 
     Args:
-        output_shape: the shape of the merged output tensor. Default to None.
+        merged_shape: the shape of the tensor required to merge the patches.
+        cropped_shape: the shape of the final merged output tensor.
+            If not provided, it will be the same as `merged_shape`.
         device: the device where Merger tensors should reside.
     """
 
-    def __init__(self, output_shape: Sequence[int] | None = None, device: torch.device | str | None = None) -> None:
-        self.output_shape = output_shape
+    def __init__(
+        self,
+        merged_shape: Sequence[int],
+        cropped_shape: Sequence[int] | None = None,
+        device: torch.device | str | None = None,
+    ) -> None:
+        self.merged_shape = merged_shape
+        self.cropped_shape = self.merged_shape if cropped_shape is None else cropped_shape
         self.device = device
         self.is_finalized = False
 
@@ -77,26 +85,29 @@ class AvgMerger(Merger):
     """Merge patches by taking average of the overlapping area
 
     Args:
-        output_shape: the shape of the merged output tensor.
+        merged_shape: the shape of the tensor required to merge the patches.
+        cropped_shape: the shape of the final merged output tensor.
+            If not provided, it will be the same as `merged_shape`.
         device: the device for aggregator tensors and final results.
         value_dtype: the dtype for value aggregating tensor and the final result.
         count_dtype: the dtype for sample counting tensor.
     """
 
     def __init__(
         self,
-        output_shape: Sequence[int],
+        merged_shape: Sequence[int],
+        cropped_shape: Sequence[int] | None = None,
         device: torch.device | str = "cpu",
         value_dtype: torch.dtype = torch.float32,
         count_dtype: torch.dtype = torch.uint8,
     ) -> None:
-        super().__init__(output_shape=output_shape, device=device)
-        if not self.output_shape:
-            raise ValueError(f"`output_shape` must be provided for `AvgMerger`. {self.output_shape} is give.")
+        super().__init__(merged_shape=merged_shape, cropped_shape=cropped_shape, device=device)
+        if not self.merged_shape:
+            raise ValueError(f"`merged_shape` must be provided for `AvgMerger`. {self.merged_shape} is give.")
         self.value_dtype = value_dtype
         self.count_dtype = count_dtype
-        self.values = torch.zeros(self.output_shape, dtype=self.value_dtype, device=self.device)
-        self.counts = torch.zeros(self.output_shape, dtype=self.count_dtype, device=self.device)
+        self.values = torch.zeros(self.merged_shape, dtype=self.value_dtype, device=self.device)
+        self.counts = torch.zeros(self.merged_shape, dtype=self.count_dtype, device=self.device)
 
     def aggregate(self, values: torch.Tensor, location: Sequence[int]) -> None:
         """
@@ -134,6 +145,8 @@ def finalize(self) -> torch.Tensor:
         if not self.is_finalized:
             # use in-place division to save space
             self.values.div_(self.counts)
+            # finalize the shape
+            self.values = self.values[tuple(slice(0, end) for end in self.cropped_shape)]
             # set finalize flag to protect performing in-place division again
             self.is_finalized = True