[2678] Add transform to fill holes and to filter (#2692)

* Add transform to fill holes and to filter (#2678)

Signed-off-by: Sebastian Penhouet <[email protected]>

* Change name of label filter class (#2678)

Signed-off-by: Sebastian Penhouet <[email protected]>

* Change fill holes to growing logic (#2678)

Signed-off-by: Sebastian Penhouet <[email protected]>

* Fix missing entry in min_tests (#2678)

Signed-off-by: Sebastian Penhouet <[email protected]>

* Fix review comments (#2678)

Signed-off-by: Sebastian Penhouet <[email protected]>

* Remove batch dim and add one-hot handling (#2678)

Signed-off-by: Sebastian Penhouet <[email protected]>

* [MONAI] python code formatting

Signed-off-by: monai-bot <[email protected]>

Co-authored-by: Sebastian Penhouet <[email protected]>

Author: Spenhouet

docs/source/transforms.rst

@@ -356,6 +356,18 @@ Post-processing
     :members:
     :special-members: __call__
 
+`LabelFilter`
+"""""""""""""
+.. autoclass:: LabelFilter
+    :members:
+    :special-members: __call__
+
+`FillHoles`
+"""""""""""
+.. autoclass:: FillHoles
+    :members:
+    :special-members: __call__
+
 `LabelToContour`
 """"""""""""""""
 .. autoclass:: LabelToContour
@@ -955,6 +967,18 @@ Post-processing (Dict)
     :members:
     :special-members: __call__
 
+`LabelFilterd`
+""""""""""""""
+.. autoclass:: LabelFilterd
+    :members:
+    :special-members: __call__
+
+`FillHolesd`
+""""""""""""
+.. autoclass:: FillHolesd
+    :members:
+    :special-members: __call__
+
 `LabelToContourd`
 """""""""""""""""
 .. autoclass:: LabelToContourd

monai/transforms/__init__.py

@@ -194,40 +194,48 @@
 from .post.array import (
     Activations,
     AsDiscrete,
+    FillHoles,
     KeepLargestConnectedComponent,
+    LabelFilter,
     LabelToContour,
     MeanEnsemble,
     ProbNMS,
     VoteEnsemble,
 )
 from .post.dictionary import (
-    Activationsd,
     ActivationsD,
+    Activationsd,
     ActivationsDict,
-    AsDiscreted,
     AsDiscreteD,
+    AsDiscreted,
     AsDiscreteDict,
     Ensembled,
-    Invertd,
+    FillHolesD,
+    FillHolesd,
+    FillHolesDict,
     InvertD,
+    Invertd,
     InvertDict,
-    KeepLargestConnectedComponentd,
     KeepLargestConnectedComponentD,
+    KeepLargestConnectedComponentd,
     KeepLargestConnectedComponentDict,
-    LabelToContourd,
+    LabelFilterD,
+    LabelFilterd,
+    LabelFilterDict,
     LabelToContourD,
+    LabelToContourd,
     LabelToContourDict,
-    MeanEnsembled,
     MeanEnsembleD,
+    MeanEnsembled,
     MeanEnsembleDict,
-    ProbNMSd,
     ProbNMSD,
+    ProbNMSd,
     ProbNMSDict,
-    SaveClassificationd,
     SaveClassificationD,
+    SaveClassificationd,
     SaveClassificationDict,
-    VoteEnsembled,
     VoteEnsembleD,
+    VoteEnsembled,
     VoteEnsembleDict,
 )
 from .spatial.array import (

monai/transforms/post/array.py

@@ -14,26 +14,29 @@
 """
 
 import warnings
-from typing import Callable, Optional, Sequence, Union
+from typing import Callable, Iterable, Optional, Sequence, Union
 
 import numpy as np
 import torch
 import torch.nn.functional as F
 
+from monai.config import NdarrayTensor
 from monai.networks import one_hot
 from monai.networks.layers import GaussianFilter
 from monai.transforms.transform import Transform
-from monai.transforms.utils import get_largest_connected_component_mask
+from monai.transforms.utils import fill_holes, get_largest_connected_component_mask
 from monai.utils import ensure_tuple
 
 __all__ = [
     "Activations",
     "AsDiscrete",
+    "FillHoles",
     "KeepLargestConnectedComponent",
+    "LabelFilter",
     "LabelToContour",
     "MeanEnsemble",
-    "VoteEnsemble",
     "ProbNMS",
+    "VoteEnsemble",
 ]
 
 
@@ -289,6 +292,137 @@ def __call__(self, img: torch.Tensor) -> torch.Tensor:
         return output
 
 
+class LabelFilter:
+    """
+    This transform filters out labels and can be used as a processing step to view only certain labels.
+
+    The list of applied labels defines which labels will be kept.
+
+    Note:
+        All labels which do not match the `applied_labels` are set to the background label (0).
+
+    For example:
+
+    Use LabelFilter with applied_labels=[1, 5, 9]::
+
+        [1, 2, 3]         [1, 0, 0]
+        [4, 5, 6]    =>   [0, 5 ,0]
+        [7, 8, 9]         [0, 0, 9]
+    """
+
+    def __init__(self, applied_labels: Union[Iterable[int], int]) -> None:
+        """
+        Initialize the LabelFilter class with the labels to filter on.
+
+        Args:
+            applied_labels: Label(s) to filter on.
+        """
+        self.applied_labels = ensure_tuple(applied_labels)
+
+    def __call__(self, img: NdarrayTensor) -> NdarrayTensor:
+        """
+        Filter the image on the `applied_labels`.
+
+        Args:
+            img: Pytorch tensor or numpy array of any shape.
+
+        Raises:
+            NotImplementedError: The provided image was not a Pytorch Tensor or numpy array.
+
+        Returns:
+            Pytorch tensor or numpy array of the same shape as the input.
+        """
+        if isinstance(img, np.ndarray):
+            return np.asarray(np.where(np.isin(img, self.applied_labels), img, 0))
+        elif isinstance(img, torch.Tensor):
+            img_arr = img.detach().cpu().numpy()
+            img_arr = self(img_arr)
+            return torch.as_tensor(img_arr, device=img.device)
+        else:
+            raise NotImplementedError(f"{self.__class__} can not handle data of type {type(img)}.")
+
+
+class FillHoles(Transform):
+    r"""
+    This transform fills holes in the image and can be used to remove artifacts inside segments.
+
+    An enclosed hole is defined as a background pixel/voxel which is only enclosed by a single class.
+    The definition of enclosed can be defined with the connectivity parameter::
+
+        1-connectivity     2-connectivity     diagonal connection close-up
+
+             [ ]           [ ]  [ ]  [ ]             [ ]
+              |               \  |  /                 |  <- hop 2
+        [ ]--[x]--[ ]      [ ]--[x]--[ ]        [x]--[ ]
+              |               /  |  \             hop 1
+             [ ]           [ ]  [ ]  [ ]
+
+    It is possible to define for which labels the hole filling should be applied.
+    The input image is assumed to be a PyTorch Tensor or numpy array with shape [C, spatial_dim1[, spatial_dim2, ...]].
+    If C = 1, then the values correspond to expected labels.
+    If C > 1, then a one-hot-encoding is expected where the index of C matches the label indexing.
+
+    Note:
+
+        The label 0 will be treated as background and the enclosed holes will be set to the neighboring class label.
+
+        The performance of this method heavily depends on the number of labels.
+        It is a bit faster if the list of `applied_labels` is provided.
+        Limiting the number of `applied_labels` results in a big decrease in processing time.
+
+    For example:
+
+        Use FillHoles with default parameters::
+
+            [1, 1, 1, 2, 2, 2, 3, 3]         [1, 1, 1, 2, 2, 2, 3, 3]
+            [1, 0, 1, 2, 0, 0, 3, 0]    =>   [1, 1 ,1, 2, 0, 0, 3, 0]
+            [1, 1, 1, 2, 2, 2, 3, 3]         [1, 1, 1, 2, 2, 2, 3, 3]
+
+        The hole in label 1 is fully enclosed and therefore filled with label 1.
+        The background label near label 2 and 3 is not fully enclosed and therefore not filled.
+    """
+
+    def __init__(
+        self, applied_labels: Optional[Union[Iterable[int], int]] = None, connectivity: Optional[int] = None
+    ) -> None:
+        """
+        Initialize the connectivity and limit the labels for which holes are filled.
+
+        Args:
+            applied_labels: Labels for which to fill holes. Defaults to None, that is filling holes for all labels.
+            connectivity: Maximum number of orthogonal hops to consider a pixel/voxel as a neighbor.
+                Accepted values are ranging from  1 to input.ndim. Defaults to a full connectivity of ``input.ndim``.
+        """
+        super().__init__()
+        self.applied_labels = ensure_tuple(applied_labels) if applied_labels else None
+        self.connectivity = connectivity
+
+    def __call__(self, img: NdarrayTensor) -> NdarrayTensor:
+        """
+        Fill the holes in the provided image.
+
+        Note:
+            The value 0 is assumed as background label.
+
+        Args:
+            img: Pytorch Tensor or numpy array of shape [C, spatial_dim1[, spatial_dim2, ...]].
+
+        Raises:
+            NotImplementedError: The provided image was not a Pytorch Tensor or numpy array.
+
+        Returns:
+            Pytorch Tensor or numpy array of shape [C, spatial_dim1[, spatial_dim2, ...]].
+        """
+        if isinstance(img, np.ndarray):
+            return fill_holes(img, self.applied_labels, self.connectivity)
+        elif isinstance(img, torch.Tensor):
+            img_arr = img.detach().cpu().numpy()
+            img_arr = self(img_arr)
+            return torch.as_tensor(img_arr, device=img.device)
+        else:
+            raise NotImplementedError(f"{self.__class__} can not handle data of type {type(img)}.")
+
+
 class LabelToContour(Transform):
     """
     Return the contour of binary input images that only compose of 0 and 1, with Laplace kernel