Adding a median filter for 3D images (#5307)

Co-authored-by: @ebrahimebrahim. Closes #5264.

### 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).
- [ ] Breaking change (fix or new feature that would cause existing
functionality to change).
- [x] New tests added to cover the changes.
- [ ] Integration tests passed locally by running `./runtests.sh -f -u
--net --coverage`.
- [ ] 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.

Author: dzenanz

docs/source/networks.rst

@@ -334,6 +334,16 @@ Layers
 .. autoclass:: GaussianFilter
     :members:
 
+`MedianFilter`
+~~~~~~~~~~~~~~
+.. autoclass:: MedianFilter
+    :members:
+
+`median_filter`
+~~~~~~~~~~~~~~~
+.. autoclass:: median_filter
+    :members:
+
 `BilateralFilter`
 ~~~~~~~~~~~~~~~~~
 .. autoclass:: BilateralFilter

docs/source/transforms.rst

@@ -336,6 +336,14 @@ Intensity
     :members:
     :special-members: __call__
 
+`MedianSmooth`
+""""""""""""""
+.. image:: https://github.com/Project-MONAI/DocImages/raw/main/transforms/MedianSmooth.png
+    :alt: example of MedianSmooth
+.. autoclass:: MedianSmooth
+    :members:
+    :special-members: __call__
+
 `GaussianSmooth`
 """"""""""""""""
 .. image:: https://github.com/Project-MONAI/DocImages/raw/main/transforms/GaussianSmooth.png
@@ -1415,6 +1423,14 @@ Intensity (Dict)
     :members:
     :special-members: __call__
 
+`MedianSmoothd`
+"""""""""""""""
+.. image:: https://github.com/Project-MONAI/DocImages/raw/main/transforms/MedianSmoothd.png
+    :alt: example of MedianSmoothd
+.. autoclass:: MedianSmoothd
+    :members:
+    :special-members: __call__
+
 `GaussianSmoothd`
 """""""""""""""""
 .. image:: https://github.com/Project-MONAI/DocImages/raw/main/transforms/GaussianSmoothd.png

monai/networks/layers/__init__.py

@@ -20,10 +20,12 @@
     Flatten,
     GaussianFilter,
     HilbertTransform,
+    MedianFilter,
     Reshape,
     SavitzkyGolayFilter,
     SkipConnection,
     apply_filter,
+    median_filter,
     separable_filtering,
 )
 from .spatial_transforms import AffineTransform, grid_count, grid_grad, grid_pull, grid_push

monai/networks/layers/simplelayers.py

@@ -11,7 +11,7 @@
 
 import math
 from copy import deepcopy
-from typing import List, Sequence, Union
+from typing import List, Optional, Sequence, Union
 
 import torch
 import torch.nn.functional as F
@@ -20,8 +20,16 @@
 
 from monai.networks.layers.convutils import gaussian_1d
 from monai.networks.layers.factories import Conv
-from monai.utils import ChannelMatching, SkipMode, look_up_option, optional_import, pytorch_after
-from monai.utils.misc import issequenceiterable
+from monai.utils import (
+    ChannelMatching,
+    SkipMode,
+    convert_to_tensor,
+    ensure_tuple_rep,
+    issequenceiterable,
+    look_up_option,
+    optional_import,
+    pytorch_after,
+)
 
 _C, _ = optional_import("monai._C")
 fft, _ = optional_import("torch.fft")
@@ -32,10 +40,12 @@
     "GaussianFilter",
     "HilbertTransform",
     "LLTM",
+    "MedianFilter",
     "Reshape",
     "SavitzkyGolayFilter",
     "SkipConnection",
     "apply_filter",
+    "median_filter",
     "separable_filtering",
 ]
 
@@ -168,7 +178,6 @@ def _separable_filtering_conv(
     paddings: List[int],
     num_channels: int,
 ) -> torch.Tensor:
-
     if d < 0:
         return input_
 
@@ -434,6 +443,121 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return torch.as_tensor(ht, device=ht.device, dtype=ht.dtype)
 
 
+def get_binary_kernel(window_size: Sequence[int], dtype=torch.float, device=None) -> torch.Tensor:
+    """
+    Create a binary kernel to extract the patches.
+    The window size HxWxD will create a (H*W*D)xHxWxD kernel.
+    """
+    win_size = convert_to_tensor(window_size, int, wrap_sequence=True)
+    prod = torch.prod(win_size)
+    s = [prod, 1, *win_size]
+    return torch.diag(torch.ones(prod, dtype=dtype, device=device)).view(s)  # type: ignore
+
+
+def median_filter(
+    in_tensor: torch.Tensor,
+    kernel_size: Sequence[int] = (3, 3, 3),
+    spatial_dims: int = 3,
+    kernel: Optional[torch.Tensor] = None,
+    **kwargs,
+) -> torch.Tensor:
+    """
+    Apply median filter to an image.
+
+    Args:
+        in_tensor: input tensor; median filtering will be applied to the last `spatial_dims` dimensions.
+        kernel_size: the convolution kernel size.
+        spatial_dims: number of spatial dimensions to apply median filtering.
+        kernel: an optional customized kernel.
+        kwargs: additional parameters to the `conv`.
+
+    Returns:
+        the filtered input tensor, shape remains the same as ``in_tensor``
+
+    Example::
+
+        >>> from monai.networks.layers import median_filter
+        >>> import torch
+        >>> x = torch.rand(4, 5, 7, 6)
+        >>> output = median_filter(x, (3, 3, 3))
+        >>> output.shape
+        torch.Size([4, 5, 7, 6])
+
+    """
+    if not isinstance(in_tensor, torch.Tensor):
+        raise TypeError(f"Input type is not a torch.Tensor. Got {type(in_tensor)}")
+
+    original_shape = in_tensor.shape
+    oshape, sshape = original_shape[: len(original_shape) - spatial_dims], original_shape[-spatial_dims:]
+    oprod = torch.prod(convert_to_tensor(oshape, int, wrap_sequence=True))
+    # prepare kernel
+    if kernel is None:
+        kernel_size = ensure_tuple_rep(kernel_size, spatial_dims)
+        kernel = get_binary_kernel(kernel_size, in_tensor.dtype, in_tensor.device)
+    else:
+        kernel = kernel.to(in_tensor)
+    # map the local window to single vector
+    conv = [F.conv1d, F.conv2d, F.conv3d][spatial_dims - 1]  # type: ignore
+    if "padding" not in kwargs:
+        if pytorch_after(1, 10):
+            kwargs["padding"] = "same"
+        else:
+            # even-sized kernels are not supported
+            kwargs["padding"] = [(k - 1) // 2 for k in kernel.shape[2:]]
+    elif kwargs["padding"] == "same" and not pytorch_after(1, 10):
+        # even-sized kernels are not supported
+        kwargs["padding"] = [(k - 1) // 2 for k in kernel.shape[2:]]
+    features: torch.Tensor = conv(in_tensor.reshape(oprod, 1, *sshape), kernel, stride=1, **kwargs)  # type: ignore
+    features = features.view(oprod, -1, *sshape)  # type: ignore
+
+    # compute the median along the feature axis
+    median: torch.Tensor = torch.median(features, dim=1)[0]
+    median = median.reshape(original_shape)
+
+    return median
+
+
+class MedianFilter(nn.Module):
+    """
+    Apply median filter to an image.
+
+    Args:
+        radius: the blurring kernel radius (radius of 1 corresponds to 3x3x3 kernel when spatial_dims=3).
+
+    Returns:
+        filtered input tensor.
+
+    Example::
+
+        >>> from monai.networks.layers import MedianFilter
+        >>> import torch
+        >>> in_tensor = torch.rand(4, 5, 7, 6)
+        >>> blur = MedianFilter([1, 1, 1])  # 3x3x3 kernel
+        >>> output = blur(in_tensor)
+        >>> output.shape
+        torch.Size([4, 5, 7, 6])
+
+    """
+
+    def __init__(self, radius: Union[Sequence[int], int], spatial_dims: int = 3, device="cpu") -> None:
+        super().__init__()
+        self.spatial_dims = spatial_dims
+        self.radius: Sequence[int] = ensure_tuple_rep(radius, spatial_dims)
+        self.window: Sequence[int] = [1 + 2 * deepcopy(r) for r in self.radius]
+        self.kernel = get_binary_kernel(self.window, device=device)
+
+    def forward(self, in_tensor: torch.Tensor, number_of_passes=1) -> torch.Tensor:
+        """
+        Args:
+            in_tensor: input tensor, median filtering will be applied to the last `spatial_dims` dimensions.
+            number_of_passes: median filtering will be repeated this many times
+        """
+        x = in_tensor
+        for _ in range(number_of_passes):
+            x = median_filter(x, kernel=self.kernel, spatial_dims=self.spatial_dims)
+        return x
+
+
 class GaussianFilter(nn.Module):
     def __init__(
         self,

monai/transforms/__init__.py

@@ -99,6 +99,7 @@
     IntensityRemap,
     KSpaceSpikeNoise,
     MaskIntensity,
+    MedianSmooth,
     NormalizeIntensity,
     RandAdjustContrast,
     RandBiasField,
@@ -152,6 +153,9 @@
     MaskIntensityd,
     MaskIntensityD,
     MaskIntensityDict,
+    MedianSmoothd,
+    MedianSmoothD,
+    MedianSmoothDict,
     NormalizeIntensityd,
     NormalizeIntensityD,
     NormalizeIntensityDict,

monai/transforms/intensity/array.py

@@ -26,7 +26,7 @@
 from monai.config.type_definitions import NdarrayOrTensor, NdarrayTensor
 from monai.data.meta_obj import get_track_meta
 from monai.data.utils import get_random_patch, get_valid_patch_size
-from monai.networks.layers import GaussianFilter, HilbertTransform, SavitzkyGolayFilter
+from monai.networks.layers import GaussianFilter, HilbertTransform, MedianFilter, SavitzkyGolayFilter
 from monai.transforms.transform import RandomizableTransform, Transform
 from monai.transforms.utils import Fourier, equalize_hist, is_positive, rescale_array
 from monai.transforms.utils_pytorch_numpy_unification import clip, percentile, where
@@ -56,6 +56,7 @@
     "MaskIntensity",
     "DetectEnvelope",
     "SavitzkyGolaySmooth",
+    "MedianSmooth",
     "GaussianSmooth",
     "RandGaussianSmooth",
     "GaussianSharpen",
@@ -1136,6 +1137,35 @@ def __call__(self, img: NdarrayOrTensor):
         return out
 
 
+class MedianSmooth(Transform):
+    """
+    Apply median filter to the input data based on specified `radius` parameter.
+    A default value `radius=1` is provided for reference.
+
+    See also: :py:func:`monai.networks.layers.median_filter`
+
+    Args:
+        radius: if a list of values, must match the count of spatial dimensions of input data,
+            and apply every value in the list to 1 spatial dimension. if only 1 value provided,
+            use it for all spatial dimensions.
+    """
+
+    backend = [TransformBackends.TORCH]
+
+    def __init__(self, radius: Union[Sequence[int], int] = 1) -> None:
+        self.radius = radius
+
+    def __call__(self, img: NdarrayTensor) -> NdarrayTensor:
+        img = convert_to_tensor(img, track_meta=get_track_meta())
+        img_t, *_ = convert_data_type(img, torch.Tensor, dtype=torch.float)
+        spatial_dims = img_t.ndim - 1
+        r = ensure_tuple_rep(self.radius, spatial_dims)
+        median_filter_instance = MedianFilter(r, spatial_dims=spatial_dims)
+        out_t: torch.Tensor = median_filter_instance(img_t)
+        out, *_ = convert_to_dst_type(out_t, dst=img, dtype=out_t.dtype)
+        return out
+
+
 class GaussianSmooth(Transform):
     """
     Apply Gaussian smooth to the input data based on specified `sigma` parameter.

monai/transforms/intensity/dictionary.py

@@ -32,6 +32,7 @@
     HistogramNormalize,
     KSpaceSpikeNoise,
     MaskIntensity,
+    MedianSmooth,
     NormalizeIntensity,
     RandAdjustContrast,
     RandBiasField,
@@ -78,6 +79,7 @@
     "ScaleIntensityRangePercentilesd",
     "MaskIntensityd",
     "SavitzkyGolaySmoothd",
+    "MedianSmoothd",
     "GaussianSmoothd",
     "RandGaussianSmoothd",
     "GaussianSharpend",
@@ -124,6 +126,8 @@
     "MaskIntensityDict",
     "SavitzkyGolaySmoothD",
     "SavitzkyGolaySmoothDict",
+    "MedianSmoothD",
+    "MedianSmoothDict",
     "GaussianSmoothD",
     "GaussianSmoothDict",
     "RandGaussianSmoothD",
@@ -988,6 +992,35 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N
         return d
 
 
+class MedianSmoothd(MapTransform):
+    """
+    Dictionary-based wrapper of :py:class:`monai.transforms.MedianSmooth`.
+
+    Args:
+        keys: keys of the corresponding items to be transformed.
+            See also: :py:class:`monai.transforms.compose.MapTransform`
+        radius: if a list of values, must match the count of spatial dimensions of input data,
+            and apply every value in the list to 1 spatial dimension. if only 1 value provided,
+            use it for all spatial dimensions.
+        allow_missing_keys: don't raise exception if key is missing.
+
+    """
+
+    backend = MedianSmooth.backend
+
+    def __init__(
+        self, keys: KeysCollection, radius: Union[Sequence[int], int], allow_missing_keys: bool = False
+    ) -> None:
+        super().__init__(keys, allow_missing_keys)
+        self.converter = MedianSmooth(radius)
+
+    def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]:
+        d = dict(data)
+        for key in self.key_iterator(d):
+            d[key] = self.converter(d[key])
+        return d
+
+
 class GaussianSmoothd(MapTransform):
     """
     Dictionary-based wrapper of :py:class:`monai.transforms.GaussianSmooth`.
@@ -1780,6 +1813,7 @@ def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, N
 ScaleIntensityRangePercentilesD = ScaleIntensityRangePercentilesDict = ScaleIntensityRangePercentilesd
 MaskIntensityD = MaskIntensityDict = MaskIntensityd
 SavitzkyGolaySmoothD = SavitzkyGolaySmoothDict = SavitzkyGolaySmoothd
+MedianSmoothD = MedianSmoothDict = MedianSmoothd
 GaussianSmoothD = GaussianSmoothDict = GaussianSmoothd
 RandGaussianSmoothD = RandGaussianSmoothDict = RandGaussianSmoothd
 GaussianSharpenD = GaussianSharpenDict = GaussianSharpend