diff --git a/monai/transforms/post/array.py b/monai/transforms/post/array.py index dcd32c5d55..6396435aa7 100644 --- a/monai/transforms/post/array.py +++ b/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,7 +302,7 @@ 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, @@ -310,8 +310,8 @@ def __init__( """ 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 diff --git a/monai/transforms/post/dictionary.py b/monai/transforms/post/dictionary.py index b89196207b..00ffe7edf7 100644 --- a/monai/transforms/post/dictionary.py +++ b/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. diff --git a/monai/transforms/utils.py b/monai/transforms/utils.py index 4fb36fabe4..847614adfe 100644 --- a/monai/transforms/utils.py +++ b/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) diff --git a/monai/transforms/utils_pytorch_numpy_unification.py b/monai/transforms/utils_pytorch_numpy_unification.py index 25cb1455dd..2103ccff58 100644 --- a/monai/transforms/utils_pytorch_numpy_unification.py +++ b/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 diff --git a/tests/test_get_unique_labels.py b/tests/test_get_unique_labels.py new file mode 100644 index 0000000000..9bc6f9b152 --- /dev/null +++ b/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() diff --git a/tests/test_keep_largest_connected_component.py b/tests/test_keep_largest_connected_component.py index fc62f0f0e9..87a5a81b75 100644 --- a/tests/test_keep_largest_connected_component.py +++ b/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,20 +336,13 @@ 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) @@ -347,12 +350,19 @@ 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__": diff --git a/tests/test_keep_largest_connected_componentd.py b/tests/test_keep_largest_connected_componentd.py index a103ac06aa..a06fb51a97 100644 --- a/tests/test_keep_largest_connected_componentd.py +++ b/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()