From 441cef2fffd95bc049349b5bf763b87e5c1bd4ce Mon Sep 17 00:00:00 2001 From: Vishnu Kannaujia Date: Fri, 17 Jul 2026 17:34:28 -0700 Subject: [PATCH] Count components via bincount over full field in get_largest_connected_component_mask Fixes #8974. Rank component sizes with a single bincount over the whole label field and zero out the background bin, instead of gathering every foreground voxel via features[nonzero(features)] first. This removes the per-spatial-dim index arrays produced by nonzero and the full gathered copy of the non-background voxels, lowering peak memory and runtime. Output is bit-identical: background is label 0, all foreground labels keep their original counts, and the top num_components labels selected by argsort are unchanged. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Vishnu Kannaujia --- monai/transforms/utils.py | 14 +++++---- .../test_keep_largest_connected_component.py | 29 +++++++++++++++++++ 2 files changed, 38 insertions(+), 5 deletions(-) diff --git a/monai/transforms/utils.py b/monai/transforms/utils.py index 2ca94617f3..1fdfabda14 100644 --- a/monai/transforms/utils.py +++ b/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 diff --git a/tests/transforms/test_keep_largest_connected_component.py b/tests/transforms/test_keep_largest_connected_component.py index 639cdf9936..8905ad8c2a 100644 --- a/tests/transforms/test_keep_largest_connected_component.py +++ b/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()