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]