Perf: faster get_largest_connected_component_mask (bincount + LUT gather)#8978
Perf: faster get_largest_connected_component_mask (bincount + LUT gather)#8978aymuos15 wants to merge 1 commit into
Conversation
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 <soumya_snigdha.kundu@kcl.ac.uk>
📝 WalkthroughWalkthroughThe change refactors get_largest_connected_component_mask in monai/transforms/utils.py to compute label voxel counts via bincount on flattened features, zero out the background count, rank labels by descending count via argsort, and build a boolean keep lookup table indexed directly into the feature map, replacing the prior nonzero-based extraction and isin masking approach. Estimated code review effort: 2 (Simple) | ~10 minutes Changes
Related issues: Addresses Related PRs: None noted. Suggested labels: performance, transforms Suggested reviewers: MONAI transforms maintainers Poem: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
monai/transforms/utils.py (1)
1227-1235: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winUse
argpartitioninstead of fullargsortfor top-k selection.
lib.argsort(counts)[::-1][:num_components]sorts the entire counts array just to pick the topnum_components. Since only the largest few are needed,argpartitionavoids the full O(n log n) sort and better matches this PR's own perf-optimization goal.♻️ Proposed refactor
- # argsort[::-1] gives labels of the largest components; keep the first n - features_to_keep = lib.argsort(counts)[::-1][:num_components] + # argpartition avoids a full sort when only the top-n labels are needed + features_to_keep = lib.argpartition(counts, -num_components)[-num_components:]🤖 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 1227 - 1235, The top-k selection in the label-count ranking is doing a full sort via lib.argsort(counts)[::-1][:num_components], which is unnecessary for selecting only the largest components. Update the selection logic in the component-ranking code around the bincount/counts handling to use lib.argpartition for the top num_components labels, then keep the existing boolean lookup-table flow (keep/features_to_keep/out) to preserve behavior while improving performance.
🤖 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.
Inline comments:
In `@monai/transforms/utils.py`:
- Around line 1227-1235: The KeepLargestConnectedComponent ranking path in
utils.py is not covered for equal-sized connected components, so the current
argsort-based selection can change silently. Add a regression test for the
KeepLargestConnectedComponent behavior with tie-sized components that exercises
the ranking branch and asserts the chosen components remain deterministic; use
the same multi-component setup already used in the existing connected-component
tests to locate the relevant case.
---
Nitpick comments:
In `@monai/transforms/utils.py`:
- Around line 1227-1235: The top-k selection in the label-count ranking is doing
a full sort via lib.argsort(counts)[::-1][:num_components], which is unnecessary
for selecting only the largest components. Update the selection logic in the
component-ranking code around the bincount/counts handling to use
lib.argpartition for the top num_components labels, then keep the existing
boolean lookup-table flow (keep/features_to_keep/out) to preserve behavior while
improving performance.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 26c10d81-e6f4-4fe3-b7ad-e8038af614ce
📒 Files selected for processing (1)
monai/transforms/utils.py
| # 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] |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Find tests covering get_largest_connected_component_mask with multi-component scenarios
rg -n "get_largest_connected_component_mask|num_components" --type=py -g '*test*'Repository: Project-MONAI/MONAI
Length of output: 157
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== Implementation locations ==\n'
rg -n "def get_largest_connected_component_mask|get_largest_connected_component_mask\(" monai -g '*.py'
printf '\n== Candidate test mentions ==\n'
rg -n "largest_connected_component|connected_component|num_components|bincount|argsort|isin" tests monai -g '*test*.py' -g '*.py' || true
printf '\n== File inventory around tests ==\n'
fd -a -t f 'test.*\.py$|.*_test\.py$' tests monai 2>/dev/null | head -n 200Repository: Project-MONAI/MONAI
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== Implementation slice ==\n'
sed -n '1191,1260p' monai/transforms/utils.py
printf '\n== Test file slice ==\n'
sed -n '300,430p' tests/transforms/test_keep_largest_connected_component.pyRepository: Project-MONAI/MONAI
Length of output: 7866
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n== Test file top ==\n'
sed -n '1,140p' tests/transforms/test_keep_largest_connected_component.py
printf '\n== grid_6 definition and nearby fixtures ==\n'
rg -n "grid_6|grid_4|grid_1" tests/transforms/test_keep_largest_connected_component.py
python3 - <<'PY'
import ast, pathlib
path = pathlib.Path('tests/transforms/test_keep_largest_connected_component.py')
src = path.read_text()
tree = ast.parse(src)
for node in ast.walk(tree):
if isinstance(node, ast.Assign):
for t in node.targets:
if isinstance(t, ast.Name) and t.id == 'grid_6':
print('\n== Parsed grid_6 value ==')
print(ast.literal_eval(node.value))
PYRepository: Project-MONAI/MONAI
Length of output: 5570
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from collections import deque
grid = [[0, 0, 1, 1, 0, 0, 1],
[0, 0, 0, 1, 0, 0, 1],
[1, 1, 0, 0, 1, 0, 1],
[0, 0, 0, 1, 0, 0, 1]]
H, W = len(grid), len(grid[0])
for conn in (4, 8):
seen = [[False]*W for _ in range(H)]
sizes = []
for i in range(H):
for j in range(W):
if grid[i][j] == 0 or seen[i][j]:
continue
q = deque([(i,j)])
seen[i][j] = True
size = 0
while q:
x,y = q.popleft()
size += 1
for dx,dy in ([(-1,0),(1,0),(0,-1),(0,1)] if conn==4 else
[(-1,-1),(-1,0),(-1,1),(0,-1),(0,1),(1,-1),(1,0),(1,1)]):
nx, ny = x+dx, y+dy
if 0 <= nx < H and 0 <= ny < W and not seen[nx][ny] and grid[nx][ny] != 0:
seen[nx][ny] = True
q.append((nx, ny))
sizes.append(size)
print(conn, sorted(sizes, reverse=True), "num_components=", len(sizes))
PYRepository: Project-MONAI/MONAI
Length of output: 223
Add a tie-case regression for KeepLargestConnectedComponent. The current tests already hit the ranking branch with multi-component inputs, but none cover equal-sized components, so the new argsort ordering stays unpinned.
🤖 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 1227 - 1235, The
KeepLargestConnectedComponent ranking path in utils.py is not covered for
equal-sized connected components, so the current argsort-based selection can
change silently. Add a regression test for the KeepLargestConnectedComponent
behavior with tie-sized components that exercises the ranking branch and asserts
the chosen components remain deterministic; use the same multi-component setup
already used in the existing connected-component tests to locate the relevant
case.
Source: Path instructions
1 participant
Fixes #8974 .
Description
get_largest_connected_component_mask(monai/transforms/utils.py) ranked component sizes by gathering every non-background voxel throughlib.nonzero(features)and then built the mask withlib.isinover the whole label field. Both allocate large transient arrays. This counts labels with a single full-fieldlib.bincount(zeroing index 0 to drop background), and builds the mask with a boolean lookup-table gather (keep[features]) instead ofisin. 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
isinreplacement is the dominant win on both time and peak memory; thebincountchange removes the redundantnonzeroindex arrays and gathered copy on top.Types of changes