KeepLargest... by default use all labels

monai/transforms/post/array.py

@@ -23,7 +23,7 @@
 from monai.networks import one_hot
 from monai.networks.layers import GaussianFilter, apply_filter
 from monai.transforms.transform import Transform
-from monai.transforms.utils import fill_holes, get_largest_connected_component_mask
+from monai.transforms.utils import fill_holes, get_largest_connected_component_mask, get_unique_labels
 from monai.transforms.utils_pytorch_numpy_unification import unravel_index
 from monai.utils import TransformBackends, convert_data_type, deprecated_arg, ensure_tuple, look_up_option
 from monai.utils.type_conversion import convert_to_dst_type
@@ -302,16 +302,16 @@ class KeepLargestConnectedComponent(Transform):
 
     def __init__(
         self,
-        applied_labels: Union[Sequence[int], int],
+        applied_labels: Optional[Union[Sequence[int], int]] = None,
         is_onehot: Optional[bool] = None,
         independent: bool = True,
         connectivity: Optional[int] = None,
     ) -> None:
         """
         Args:
             applied_labels: Labels for applying the connected component analysis on.
-                If not OneHot. The pixel whose value is in this list will be analyzed.
-                If the data is in OneHot format, this is used to determine which channels to apply.
+                If given, voxels whose value is in this list will be analyzed.
+                If `None`, all non-zero values will be analyzed.
             is_onehot: if `True`, treat the input data as OneHot format data, otherwise, not OneHot format data.
                 default to None, which treats multi-channel data as OneHot and single channel data as not OneHot.
             independent: whether to treat ``applied_labels`` as a union of foreground labels.
@@ -326,7 +326,7 @@ def __init__(
 
         """
         super().__init__()
-        self.applied_labels = ensure_tuple(applied_labels)
+        self.applied_labels = ensure_tuple(applied_labels) if applied_labels is not None else None
         self.is_onehot = is_onehot
         self.independent = independent
         self.connectivity = connectivity
