Count components via bincount over full field in get_largest_connected_component_mask

monai/transforms/utils.py

@@ -1224,11 +1224,15 @@ def get_largest_connected_component_mask(
     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]
+        # count voxels per feature directly over the whole field; background is
+        # label 0. This avoids materializing the ``nonzero`` index arrays (one per
+        # spatial dim) and the gathered copy of every foreground voxel that
+        # ``features[lib.nonzero(features)]`` would create.
+        counts = lib.bincount(features.reshape(-1))
+        # ignore background so it is never selected as a component to keep.
+        counts[0] = 0
+        # argsort[::-1] to get indices of the largest components.
+        features_to_keep = lib.argsort(counts)[::-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

tests/transforms/test_keep_largest_connected_component.py

@@ -19,6 +19,7 @@
 from parameterized import parameterized
 
 from monai.transforms import KeepLargestConnectedComponent
+from monai.transforms.utils import get_largest_connected_component_mask
 from monai.transforms.utils_pytorch_numpy_unification import moveaxis
 from monai.utils.type_conversion import convert_to_dst_type
 from tests.test_utils import TEST_NDARRAYS, assert_allclose
@@ -415,5 +416,33 @@ def test_correct_results_before_after_onehot(self, _, args, input_image, expecte
             assert_allclose(result.argmax(0)[None], result2, type_test="tensor")
 
 
+class TestGetLargestConnectedComponentMask(unittest.TestCase):
+    @parameterized.expand([(p,) for p in TEST_NDARRAYS])
+    def test_num_components_selection(self, in_type):
+        # a field with four separate foreground blobs of sizes 6, 4, 3 and 1
+        # voxels (plus background). Selecting the ``num_components`` largest must
+        # keep exactly the largest ``num_components`` blobs, regardless of label
+        # order. This guards the bincount-based ranking in
+        # ``get_largest_connected_component_mask``.
+        img = in_type(
+            torch.tensor(
+                [[1, 1, 0, 2, 2, 0], [1, 1, 0, 2, 2, 0], [1, 1, 0, 0, 0, 0], [0, 0, 0, 3, 3, 0], [4, 0, 0, 3, 0, 0]]
+            )
+        )
+
+        # sizes: blob A=6, B=4, C=3, D=1 -> expected kept voxels for each n
+        expected_counts = {1: 6, 2: 10, 3: 13, 4: 14}
+        for num_components, expected in expected_counts.items():
+            mask = get_largest_connected_component_mask(img, num_components=num_components)
+            assert_allclose(mask.sum(), torch.tensor(expected), type_test=False)
+
+    @parameterized.expand([(p,) for p in TEST_NDARRAYS])
+    def test_background_never_selected(self, in_type):
+        # background (label 0) dominates by voxel count but must never be kept.
+        img = in_type(torch.tensor([[0, 0, 0, 0], [0, 1, 1, 0], [0, 0, 0, 0]]))
+        mask = get_largest_connected_component_mask(img, num_components=1)
+        assert_allclose(mask.sum(), torch.tensor(2), type_test=False)
+
+
 if __name__ == "__main__":
     unittest.main()