adding max_size_fraction option (#796)

cellpose/dynamics.py

@@ -663,7 +663,7 @@ def mem_info():
     return masks
 
 
-def get_masks(p, iscell=None, rpad=20):
+def get_masks(p, iscell=None, rpad=20, max_size_fraction=0.4):
     """Create masks using pixel convergence after running dynamics.
 
     Makes a histogram of final pixel locations p, initializes masks 
@@ -677,6 +677,8 @@ def get_masks(p, iscell=None, rpad=20):
         iscell (bool, 2D or 3D array): If iscell is not None, set pixels that are 
             iscell False to stay in their original location.
         rpad (int, optional): Histogram edge padding. Default is 20.
+        max_size_fraction (float, optional): Masks larger than max_size_fraction of
+            total image size are removed. Default is 0.4.
 
     Returns:
         M0 (int, 2D or 3D array): Masks with inconsistent flow masks removed, 
@@ -750,7 +752,7 @@ def get_masks(p, iscell=None, rpad=20):
 
     # remove big masks
     uniq, counts = fastremap.unique(M0, return_counts=True)
-    big = np.prod(shape0) * 0.4
+    big = np.prod(shape0) * max_size_fraction
     bigc = uniq[counts > big]
     if len(bigc) > 0 and (len(bigc) > 1 or bigc[0] != 0):
         M0 = fastremap.mask(M0, bigc)
@@ -761,7 +763,7 @@ def get_masks(p, iscell=None, rpad=20):
 
 def resize_and_compute_masks(dP, cellprob, p=None, niter=200, cellprob_threshold=0.0,
                              flow_threshold=0.4, interp=True, do_3D=False, min_size=15,
-                             resize=None, device=None):
+                             max_size_fraction=0.4, resize=None, device=None):
     """Compute masks using dynamics from dP and cellprob, and resizes masks if resize is not None.
 
     Args:
@@ -774,6 +776,8 @@ def resize_and_compute_masks(dP, cellprob, p=None, niter=200, cellprob_threshold
         interp (bool, optional): Whether to interpolate during dynamics computation. Defaults to True.
         do_3D (bool, optional): Whether to perform mask computation in 3D. Defaults to False.
         min_size (int, optional): The minimum size of the masks. Defaults to 15.
+        max_size_fraction (float, optional): Masks larger than max_size_fraction of
+            total image size are removed. Default is 0.4.
         resize (tuple, optional): The desired size for resizing the masks. Defaults to None.
         device (str, optional): The torch device to use for computation. Defaults to None.
 
@@ -783,7 +787,8 @@ def resize_and_compute_masks(dP, cellprob, p=None, niter=200, cellprob_threshold
     mask, p = compute_masks(dP, cellprob, p=p, niter=niter,
                             cellprob_threshold=cellprob_threshold,
                             flow_threshold=flow_threshold, interp=interp, do_3D=do_3D,
-                            min_size=min_size, device=device)
+                            min_size=min_size, max_size_fraction=max_size_fraction, 
+                            device=device)
 
     if resize is not None:
         mask = transforms.resize_image(mask, resize[0], resize[1],
@@ -798,7 +803,7 @@ def resize_and_compute_masks(dP, cellprob, p=None, niter=200, cellprob_threshold
 
 def compute_masks(dP, cellprob, p=None, niter=200, cellprob_threshold=0.0,
                   flow_threshold=0.4, interp=True, do_3D=False, min_size=15,
-                  device=None):
+                  max_size_fraction=0.4, device=None):
     """Compute masks using dynamics from dP and cellprob.
 
     Args:
@@ -811,6 +816,8 @@ def compute_masks(dP, cellprob, p=None, niter=200, cellprob_threshold=0.0,
         interp (bool, optional): Whether to interpolate during dynamics computation. Defaults to True.
         do_3D (bool, optional): Whether to perform mask computation in 3D. Defaults to False.
         min_size (int, optional): The minimum size of the masks. Defaults to 15.
+        max_size_fraction (float, optional): Masks larger than max_size_fraction of
+            total image size are removed. Default is 0.4.
         device (str, optional): The torch device to use for computation. Defaults to None.
 
     Returns:
@@ -831,7 +838,7 @@ def compute_masks(dP, cellprob, p=None, niter=200, cellprob_threshold=0.0,
                 return mask, p
 
         #calculate masks
-        mask = get_masks(p, iscell=cp_mask)
+        mask = get_masks(p, iscell=cp_mask, max_size_fraction=max_size_fraction)
 
         # flow thresholding factored out of get_masks
         if not do_3D:

cellpose/models.py

@@ -140,7 +140,7 @@ def __init__(self, gpu=False, model_type="cyto3", nchan=2, device=None,
         self.sz.model_type = model_type
 
     def eval(self, x, batch_size=8, channels=[0, 0], channel_axis=None, invert=False,
-             normalize=True, diameter=30., do_3D=False, find_masks=True, **kwargs):
+             normalize=True, diameter=30., do_3D=False, **kwargs):
         """Run cellpose size model and mask model and get masks.
 
         Args:
@@ -353,9 +353,9 @@ def __init__(self, gpu=False, pretrained_model=False, model_type=None,
     def eval(self, x, batch_size=8, resample=True, channels=None, channel_axis=None,
              z_axis=None, normalize=True, invert=False, rescale=None, diameter=None,
              flow_threshold=0.4, cellprob_threshold=0.0, do_3D=False, anisotropy=None,
-             stitch_threshold=0.0, min_size=15, niter=None, augment=False, tile=True,
-             tile_overlap=0.1, bsize=224, interp=True, compute_masks=True,
-             progress=None):
+             stitch_threshold=0.0, min_size=15, max_size_fraction=0.4, niter=None, 
+             augment=False, tile=True, tile_overlap=0.1, bsize=224, 
+             interp=True, compute_masks=True, progress=None):
         """ segment list of images x, or 4D array - Z x nchan x Y x X
 
         Args:
@@ -394,6 +394,8 @@ def eval(self, x, batch_size=8, resample=True, channels=None, channel_axis=None,
             anisotropy (float, optional): for 3D segmentation, optional rescaling factor (e.g. set to 2.0 if Z is sampled half as dense as X or Y). Defaults to None.
             stitch_threshold (float, optional): if stitch_threshold>0.0 and not do_3D, masks are stitched in 3D to return volume segmentation. Defaults to 0.0.
             min_size (int, optional): all ROIs below this size, in pixels, will be discarded. Defaults to 15.
+            max_size_fraction (float, optional): max_size_fraction (float, optional): Masks larger than max_size_fraction of
+                total image size are removed. Default is 0.4.
             niter (int, optional): number of iterations for dynamics computation. if None, it is set proportional to the diameter. Defaults to None.
             augment (bool, optional): tiles image with overlapping tiles and flips overlapped regions to augment. Defaults to False.
             tile (bool, optional): tiles image to ensure GPU/CPU memory usage limited (recommended). Defaults to True.
@@ -435,7 +437,8 @@ def eval(self, x, batch_size=8, resample=True, channels=None, channel_axis=None,
                     tile_overlap=tile_overlap, bsize=bsize, resample=resample,
                     interp=interp, flow_threshold=flow_threshold,
                     cellprob_threshold=cellprob_threshold, compute_masks=compute_masks,
-                    min_size=min_size, stitch_threshold=stitch_threshold,
+                    min_size=min_size, max_size_fraction=max_size_fraction, 
+                    stitch_threshold=stitch_threshold,
                     progress=progress, niter=niter)
                 masks.append(maski)
                 flows.append(flowi)
@@ -464,7 +467,7 @@ def eval(self, x, batch_size=8, resample=True, channels=None, channel_axis=None,
                 rescale=rescale, resample=resample, augment=augment, tile=tile,
                 tile_overlap=tile_overlap, bsize=bsize, flow_threshold=flow_threshold,
                 cellprob_threshold=cellprob_threshold, interp=interp, min_size=min_size,
-                do_3D=do_3D, anisotropy=anisotropy, niter=niter,
+                max_size_fraction=max_size_fraction, do_3D=do_3D, anisotropy=anisotropy, niter=niter,
                 stitch_threshold=stitch_threshold)
 
             flows = [plot.dx_to_circ(dP), dP, cellprob, p]
@@ -473,7 +476,8 @@ def eval(self, x, batch_size=8, resample=True, channels=None, channel_axis=None,
     def _run_cp(self, x, compute_masks=True, normalize=True, invert=False, niter=None,
                 rescale=1.0, resample=True, augment=False, tile=True, tile_overlap=0.1,
                 cellprob_threshold=0.0, bsize=224, flow_threshold=0.4, min_size=15,
-                interp=True, anisotropy=1.0, do_3D=False, stitch_threshold=0.0):
+                max_size_fraction=0.4, interp=True, anisotropy=1.0, do_3D=False, 
+                stitch_threshold=0.0):
 
         if isinstance(normalize, dict):
             normalize_params = {**normalize_default, **normalize}
@@ -538,7 +542,7 @@ def _run_cp(self, x, compute_masks=True, normalize=True, invert=False, niter=Non
                 masks, p = dynamics.resize_and_compute_masks(
                     dP, cellprob, niter=niter, cellprob_threshold=cellprob_threshold,
                     flow_threshold=flow_threshold, interp=interp, do_3D=do_3D,
-                    min_size=min_size, resize=None,
+                    min_size=min_size, max_size_fraction=max_size_fraction, resize=None,
                     device=self.device if self.gpu else None)
             else:
                 masks, p = [], []
@@ -557,6 +561,7 @@ def _run_cp(self, x, compute_masks=True, normalize=True, invert=False, niter=Non
                         resize=resize,
                         min_size=min_size if stitch_threshold == 0 or nimg == 1 else
                         -1,  # turn off for 3D stitching
+                        max_size_fraction=max_size_fraction,
                         device=self.device if self.gpu else None)
                     masks.append(outputs[0])
                     p.append(outputs[1])