Skip to content

Perf: faster get_largest_connected_component_mask (bincount + LUT gather) #8978

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
aymuos15 wants to merge 2 commits into
base: dev
Choose a base branch
from
+9 −9
Open

Perf: faster get_largest_connected_component_mask (bincount + LUT gather) #8978

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6c354ca
Perf: speed up get_largest_connected_component_mask ranking
aymuos15 Jul 2, 2026
5968dfc
Merge branch 'dev' into perf/largest-cc-bincount-lut
ericspod Jul 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 9 additions & 9 deletions monai/transforms/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Comment on lines +1227 to +1235

@coderabbitai coderabbitai Bot Jul 2, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 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 200

Repository: 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.py

Repository: 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))
PY

Repository: 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))
PY

Repository: 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


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

Expand Down
Loading