fix(utils): fixes a bug in fill_holes_and_remove_small_masks

cellpose/utils.py

@@ -642,10 +642,15 @@ def fill_holes_and_remove_small_masks(masks, min_size=15):
 
     # Filter small masks
     if min_size > 0:
-        counts = fastremap.unique(masks, return_counts=True)[1][1:]
-        masks = fastremap.mask(masks, np.nonzero(counts < min_size)[0] + 1)
+        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)
-        
+
     slices = find_objects(masks)
     j = 0
     for i, slc in enumerate(slices):
@@ -656,8 +661,11 @@ def fill_holes_and_remove_small_masks(masks, min_size=15):
             j += 1
 
     if min_size > 0:
-        counts = fastremap.unique(masks, return_counts=True)[1][1:]
-        masks = fastremap.mask(masks, np.nonzero(counts < min_size)[0] + 1)
+        uniq, counts = fastremap.unique(masks, return_counts=True)
+        small_mask_indices = np.nonzero(counts[1:] < min_size)[0]
+        labels_to_remove = uniq[1:][small_mask_indices]
+        masks = fastremap.mask(masks, labels_to_remove)
         fastremap.renumber(masks, in_place=True)
 
+
     return masks

tests/test_utils.py

@@ -0,0 +1,23 @@
+import numpy as np
+from cellpose.utils import fill_holes_and_remove_small_masks
+import fastremap
+
+
+def test_fill_holes_and_remove_small_masks():
+    # make a 2-channel mask with holes and small objects. The first channel is the
+    # "ground truth" without holes or small objects, the second channel needs to be cleaned.
+    masks = np.zeros((2, 100, 100), dtype=np.uint16)
+    masks[:, 10:30, 10:30] = 1  # object 1
+    masks[1, 15:25, 15:25] = 0  # hole in object 1
+    masks[1, 40:45, 40:45] = 2  # small object 2
+    masks[:, 60:90, 60:90] = 4  # object 4 (skip 3)
+    masks[1, 70:80, 70:80] = 0  # hole in object 4
+    masks[1, 10:15, 80:82] = 5  # small object 4
+
+    # apply function
+    min_size = 30
+    masks_cleaned = fill_holes_and_remove_small_masks(masks[1], min_size=min_size)
+
+    gt_masks = fastremap.renumber(masks[0], in_place=False)[0]
+
+    assert (gt_masks == masks_cleaned).all()