@@ -340,8 +340,13 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
             An array with shape (C, spatial_dim1[, spatial_dim2, ...]).
         """
         is_onehot = img.shape[0] > 1 if self.is_onehot is None else self.is_onehot
+        if self.applied_labels is not None:
+            applied_labels = self.applied_labels
+        else:
+            applied_labels = tuple(get_unique_labels(img, is_onehot, discard=0))
+
         if self.independent:
-            for i in self.applied_labels:
+            for i in applied_labels:
                 foreground = img[i] > 0 if is_onehot else img[0] == i
                 mask = get_largest_connected_component_mask(foreground, self.connectivity)
                 if is_onehot:
@@ -350,15 +355,15 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
                     img[0][foreground != mask] = 0
             return img
         if not is_onehot:  # not one-hot, union of labels
-            labels, *_ = convert_to_dst_type(self.applied_labels, dst=img, wrap_sequence=True)
+            labels, *_ = convert_to_dst_type(applied_labels, dst=img, wrap_sequence=True)
             foreground = (img[..., None] == labels).any(-1)[0]
             mask = get_largest_connected_component_mask(foreground, self.connectivity)
             img[0][foreground != mask] = 0
             return img
         # one-hot, union of labels
-        foreground = (img[self.applied_labels, ...] == 1).any(0)
+        foreground = (img[applied_labels, ...] == 1).any(0)
         mask = get_largest_connected_component_mask(foreground, self.connectivity)
-        for i in self.applied_labels:
+        for i in applied_labels:
             img[i][foreground != mask] = 0
         return img
 

monai/transforms/post/dictionary.py

@@ -220,7 +220,7 @@ class KeepLargestConnectedComponentd(MapTransform):
     def __init__(
         self,
         keys: KeysCollection,
-        applied_labels: Union[Sequence[int], int],
+        applied_labels: Optional[Union[Sequence[int], int]] = None,
         is_onehot: Optional[bool] = None,
         independent: bool = True,
         connectivity: Optional[int] = None,
@@ -231,8 +231,8 @@ def __init__(
             keys: keys of the corresponding items to be transformed.
                 See also: :py:class:`monai.transforms.compose.MapTransform`
             applied_labels: Labels for applying the connected component analysis on.
-                If not OneHot. The pixel whose value is in this list will be analyzed.
-                If the data is in OneHot format, this is used to determine which channels to apply.
+                If given, voxels whose value is in this list will be analyzed.
+                If `None`, all non-zero values will be analyzed.
             is_onehot: if `True`, treat the input data as OneHot format data, otherwise, not OneHot format data.
                 default to None, which treats multi-channel data as OneHot and single channel data as not OneHot.
             independent: whether to treat ``applied_labels`` as a union of foreground labels.

monai/transforms/utils.py

@@ -14,7 +14,7 @@
 import warnings
 from contextlib import contextmanager
 from inspect import getmembers, isclass
-from typing import Any, Callable, Hashable, Iterable, List, Mapping, Optional, Sequence, Tuple, Union
+from typing import Any, Callable, Hashable, Iterable, List, Mapping, Optional, Sequence, Set, Tuple, Union
 
 import numpy as np
 import torch
@@ -34,6 +34,7 @@
     nonzero,
     ravel,
     searchsorted,
+    unique,
     unravel_index,
     where,
 )
@@ -103,6 +104,7 @@
     "print_transform_backends",
     "convert_pad_mode",
     "convert_to_contiguous",
+    "get_unique_labels",
 ]
 
 
@@ -968,6 +970,34 @@ def get_largest_connected_component_mask(img: NdarrayTensor, connectivity: Optio
     return convert_to_dst_type(largest_cc, dst=img, dtype=largest_cc.dtype)[0]
 
 
+def get_unique_labels(
+    img: NdarrayOrTensor, is_onehot: bool, discard: Optional[Union[int, Iterable[int]]] = None
+) -> Set[int]:
+    """Get list of non-background labels in an image.
+
+    Args:
+        img: Image to be processed. Shape should be [C, W, H, [D]] with C=1 if not onehot else `num_classes`.
+        is_onehot: Boolean as to whether input image is one-hotted. If one-hotted, only return channels with
+        discard: Can be used to remove labels (e.g., background). Can be any value, sequence of values, or
+            `None` (nothing is discarded).
+
+    Returns:
+        Set of labels
+    """
+    applied_labels: Set[int]
+    n_channels = img.shape[0]
+    if is_onehot:
+        applied_labels = {i for i, s in enumerate(img) if s.sum() > 0}
+    else:
+        if n_channels != 1:
+            raise ValueError("If input not one-hotted, should only be 1 channel.")
+        applied_labels = set(unique(img).tolist())
+    if discard is not None:
+        for i in ensure_tuple(discard):
+            applied_labels.discard(i)
+    return applied_labels
+
+
 def fill_holes(
     img_arr: np.ndarray, applied_labels: Optional[Iterable[int]] = None, connectivity: Optional[int] = None
 ) -> np.ndarray:
@@ -1004,7 +1034,7 @@ def fill_holes(
     structure = ndimage.generate_binary_structure(spatial_dims, connectivity or spatial_dims)
 
     # Get labels if not provided. Exclude background label.
-    applied_labels = set(applied_labels or (range(num_channels) if is_one_hot else np.unique(img_arr)))
+    applied_labels = set(applied_labels) if applied_labels is not None else get_unique_labels(img_arr, is_one_hot)
     background_label = 0
     applied_labels.discard(background_label)
 

monai/transforms/utils_pytorch_numpy_unification.py

@@ -41,6 +41,7 @@
     "ascontiguousarray",
     "stack",
     "mode",
+    "unique",
 ]
 
 
@@ -417,3 +418,12 @@ def mode(x: NdarrayTensor, dim: int = -1, to_long: bool = True) -> NdarrayTensor
     o_t = torch.mode(x_t, dim).values
     o, *_ = convert_to_dst_type(o_t, x)
     return o
+
+
+def unique(x: NdarrayTensor) -> NdarrayTensor:
+    """`torch.unique` with equivalent implementation for numpy.
+
+    Args:
+        x: array/tensor
+    """
+    return torch.unique(x) if isinstance(x, torch.Tensor) else np.unique(x)  # type: ignore

tests/test_get_unique_labels.py

@@ -0,0 +1,44 @@
+# 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
+import torch.nn.functional as F
+from parameterized import parameterized
+
+from monai.transforms.utils import get_unique_labels
+from monai.transforms.utils_pytorch_numpy_unification import moveaxis
+from tests.utils import TEST_NDARRAYS
+
+grid_raw = [[0, 0, 0], [0, 0, 1], [2, 2, 3], [5, 5, 6], [3, 6, 2], [5, 6, 6]]
+grid = torch.Tensor(grid_raw).unsqueeze(0).to(torch.int64)
+grid_onehot = moveaxis(F.one_hot(grid)[0], -1, 0)
+
+TESTS = []
+for p in TEST_NDARRAYS:
+    for o_h in (False, True):
+        im = grid_onehot if o_h else grid
+        TESTS.append([dict(img=p(im), is_onehot=o_h), {0, 1, 2, 3, 5, 6}])
+        TESTS.append([dict(img=p(im), is_onehot=o_h, discard=0), {1, 2, 3, 5, 6}])
+        TESTS.append([dict(img=p(im), is_onehot=o_h, discard=[1, 2]), {0, 3, 5, 6}])
+
+
+class TestGetUniqueLabels(unittest.TestCase):
+    @parameterized.expand(TESTS)
+    def test_correct_results(self, args, expected):
+        result = get_unique_labels(**args)
+        self.assertEqual(result, expected)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_keep_largest_connected_component.py

@@ -12,11 +12,21 @@
 import unittest
 
 import torch
+import torch.nn.functional as F
 from parameterized import parameterized
 
 from monai.transforms import KeepLargestConnectedComponent
+from monai.transforms.utils_pytorch_numpy_unification import moveaxis
+from monai.utils.type_conversion import convert_to_dst_type
 from tests.utils import TEST_NDARRAYS, assert_allclose
 
+
+def to_onehot(x):
+    out = moveaxis(F.one_hot(torch.as_tensor(x).long())[0], -1, 0)
+    out, *_ = convert_to_dst_type(out, x)
+    return out
+
+
 grid_1 = [[[0, 0, 1, 0, 0], [0, 2, 1, 1, 1], [1, 2, 1, 0, 0], [1, 2, 0, 1, 0], [2, 2, 0, 0, 2]]]
 grid_2 = [[[0, 0, 0, 0, 1], [0, 0, 1, 1, 1], [1, 0, 1, 1, 2], [1, 0, 1, 2, 2], [0, 0, 0, 0, 1]]]
 grid_3 = [
@@ -326,33 +336,33 @@
 
     TESTS.append(
         [
-            "single_channel_onehot",
-            {"independent": False, "applied_labels": 0, "connectivity": 1, "is_onehot": True},
-            p(grid_5),
-            torch.tensor([[[0, 0, 1, 0, 0], [0, 1, 1, 1, 1], [1, 1, 1, 0, 0], [1, 1, 0, 0, 0], [1, 1, 0, 0, 0]]]),
+            "all_non_zero_labels",
+            {"independent": True},
+            p(grid_1),
+            torch.tensor([[[0, 0, 1, 0, 0], [0, 2, 1, 1, 1], [0, 2, 1, 0, 0], [0, 2, 0, 1, 0], [2, 2, 0, 0, 0]]]),
         ]
     )
 
-INVALID_CASES = []
-for p in TEST_NDARRAYS:
-    INVALID_CASES.append(
-        ["no_applied_labels_for_single_channel", {"independent": False, "is_onehot": False}, p(grid_1), TypeError]
-    )
-    INVALID_CASES.append(["no_applied_labels_for_multi_channel", {"independent": False}, p(grid_3), TypeError])
-
 
 class TestKeepLargestConnectedComponent(unittest.TestCase):
     @parameterized.expand(TESTS)
     def test_correct_results(self, _, args, input_image, expected):
         converter = KeepLargestConnectedComponent(**args)
         result = converter(input_image)
         assert_allclose(result, expected, type_test=False)
-
-    @parameterized.expand(INVALID_CASES)
-    def test_raise_exception(self, _, args, input_image, expected_error):
-        with self.assertRaises(expected_error):
-            converter = KeepLargestConnectedComponent(**args)
-            _ = converter(input_image)
+        if "is_onehot" in args:
+            args["is_onehot"] = not args["is_onehot"]
+        # if not onehotted, onehot it and make sure result stays the same
+        if input_image.shape[0] == 1:
+            img = to_onehot(input_image)
+            result2 = KeepLargestConnectedComponent(**args)(img)
+            result2 = result2.argmax(0)[None]
+            assert_allclose(result, result2)
+        # if onehotted, un-onehot and check result stays the same
+        else:
+            img = input_image.argmax(0)[None]
+            result2 = KeepLargestConnectedComponent(**args)(img)
+            assert_allclose(result.argmax(0)[None], result2)
 
 
 if __name__ == "__main__":

tests/test_keep_largest_connected_componentd.py

@@ -333,20 +333,6 @@
         ]
     )
 
-INVALID_CASES = []
-for p in TEST_NDARRAYS:
-    INVALID_CASES.append(
-        [
-            "no_applied_labels_for_single_channel",
-            {"keys": ["img"], "independent": False, "is_onehot": False},
-            {"img": p(grid_1)},
-            TypeError,
-        ]
-    )
-    INVALID_CASES.append(
-        ["no_applied_labels_for_multi_channel", {"keys": ["img"], "independent": False}, {"img": p(grid_3)}, TypeError]
-    )
-
 
 class TestKeepLargestConnectedComponentd(unittest.TestCase):
     @parameterized.expand(VALID_CASES)
@@ -355,12 +341,6 @@ def test_correct_results(self, _, args, input_dict, expected):
         result = converter(input_dict)
         assert_allclose(result["img"], expected, type_test=False)
 
-    @parameterized.expand(INVALID_CASES)
-    def test_raise_exception(self, _, args, input_dict, expected_error):
-        with self.assertRaises(expected_error):
-            converter = KeepLargestConnectedComponentd(**args)
-            _ = converter(input_dict)
-
 
 if __name__ == "__main__":
     unittest.main()