Keep n components

monai/transforms/post/array.py

@@ -273,6 +273,7 @@ def __init__(
         is_onehot: Optional[bool] = None,
         independent: bool = True,
         connectivity: Optional[int] = None,
+        num_components: int = 1,
     ) -> None:
         """
         Args:
@@ -290,13 +291,15 @@ def __init__(
                 Accepted values are ranging from  1 to input.ndim. If ``None``, a full
                 connectivity of ``input.ndim`` is used. for more details:
                 https://scikit-image.org/docs/dev/api/skimage.measure.html#skimage.measure.label.
+            num_components: The number of largest components to preserve.
 
         """
         super().__init__()
         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
+        self.num_components = num_components
 
     def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
         """
@@ -316,7 +319,7 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
         if self.independent:
             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)
+                mask = get_largest_connected_component_mask(foreground, self.connectivity, self.num_components)
                 if is_onehot:
                     img_[i][foreground != mask] = 0
                 else:
@@ -325,12 +328,12 @@ def __call__(self, img: NdarrayOrTensor) -> NdarrayOrTensor:
         if not is_onehot:  # not one-hot, union of labels
             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)
+            mask = get_largest_connected_component_mask(foreground, self.connectivity, self.num_components)
             img_[0][foreground != mask] = 0
             return convert_to_dst_type(img_, dst=img)[0]
         # one-hot, union of labels
         foreground = (img_[applied_labels, ...] == 1).any(0)
-        mask = get_largest_connected_component_mask(foreground, self.connectivity)
+        mask = get_largest_connected_component_mask(foreground, self.connectivity, self.num_components)
         for i in applied_labels:
             img_[i][foreground != mask] = 0
         return convert_to_dst_type(img_, dst=img)[0]

monai/transforms/post/dictionary.py

@@ -204,6 +204,7 @@ def __init__(
         is_onehot: Optional[bool] = None,
         independent: bool = True,
         connectivity: Optional[int] = None,
+        num_components: int = 1,
         allow_missing_keys: bool = False,
     ) -> None:
         """
