From 6c354cafed976b56be0e3f3ca84fb149f9c8e54c Mon Sep 17 00:00:00 2001 From: Soumya Snigdha Kundu Date: Thu, 2 Jul 2026 22:35:08 +0100 Subject: [PATCH] Perf: speed up get_largest_connected_component_mask ranking Rank component sizes with a full-field bincount (zeroing background at index 0) instead of gathering every non-background voxel through lib.nonzero, and build the output mask with a boolean lookup-table gather instead of lib.isin over the whole label field. Both avoid large transient allocations; output is bit-identical. Covers the numpy and cupy/cucim paths unchanged. Signed-off-by: Soumya Snigdha Kundu --- monai/transforms/utils.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/monai/transforms/utils.py b/monai/transforms/utils.py index 2ca94617f3..377d2b60e9 100644 --- a/monai/transforms/utils.py +++ b/monai/transforms/utils.py @@ -1224,15 +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] - # 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) + # bincount counts every label; index 0 is background, so drop it before ranking + counts = lib.bincount(features.reshape(-1)) + counts[0] = 0 + # argsort[::-1] gives labels of the largest components; keep the first n + features_to_keep = lib.argsort(counts)[::-1][:num_components] + # boolean lookup-table gather over the label field, cheaper than isin + keep = lib.zeros(counts.shape[0], dtype=bool) + keep[features_to_keep] = True + out = keep[features] return convert_to_dst_type(out, dst=img, dtype=out.dtype)[0]