From eea0466a65220da8af16adf8244365df1d2a68ea Mon Sep 17 00:00:00 2001
From: Tom van Leeuwen <tom@tleeuwen.nl>
Date: Fri, 21 Feb 2025 13:10:00 +0100
Subject: [PATCH 1/2] Use np.bincount and np.isin to filter small masks before
 filling holes

---
 cellpose/utils.py | 24 +++++++++++++-----------
 1 file changed, 13 insertions(+), 11 deletions(-)

diff --git a/cellpose/utils.py b/cellpose/utils.py
index ac654960..11959d98 100644
--- a/cellpose/utils.py
+++ b/cellpose/utils.py
@@ -640,20 +640,22 @@ def fill_holes_and_remove_small_masks(masks, min_size=15):
         raise ValueError("masks_to_outlines takes 2D or 3D array, not %dD array" %
                          masks.ndim)
 
+    # Filter small masks
+    if min_size > 0:
+        counts = np.bincount(masks.ravel())
+        filter = np.isin(masks, np.where(counts < min_size)[0])
+        masks[filter] = 0
+
     slices = find_objects(masks)
     j = 0
     for i, slc in enumerate(slices):
         if slc is not None:
             msk = masks[slc] == (i + 1)
-            npix = msk.sum()
-            if min_size > 0 and npix < min_size:
-                masks[slc][msk] = 0
-            elif npix > 0:
-                if msk.ndim == 3:
-                    for k in range(msk.shape[0]):
-                        msk[k] = binary_fill_holes(msk[k])
-                else:
-                    msk = binary_fill_holes(msk)
-                masks[slc][msk] = (j + 1)
-                j += 1
+            if msk.ndim == 3:
+                for k in range(msk.shape[0]):
+                    msk[k] = binary_fill_holes(msk[k])
+            else:
+                msk = binary_fill_holes(msk)
+            masks[slc][msk] = (j + 1)
+            j += 1
     return masks

From 7d916ca5dbce0c31f0dff3961155957e2ec136ee Mon Sep 17 00:00:00 2001
From: Tom van Leeuwen <tom@tleeuwen.nl>
Date: Fri, 21 Feb 2025 13:16:05 +0100
Subject: [PATCH 2/2] Replace scipy.ndimage.morphology.binary_fill_holes with
 fill_voids.fill, a more optimized function for filling in 2D and 3D contours

---
 cellpose/utils.py | 8 ++++----
 environment.yml   | 1 +
 setup.py          | 1 +
 3 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/cellpose/utils.py b/cellpose/utils.py
index 11959d98..cd42495e 100644
--- a/cellpose/utils.py
+++ b/cellpose/utils.py
@@ -11,6 +11,7 @@
 import numpy as np
 import colorsys
 import fastremap
+import fill_voids
 from multiprocessing import Pool, cpu_count
 
 from . import metrics
@@ -619,7 +620,7 @@ def size_distribution(masks):
 def fill_holes_and_remove_small_masks(masks, min_size=15):
     """ Fills holes in masks (2D/3D) and discards masks smaller than min_size.
 
-    This function fills holes in each mask using scipy.ndimage.morphology.binary_fill_holes.
+    This function fills holes in each mask using fill_voids.fill.
     It also removes masks that are smaller than the specified min_size.
 
     Parameters:
@@ -652,10 +653,9 @@ def fill_holes_and_remove_small_masks(masks, min_size=15):
         if slc is not None:
             msk = masks[slc] == (i + 1)
             if msk.ndim == 3:
-                for k in range(msk.shape[0]):
-                    msk[k] = binary_fill_holes(msk[k])
+                msk = np.array([fill_voids.fill(msk[k]) for k in range(msk.shape[0])])
             else:
-                msk = binary_fill_holes(msk)
+                msk = fill_voids.fill(msk)
             masks[slc][msk] = (j + 1)
             j += 1
     return masks
diff --git a/environment.yml b/environment.yml
index 813664a1..ac97d053 100644
--- a/environment.yml
+++ b/environment.yml
@@ -26,4 +26,5 @@ dependencies:
     - zarr
     - dask_jobqueue
     - bokeh
+    - fill-voids
   
diff --git a/setup.py b/setup.py
index e027cb82..f8df2aa5 100644
--- a/setup.py
+++ b/setup.py
@@ -14,6 +14,7 @@
     'fastremap',
     'imagecodecs',
     'roifile',
+    'fill-voids',
 ]
 
 image_deps = ['nd2', 'pynrrd']