@@ -224,12 +225,17 @@ def __init__(
                 Accepted values are ranging from  1 to input.ndim. If ``None``, a full
                 connectivity of ``input.ndim`` is used. for more details:
                 https://scikit-image.org/docs/dev/api/skimage.measure.html#skimage.measure.label.
+            num_components: The number of largest components to preserve.
             allow_missing_keys: don't raise exception if key is missing.
 
         """
         super().__init__(keys, allow_missing_keys)
         self.converter = KeepLargestConnectedComponent(
-            applied_labels=applied_labels, is_onehot=is_onehot, independent=independent, connectivity=connectivity
+            applied_labels=applied_labels,
+            is_onehot=is_onehot,
+            independent=independent,
+            connectivity=connectivity,
+            num_components=num_components,
         )
 
     def __call__(self, data: Mapping[Hashable, NdarrayOrTensor]) -> Dict[Hashable, NdarrayOrTensor]:

monai/transforms/utils.py

@@ -57,9 +57,9 @@
     optional_import,
 )
 from monai.utils.enums import TransformBackends
-from monai.utils.type_conversion import convert_data_type, convert_to_dst_type, convert_to_tensor
+from monai.utils.type_conversion import convert_data_type, convert_to_cupy, convert_to_dst_type, convert_to_tensor
 
-measure, _ = optional_import("skimage.measure", "0.14.2", min_version)
+measure, has_measure = optional_import("skimage.measure", "0.14.2", min_version)
 morphology, has_morphology = optional_import("skimage.morphology")
 ndimage, _ = optional_import("scipy.ndimage")
 cp, has_cp = optional_import("cupy")
@@ -951,7 +951,9 @@ def generate_spatial_bounding_box(
     return box_start, box_end
 
 
-def get_largest_connected_component_mask(img: NdarrayTensor, connectivity: Optional[int] = None) -> NdarrayTensor:
+def get_largest_connected_component_mask(
+    img: NdarrayTensor, connectivity: Optional[int] = None, num_components: int = 1
+) -> NdarrayTensor:
     """
     Gets the largest connected component mask of an image.
 
@@ -961,24 +963,40 @@ def get_largest_connected_component_mask(img: NdarrayTensor, connectivity: Optio
             Accepted values are ranging from  1 to input.ndim. If ``None``, a full
             connectivity of ``input.ndim`` is used. for more details:
             https://scikit-image.org/docs/dev/api/skimage.measure.html#skimage.measure.label.
-    """
-    if isinstance(img, torch.Tensor) and has_cp and has_cucim:
-        x_cupy = monai.transforms.ToCupy()(img.short())
-        x_label = cucim.skimage.measure.label(x_cupy, connectivity=connectivity)
-        vals, counts = cp.unique(x_label[cp.nonzero(x_label)], return_counts=True)
-        comp = x_label == vals[cp.ndarray.argmax(counts)]
-        out_tensor = monai.transforms.ToTensor(device=img.device)(comp)
-        out_tensor = out_tensor.bool()
-
-        return out_tensor  # type: ignore
-
-    img_arr = convert_data_type(img, np.ndarray)[0]
-    largest_cc: np.ndarray = np.zeros(shape=img_arr.shape, dtype=img_arr.dtype)
-    img_arr = measure.label(img_arr, connectivity=connectivity)
-    if img_arr.max() != 0:
-        largest_cc[...] = img_arr == (np.argmax(np.bincount(img_arr.flat)[1:]) + 1)
-
-    return convert_to_dst_type(largest_cc, dst=img, dtype=largest_cc.dtype)[0]
+        num_components: The number of largest components to preserve.
+    """
+    # use skimage/cucim.skimage and np/cp depending on whether packages are
+    # available and input is non-cpu torch.tensor
+    use_cp = has_cp and has_cucim and isinstance(img, torch.Tensor) and img.device != torch.device("cpu")
+    if use_cp:
+        img_ = convert_to_cupy(img.short())  # type: ignore
+        label = cucim.skimage.measure.label
+        lib = cp
+    else:
+        if not has_measure:
+            raise RuntimeError("Skimage.measure required.")
+        img_, *_ = convert_data_type(img, np.ndarray)
+        label = measure.label
+        lib = np
+
+    # features will be an image -- 0 for background and then each different
+    # feature will have its own index.
+    features, num_features = label(img_, connectivity=connectivity, return_num=True)
+    # if num features less than max desired, nothing to do.
+    if num_features <= num_components:
+        out = img_.astype(bool)
+    else:
+        # ignore background
+        nonzeros = features[lib.nonzero(features)]
+        # get number voxels per feature (bincount). argsort[::-1] to get indices
+        # of largest components.
+        features_to_keep = lib.argsort(lib.bincount(nonzeros))[::-1]
+        # only keep the first n non-background indices
+        features_to_keep = features_to_keep[:num_components]
+        # generate labelfield. True if in list of features to keep
+        out = lib.isin(features, features_to_keep)
+
+    return convert_to_dst_type(out, dst=img, dtype=out.dtype)[0]
 
 
 def remove_small_objects(

tests/test_keep_largest_connected_component.py

@@ -78,6 +78,8 @@ def to_onehot(x):
 ]
 grid_5 = [[[0, 0, 1, 0, 0], [0, 1, 1, 1, 1], [1, 1, 1, 0, 0], [1, 1, 0, 1, 0], [1, 1, 0, 0, 1]]]
 
+grid_6 = [[[0, 0, 1, 1, 0, 0, 1], [0, 0, 0, 1, 0, 0, 1], [1, 1, 0, 0, 1, 0, 1], [0, 0, 0, 1, 0, 0, 1]]]
+
 TESTS = []
 for p in TEST_NDARRAYS:
     TESTS.append(
@@ -343,6 +345,37 @@ def to_onehot(x):
             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]]]),
         ]
     )
+    # no connected regions
+    TESTS.append(["0 regions", {"num_components": 0}, p(grid_6), p(torch.zeros(1, 4, 7))])
+    # 1 connected region
+    TESTS.append(
+        [
+            "1 region",
+            {"num_components": 1},
+            p(grid_6),
+            p(
+                torch.tensor(
+                    [[[0, 0, 1, 1, 0, 0, 0], [0, 0, 0, 1, 0, 0, 0], [0, 0, 0, 0, 1, 0, 0], [0, 0, 0, 1, 0, 0, 0]]]
+                )
+            ),
+        ]
+    )
+    # 2 connected regions
+    TESTS.append(
+        [
+            "2 regions",
+            {"num_components": 2},
+            p(grid_6),
+            p(
+                torch.tensor(
+                    [[[0, 0, 1, 1, 0, 0, 1], [0, 0, 0, 1, 0, 0, 1], [0, 0, 0, 0, 1, 0, 1], [0, 0, 0, 1, 0, 0, 1]]]
+                )
+            ),
+        ]
+    )
+    # 3+ connected regions unchanged (as input has 3)
+    for num_connected in (3, 4):
+        TESTS.append([f"{num_connected} regions", {"num_components": num_connected}, p(grid_6), p(grid_6)])
 
 
 class TestKeepLargestConnectedComponent(unittest.TestCase):