add support for anisotropic values of the flow3D_smooth

cellpose/cli.py

@@ -109,8 +109,8 @@ def get_arg_parser():
         "--min_size", required=False, default=15, type=int,
         help="minimum number of pixels per mask, can turn off with -1")
     algorithm_args.add_argument(
-        "--flow3D_smooth", required=False, default=0, type=float,
-        help="stddev of gaussian for smoothing of dP for dynamics in 3D, default of 0 means no smoothing")
+        "--flow3D_smooth", required=False, default=0, type=float, nargs='+',
+        help="stddev of gaussian for smoothing of dP for dynamics in 3D, default of 0 means no smoothing. Pass a list of values to allow smoothing of the ZYX axes independently")
     algorithm_args.add_argument(
         "--flow_threshold", default=0.4, type=float, help=
         "flow error threshold, 0 turns off this optional QC step. Default: %(default)s")

cellpose/models.py

@@ -187,7 +187,7 @@ def eval(self, x, batch_size=8, resample=True, channels=None, channel_axis=None,
             flow_threshold (float, optional): flow error threshold (all cells with errors below threshold are kept) (not used for 3D). Defaults to 0.4.
             cellprob_threshold (float, optional): all pixels with value above threshold kept for masks, decrease to find more and larger masks. Defaults to 0.0.
             do_3D (bool, optional): set to True to run 3D segmentation on 3D/4D image input. Defaults to False.
-            flow3D_smooth (int, optional): if do_3D and flow3D_smooth>0, smooth flows with gaussian filter of this stddev. Defaults to 0.
+            flow3D_smooth (int or float or list of (int or float), optional): if do_3D and flow3D_smooth>0, smooth flows with gaussian filter of this stddev. List smooths the ZYX axes independently and must be length 3. Defaults to 0.
             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.
@@ -319,18 +319,11 @@ def eval(self, x, batch_size=8, resample=True, channels=None, channel_axis=None,
             do_3D=do_3D, 
             anisotropy=anisotropy)
 
-        if do_3D:    
-            if flow3D_smooth > 0:
-                models_logger.info(f"smoothing flows with sigma={flow3D_smooth}")
-                dP = gaussian_filter(dP, (0, flow3D_smooth, flow3D_smooth, flow3D_smooth))
+        if do_3D:
             torch.cuda.empty_cache()
             gc.collect()
 
             if resample:
-                # upsample flows flows before computing them:
-                # dP = self._resize_gradients(dP, to_y_size=Ly_0, to_x_size=Lx_0, to_z_size=Lz_0)
-                # cellprob = self._resize_cellprob(cellprob, to_x_size=Lx_0, to_y_size=Ly_0, to_z_size=Lz_0)
-
                 # resize XY then YZ and then put channels first
                 dP = transforms.resize_image(dP.transpose(1, 2, 3, 0), Ly=Ly_0, Lx=Lx_0, no_channels=False)
                 dP = transforms.resize_image(dP.transpose(1, 0, 2, 3), Lx=Lx_0, Ly=Lz_0, no_channels=False)
@@ -341,16 +334,22 @@ def eval(self, x, batch_size=8, resample=True, channels=None, channel_axis=None,
                 cellprob = transforms.resize_image(cellprob.transpose(1, 0, 2), Lx=Lx_0, Ly=Lz_0, no_channels=True)
                 cellprob = cellprob.transpose(1, 0, 2)
 
-
         # 2d case:
         if resample and not do_3D:
-            # upsample flows before computing them: 
-            # dP = self._resize_gradients(dP, to_y_size=Ly_0, to_x_size=Lx_0, to_z_size=Lz_0)
-            # cellprob = self._resize_cellprob(cellprob, to_x_size=Lx_0, to_y_size=Ly_0, to_z_size=Lz_0)
-
             # 2D images have N = 1 in batch dimension:
             dP = transforms.resize_image(dP.transpose(1, 2, 3, 0), Ly=Ly_0, Lx=Lx_0, no_channels=False).transpose(3, 0, 1, 2)
             cellprob = transforms.resize_image(cellprob, Ly=Ly_0, Lx=Lx_0, no_channels=True)
+
+        if do_3D and flow3D_smooth:
+            if isinstance(flow3D_smooth, (int, float)):
+                flow3D_smooth = [flow3D_smooth]*3 
+            if isinstance(flow3D_smooth, list) and len(flow3D_smooth) == 1:
+                flow3D_smooth = flow3D_smooth*3
+            if len(flow3D_smooth) == 3 and any(v > 0 for v in flow3D_smooth):
+                models_logger.info(f"smoothing flows with ZYX sigma={flow3D_smooth}")
+                dP = gaussian_filter(dP, [0, *flow3D_smooth])
+            else: 
+                models_logger.warning(f"Could not do flow smoothing with {flow3D_smooth} either because its len was not 3 or no items were > 0, skipping flow3D_smoothing")
 
         if compute_masks:
             # use user niter if specified, otherwise scale niter (200) with diameter

docs/do3d.rst

@@ -57,6 +57,9 @@ If you see many cells that are fragmented, you can smooth the flows before the d
 are run in 3D using the ``flow3D_smooth`` parameter, which specifies the standard deviation of 
 a Gaussian for smoothing the flows. The default is 0.0, which means no smoothing. Alternatively/additionally,
 you may want to train a model on 2D slices from your 3D data to improve the segmentation (see below).
+*If there are ring-like artifacts in your masks*, increasing ``flow3D_smooth`` can help remove them. 
+You can specify the ZYX flow smoothing independently for each axis by passing a list of values to the ``flow3D_smooth`` 
+argument. For example: ``flow3D_smooth = [2, 0, 0]`` 
 
 The network can rescale images using the user diameter and the model ``diam_mean`` (30),
 so for example if you input a diameter of 90, 

tests/test_output.py

@@ -46,7 +46,7 @@ def clear_output(data_dir, image_names):
                              (False, True, None),
                              (False, False, None),
                              (True, False, None),
-                             (True, False, 40)
+                             (True, False, 40),
                          ]
 )
 def test_class_2D_one_img(data_dir, image_names, cellposemodel_fixture_24layer, compute_masks, resample, diameter):
