diff --git a/cellpose/utils.py b/cellpose/utils.py index ac654960..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: @@ -640,20 +641,21 @@ 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: + msk = np.array([fill_voids.fill(msk[k]) for k in range(msk.shape[0])]) + else: + 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']