Count components via bincount over full field in get_largest_connected_component_mask

[vishnukannaujia]

Fixes #8974.

Summary

get_largest_connected_component_mask (monai/transforms/utils.py) ranked component sizes with:

nonzeros = features[lib.nonzero(features)]
features_to_keep = lib.argsort(lib.bincount(nonzeros))[::-1][:num_components]

lib.nonzero(features) materializes one full index array per spatial dim, and features[...] then gathers a full copy of every non-background voxel — all purely to drop background before bincount. That is unnecessary: a bincount over the whole field already counts every label, and background is label 0.

This PR counts directly over the field and masks background out:

counts = lib.bincount(features.reshape(-1))
counts[0] = 0  # ignore background
features_to_keep = lib.argsort(counts)[::-1][:num_components]

This avoids both the nonzero index arrays and the gathered copy, so it lowers peak memory and runtime on large volumes. The change applies to both the NumPy and CuPy (lib) paths.

Why the output is unchanged (bit-identical)

• Background is label 0; setting counts[0] = 0 guarantees it is never selected.
• Every foreground label keeps the exact same voxel count as before (the old code excluded background from bincount, which only affected bin 0, whose count was already 0).
• This branch only runs when num_features > num_components, so all foreground labels have count >= 1 and the num_components largest selected by argsort(...)[::-1] are identical. isin membership is order-independent.

Tests

Existing test_keep_largest_connected_component[d] suites (174 cases) still pass. Added direct get_largest_connected_component_mask tests across TEST_NDARRAYS backends:

• test_num_components_selection: a field with four blobs (6/4/3/1 voxels) keeps exactly the largest num_components for n = 1..4.
• test_background_never_selected: background dominates by voxel count but is never kept.

[coderabbitai]

📝 Walkthrough

Walkthrough

get_largest_connected_component_mask now uses bincount directly on the flattened labeled feature array, zeroes the background count, and selects the largest foreground components without intermediate nonzero-based arrays. Tests cover selection of the largest components regardless of label order and confirm that background label 0 is excluded.

Estimated code review effort: 2 (Simple) | ~10 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly describes the main change: counting connected components with bincount over the full field.
Description check	✅ Passed	The description includes the fix, rationale, unchanged-output argument, and tests; only the checklist section is omitted.
Linked Issues check	✅ Passed	The code implements issue #8974’s full-field bincount approach, zeros background, and preserves largest-component selection.
Out of Scope Changes check	✅ Passed	Only the targeted utility and tests changed; no unrelated or out-of-scope edits are present.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (2)

monai/transforms/utils.py (2)

1231-1231: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low value

Prefer .ravel() over .reshape(-1).

Use .ravel() to flatten the array. It is more idiomatic and explicitly conveys the intent to return a flattened view.

♻️ Proposed refactor

-        counts = lib.bincount(features.reshape(-1))
+        counts = lib.bincount(features.ravel())

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@monai/transforms/utils.py` at line 1231, Update the flattening expression in
the counts computation to use features.ravel() instead of features.reshape(-1),
preserving the existing lib.bincount behavior.

---

1194-1204: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Add Returns and Raises sections to the docstring.

As per path instructions, describe the return value and raised exceptions using the appropriate sections of the Google-style docstrings.

📝 Proposed docstring update

     """
     Gets the largest connected component mask of an image.
 
     Args:
         img: Image to get largest connected component from. Shape is (spatial_dim1 [, spatial_dim2, ...])
         connectivity: Maximum number of orthogonal hops to consider a pixel/voxel as a neighbor.
             Accepted values are ranging from  1 to input.ndim. If ``None``, a full
             connectivity of ``input.ndim`` is used. for more details:
             https://scikit-image.org/docs/dev/api/skimage.measure.html#skimage.measure.label.
         num_components: The number of largest components to preserve.
+
+    Returns:
+        A boolean mask tensor or array of the same shape as `img` containing the largest connected components.
+
+    Raises:
+        RuntimeError: If `skimage.measure` is required but not installed.
     """

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@monai/transforms/utils.py` around lines 1194 - 1204, Update the docstring for
the largest connected component utility to add Google-style Returns and Raises
sections, documenting the returned component mask and the exceptions raised for
invalid inputs or processing failures. Keep the existing Args documentation
unchanged.

Source: Path instructions

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@monai/transforms/utils.py`:
- Line 1231: Update the flattening expression in the counts computation to use
features.ravel() instead of features.reshape(-1), preserving the existing
lib.bincount behavior.
- Around line 1194-1204: Update the docstring for the largest connected
component utility to add Google-style Returns and Raises sections, documenting
the returned component mask and the exceptions raised for invalid inputs or
processing failures. Keep the existing Args documentation unchanged.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 0aec11f9-d860-46fd-aaae-bd29eefc365f

📥 Commits

Reviewing files that changed from the base of the PR and between 5da2472 and 441cef2.

📒 Files selected for processing (2)

• monai/transforms/utils.py
• tests/transforms/test_keep_largest_connected_component.py