@@ -172,13 +172,21 @@ def test_cli_3D_diam_anisotropy_shape(data_dir, image_names_3d, diam, aniso):
     compare_mask_shapes(data_dir, image_names_3d[0], "3D")
     clear_output(data_dir, image_names_3d)
 
-
+@pytest.mark.parametrize('flow3D_smooth',
+                         [None, 2, [1., 0., 0.]])
 @pytest.mark.slow
-def test_cli_3D_one_img(data_dir, image_names_3d):
+def test_cli_3D_one_img(data_dir, image_names_3d, flow3D_smooth):
     clear_output(data_dir, image_names_3d)
     use_gpu = torch.cuda.is_available() or torch.backends.mps.is_available() 
     gpu_string = "--use_gpu" if use_gpu else ""
-    cmd = f"python -m cellpose --image_path {str(data_dir / '3D' / image_names_3d[0])} --do_3D --save_tif {gpu_string} --verbose"
+
+    flow_string = ''
+    if isinstance(flow3D_smooth, (float, int)):
+        flow_string = f" --flow3D_smooth {flow3D_smooth}"
+    elif isinstance(flow3D_smooth, list):
+        flow_string = f" --flow3D_smooth {' '.join([str(f) for f in flow3D_smooth])}"
+
+    cmd = f"python -m cellpose --image_path {str(data_dir / '3D' / image_names_3d[0])} --do_3D --save_tif {gpu_string} --verbose{flow_string}"
     print(cmd)
     try:
         cmd_stdout = check_output(cmd, stderr=STDOUT, shell=True).decode()