Skip to content

Commit a3d5160

Browse files
aymuos15ericspod
andauthored
Perf: faster get_largest_connected_component_mask (bincount + LUT gather) (#8978)
Fixes #8974 . ### Description `get_largest_connected_component_mask` (`monai/transforms/utils.py`) ranked component sizes by gathering every non-background voxel through `lib.nonzero(features)` and then built the mask with `lib.isin` over the whole label field. Both allocate large transient arrays. This counts labels with a single full-field `lib.bincount` (zeroing index 0 to drop background), and builds the mask with a boolean lookup-table gather (`keep[features]`) instead of `isin`. Output is bit-identical, verified across 2D and 3D at several foreground fractions and component counts, and the numpy and cupy/cucim paths are covered unchanged. The `isin` replacement is the dominant win on both time and peak memory; the `bincount` change removes the redundant `nonzero` index arrays and gathered copy on top. | case | metric | before | after | improvement | |---|---|---|---|---| | 3D 192^3, fg 50% | time | 82.0 ms | 8.6 ms | 9.5x | | 3D 192^3, fg 50% | peak transient mem | 108.0 MB | 7.2 MB | 15x | | 2D 1024^2, fg 60% | time | 21.5 ms | 1.8 ms | 11.7x | | 2D 1024^2, fg 60% | peak transient mem | 18.0 MB | 1.4 MB | 12.9x | | 3D 128^3, fg 5% (48k comps) | time | 10.7 ms | 6.2 ms | 1.7x | ### Types of changes - [x] Non-breaking change (fix or new feature that would not break existing functionality). Signed-off-by: Soumya Snigdha Kundu <soumya_snigdha.kundu@kcl.ac.uk> Co-authored-by: Eric Kerfoot <17726042+ericspod@users.noreply.github.com>
1 parent 8885fcb commit a3d5160

1 file changed

Lines changed: 9 additions & 9 deletions

File tree

monai/transforms/utils.py

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1224,15 +1224,15 @@ def get_largest_connected_component_mask(
12241224
if num_features <= num_components:
12251225
out = img_.astype(bool)
12261226
else:
1227-
# ignore background
1228-
nonzeros = features[lib.nonzero(features)]
1229-
# get number voxels per feature (bincount). argsort[::-1] to get indices
1230-
# of largest components.
1231-
features_to_keep = lib.argsort(lib.bincount(nonzeros))[::-1]
1232-
# only keep the first n non-background indices
1233-
features_to_keep = features_to_keep[:num_components]
1234-
# generate labelfield. True if in list of features to keep
1235-
out = lib.isin(features, features_to_keep)
1227+
# bincount counts every label; index 0 is background, so drop it before ranking
1228+
counts = lib.bincount(features.reshape(-1))
1229+
counts[0] = 0
1230+
# argsort[::-1] gives labels of the largest components; keep the first n
1231+
features_to_keep = lib.argsort(counts)[::-1][:num_components]
1232+
# boolean lookup-table gather over the label field, cheaper than isin
1233+
keep = lib.zeros(counts.shape[0], dtype=bool)
1234+
keep[features_to_keep] = True
1235+
out = keep[features]
12361236

12371237
return convert_to_dst_type(out, dst=img, dtype=out.dtype)[0]
12381238

0 commit comments

Comments
 (0)