SliceInferer for slice-by-slice infererence

docs/source/inferers.rst

@@ -36,3 +36,9 @@ Inferers
 .. autoclass:: SaliencyInferer
     :members:
     :special-members: __call__
+
+`SliceInferer`
+~~~~~~~~~~~~~~
+.. autoclass:: SliceInferer
+    :members:
+    :special-members: __call__

monai/inferers/__init__.py

@@ -9,5 +9,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .inferer import Inferer, SaliencyInferer, SimpleInferer, SlidingWindowInferer
+from .inferer import Inferer, SaliencyInferer, SimpleInferer, SliceInferer, SlidingWindowInferer
 from .utils import sliding_window_inference

monai/inferers/inferer.py

@@ -16,10 +16,10 @@
 import torch.nn as nn
 
 from monai.inferers.utils import sliding_window_inference
-from monai.utils import BlendMode, PytorchPadMode
+from monai.utils import BlendMode, PytorchPadMode, ensure_tuple
 from monai.visualize import CAM, GradCAM, GradCAMpp
 
-__all__ = ["Inferer", "SimpleInferer", "SlidingWindowInferer", "SaliencyInferer"]
+__all__ = ["Inferer", "SimpleInferer", "SlidingWindowInferer", "SaliencyInferer", "SliceInferer"]
 
 
 class Inferer(ABC):
@@ -222,3 +222,59 @@ def __call__(self, inputs: torch.Tensor, network: nn.Module, *args: Any, **kwarg
             cam = GradCAMpp(network, self.target_layers, *self.args, **self.kwargs)
 
         return cam(inputs, self.class_idx, *args, **kwargs)
+
+
+class SliceInferer(SlidingWindowInferer):
+    """
+    SliceInferer extends SlidingWindowInferer to provide slice-by-slice (2D) inference
+    when provided a 3D volume.
+
+    Args:
+        spatial_dim: Spatial dimension over which the slice-by-slice inference runs on the 3D volume.
+            For example ``0`` could slide over axial slices. ``1`` over coronal slices and ``2`` over sagittal slices.
+        args: other optional args to be passed to the `__init__` of base class SlidingWindowInferer.
+        kwargs: other optional keyword args to be passed to `__init__` of base class SlidingWindowInferer.
+
+    Note:
+        ``roi_size`` in SliceInferer is expected to be a 2D tuple when a 3D volume is provided. This allows
+        sliding across slices along the 3D volume using a selected ``spatial_dim``.
+
+    """
+
+    def __init__(self, spatial_dim: int = 0, *args, **kwargs) -> None:
+        self.spatial_dim = spatial_dim
+        super().__init__(*args, **kwargs)
+
+    def __call__(
+        self, inputs: torch.Tensor, network: Callable[..., torch.Tensor], *args: Any, **kwargs: Any
+    ) -> torch.Tensor:
+        """
+        Args:
+            inputs: 3D input for inference
+            network: 2D model to execute inference on slices in the 3D input
+            args: optional args to be passed to ``network``.
+            kwargs: optional keyword args to be passed to ``network``.
+        """
+        if self.spatial_dim > 2:
+            raise ValueError("`spatial_dim` can only be `0, 1, 2` with `[H, W, D]` respectively.")
+
+        # Check if ``roi_size`` tuple is 2D and ``inputs`` tensor is 3D
+        self.roi_size = ensure_tuple(self.roi_size)
+        if len(self.roi_size) == 2 and len(inputs.shape[2:]) == 3:
+            self.roi_size = list(self.roi_size)
+            self.roi_size.insert(self.spatial_dim, 1)
+        else:
+            raise RuntimeError("Currently, only 2D `roi_size` with 3D `inputs` tensor is supported.")
+
+        return super().__call__(inputs=inputs, network=lambda x: self.network_wrapper(network, x, *args, **kwargs))
+
+    def network_wrapper(self, network: Callable[..., torch.Tensor], x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
+        """
+        Wrapper handles inference for 2D models over 3D volume inputs.
+        """
+        #  Pass 4D input [N, C, H, W]/[N, C, D, W]/[N, C, D, H] to the model as it is 2D.
+        x = x.squeeze(dim=self.spatial_dim + 2)
+        out = network(x, *args, **kwargs)
+        #  Unsqueeze the network output so it is [N, C, D, H, W] as expected by
+        # the default SlidingWindowInferer class
+        return out.unsqueeze(dim=self.spatial_dim + 2)

tests/test_slice_inferer.py

@@ -0,0 +1,51 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import torch
+from parameterized import parameterized
+
+from monai.inferers import SliceInferer
+from monai.networks.nets import UNet
+
+TEST_CASES = ["0", "1", "2"]
+
+
+class TestSliceInferer(unittest.TestCase):
+    @parameterized.expand(TEST_CASES)
+    def test_shape(self, spatial_dim):
+        spatial_dim = int(spatial_dim)
+
+        model = UNet(
+            spatial_dims=2, in_channels=1, out_channels=1, channels=(4, 8, 16), strides=(2, 2), num_res_units=2
+        )
+
+        device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        model.to(device)
+        model.eval()
+
+        # Initialize a dummy 3D tensor volume with shape (N,C,D,H,W)
+        input_volume = torch.ones(1, 1, 64, 256, 256, device=device)
+
+        # Remove spatial dim to slide across from the roi_size
+        roi_size = list(input_volume.shape[2:])
+        roi_size.pop(spatial_dim)
+
+        # Initialize and run inferer
+        inferer = SliceInferer(roi_size=roi_size, spatial_dim=spatial_dim, sw_batch_size=1, cval=-1)
+        result = inferer(input_volume, model)
+
+        self.assertTupleEqual(result.shape, input_volume.shape)
+
+
+if __name__ == "__main__":
+    unittest.main()