fix(utils): fixes a bug in fill_holes_and_remove_small_masks

[rubencardenes]

Bug Fix: fill_holes_and_remove_small_masks - Incorrect Label Removal

Problem Identified

When running celposeSAM with different cellprob_threshold I found that with lower thresholds I was missing cells, which was quite odd and should not be the case. The problem seems to be related the function fill_holes_and_remove_small_masks in cellpose/utils.py, that I think has a critical bug when removing small masks.

Here an example with two thresholds where some cells are missing for cellprob_threshold = -5, both using min_size=500 and flow_threshold=0.95

cellprob_threshold = 0

cellprob_threshold = -5 (see that some obvious cells are missing)

The Buggy Code seems to be in (lines 645-647):

counts = fastremap.unique(masks, return_counts=True)[1][1:]
masks = fastremap.mask(masks, np.nonzero(counts < min_size)[0] + 1)
fastremap.renumber(masks, in_place=True)

The code assumed mask labels are always sequential (1, 2, 3, ...), but this assumption is violated when remove_bad_flow_masks() in dynamics.py removes masks with poor flow quality.

Example scenario:

• After get_masks_torch: labels = [1, 2, 3, 4]
• After remove_bad_flow_masks removes mask 2: labels = [1, 3, 4] (non-sequential!)
• fastremap.unique() returns counts in order of labels: [count_1, count_3, count_4]
• Finding small mask at index 1 (label 3 is small)
• Buggy code: index + 1 = 1 + 1 = 2 → tries to remove label 2 (which doesn't exist!)
• Should remove: label 3

Impact

• Wrong masks could be removed when labels are non-sequential
• Correct small masks might not be removed if the wrong index is used
• This could lead to incorrect cell counts and missed small debris removal

Proposed Fix

uniq, counts = fastremap.unique(masks, return_counts=True)
# uniq[0] is background (0), so uniq[1:] are the actual mask labels
# counts[1:] are the corresponding counts
small_mask_indices = np.nonzero(counts[1:] < min_size)[0]
# Get the actual label values to remove (not indices)
labels_to_remove = uniq[1:][small_mask_indices]
masks = fastremap.mask(masks, labels_to_remove)
fastremap.renumber(masks, in_place=True)

This fix is applied to both filtering locations (before and after fill_voids).

Example image with cellprob_threshold = -5 after the fix

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 42.39%. Comparing base (2932984) to head (8b3930e).
⚠️ Report is 7 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1367      +/-   ##
==========================================
+ Coverage   42.35%   42.39%   +0.04%     
==========================================
  Files          16       16              
  Lines        3778     3781       +3     
==========================================
+ Hits         1600     1603       +3     
  Misses       2178     2178              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[mrariden]

Look good @rubencardenes thanks for pointing this out. I added a test that caught the incorrect behavior before the fix, and passed after the fix was implemented.

Will merge into main when tests are passing

• GH workflows (partial)
• Full tests on Ubuntu
• Mac (partial)