Perf: faster get_largest_connected_component_mask (bincount + LUT gather)

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]