Merge pull request #1 from MouseLand/ring_artifacts-refactor

## Summary of Changes

- Modularization of ring artifact detection code.
- Improved logic for artifact removal and filtering.
- Enhanced documentation and inline comments for clarity.
- Compatibility updates to work with upstream changes from MouseLand.
- Bug fixes identified during code review and testing.
- Refined parameter settings and default behaviors for artifact processing.

Author: yuriyzubov

.github/workflows/test_and_deploy.yml

@@ -3,7 +3,7 @@
 
 name: tests
 
-on: 
+on:
   push:
     branches:
       - main
@@ -14,7 +14,7 @@ on:
       - main
   workflow_dispatch:
 
-jobs:          
+jobs:
   test:
     name: ${{ matrix.platform }} py${{ matrix.python-version }}
     runs-on: ${{ matrix.platform }}
@@ -26,10 +26,12 @@ jobs:
 
     steps:
       - uses: actions/checkout@v4
+
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
+
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
@@ -45,21 +47,25 @@ jobs:
 
   deploy:
     # this will run when you have tagged a commit, starting with "v*"
-    # and requires that you have put your twine API key in your 
+    # and requires that you have put your twine API key in your
     # github secrets (see readme for details)
     needs: [test]
     runs-on: ubuntu-latest
     if: contains(github.ref, 'tags')
+
     steps:
       - uses: actions/checkout@v4
+
       - name: Set up Python
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@v5
         with:
           python-version: "3.x"
+
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
           pip install -U setuptools setuptools_scm wheel twine
+
       - name: Build and publish
         env:
           TWINE_USERNAME: __token__

cellpose/gui/io.py

@@ -139,6 +139,7 @@ def _load_image(parent, filename=None, load_seg=True, load_3D=False):
     except Exception as e:
         print("ERROR: images not compatible")
         print(f"ERROR: {e}")
+        return
 
     if parent.loaded:
         parent.reset()

cellpose/io.py

@@ -47,6 +47,33 @@
 io_logger = logging.getLogger(__name__)
 
 def logger_setup(cp_path=".cellpose", logfile_name="run.log", stdout_file_replacement=None):
+    """Set up logging to a file and stdout (or a file replacement).
+
+    Creates the log directory if it doesn't exist, removes any existing log
+    file, and configures the root logger to write INFO-level and above messages
+    to both a log file and stdout (or a replacement file).
+
+    Parameters
+    ----------
+    cp_path : str, optional
+        Directory name under the user's home directory for log output.
+        Default is ".cellpose".
+    logfile_name : str, optional
+        Name of the log file created inside cp_path. Default is "run.log".
+    stdout_file_replacement : str or None, optional
+        If provided, log output is written to this file path instead of stdout.
+
+    Returns
+    -------
+    logger : logging.Logger
+        Configured logger for this module. Only INFO and above messages are
+        emitted by default. To enable debug output, call
+        ``logger.setLevel(logging.DEBUG)`` on the returned logger.
+
+    Notes
+    -----
+    The log file is deleted and recreated on each call.
+    """
     cp_dir = pathlib.Path.home().joinpath(cp_path)
     cp_dir.mkdir(exist_ok=True)
     log_file = cp_dir.joinpath(logfile_name)
@@ -189,6 +216,28 @@ def imread(filename):
         if not ND2:
             io_logger.critical("ERROR: need to 'pip install nd2' to load in .nd2 file")
             return None
+        else:
+            with nd2.ND2File(filename) as nd2_file:
+                img = nd2_file.asarray()
+                sizes = nd2_file.sizes
+
+            kept_axes = [nd2.AXIS.Y, nd2.AXIS.X, nd2.AXIS.CHANNEL, nd2.AXIS.Z]
+            # For multi-dimensional data (T, P, etc.), take first frame/position
+            # Work backwards through axes to avoid index shifting
+            for i, (ax_name, size) in reversed(list(enumerate(sizes.items()))):
+                # Keep Y, X, C, Z; remove or reduce everything else
+                if ax_name not in kept_axes:
+                    if size > 1:
+                        io_logger.warning(
+                            f"ND2 file has {size} {ax_name} - using first only"
+                        )
+                    # Take first element (works for both size=1 and size>1)
+                    img = np.take(img, 0, axis=i)
+
+            # Result should now be YX, CYX, ZYX, or CZYX depending on original axes
+            # nd2 preserves axis order from sizes dict (usually C, Z, Y, X)
+            return img
+
     elif ext == ".nrrd":
         if not NRRD:
             io_logger.critical(
@@ -230,40 +279,47 @@ def imread_2D(img_file):
         img_out (numpy.ndarray): The 3-channel image data as a NumPy array.
     """
     img = imread(img_file)
+    if img is None:
+        raise ValueError(f"could not read image file {img_file}")
     return transforms.convert_image(img, do_3D=False)
 
 
 def imread_3D(img_file):
     """
     Read in a 3D image file and convert it to have a channel axis last automatically. Attempts to do this for multi-channel and grayscale images.
 
-    If multichannel image, the channel axis is assumed to be the smallest dimension, and the z axis is the next smallest dimension. 
-    Use `cellpose.io.imread()` to load the full image without selecting the z and channel axes. 
-    
+    For grayscale images (3D array), axis 0 is assumed to be the Z axis (e.g., Z x Y x X).
+    For multichannel images (4D array), the channel axis is assumed to be the smallest dimension,
+    and the Z axis is assumed to be the first remaining axis after the channel axis is removed.
+
+    Use ``cellpose.io.imread()`` to load the full image without automatic axis selection,
+    then specify ``z_axis`` and ``channel_axis`` manually when calling ``model.eval``.
+
     Args:
         img_file (str): The path to the image file.
 
     Returns:
-        img_out (numpy.ndarray): The image data as a NumPy array.
+        img_out (numpy.ndarray): The image data as a NumPy array with channels last, or None if loading fails.
     """
     img = imread(img_file)
+    if img is None:
+        raise ValueError(f"could not read image file {img_file}")
 
     dimension_lengths = list(img.shape)
 
     # grayscale images:
     if img.ndim == 3:
         channel_axis = None
         # guess at z axis:
-        z_axis = np.argmin(dimension_lengths)
+        z_axis = 0
 
     elif img.ndim == 4:
         # guess at channel axis:
         channel_axis = np.argmin(dimension_lengths)
-
-        # guess at z axis: 
-        # set channel axis to max so argmin works:
-        dimension_lengths[channel_axis] = max(dimension_lengths)
-        z_axis = np.argmin(dimension_lengths)
+        dimensions = list(range(img.ndim))
+        dimensions.pop(channel_axis)
+        # guess at z axis as the first remaining dimension: 
+        z_axis = dimensions[0]
 
     else: 
         raise ValueError(f'image shape error, 3D image must 3 or 4 dimensional. Number of dimensions: {img.ndim}')

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 or list[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.
@@ -320,20 +320,10 @@ def eval(self, x, batch_size=8, resample=True, channels=None, channel_axis=None,
             anisotropy=anisotropy)
 
         if do_3D:
-            
-            if isinstance(flow3D_smooth, int):
-                flow3D_smooth = [flow3D_smooth]*3 
-            if any(v > 0 for v in flow3D_smooth):
-                models_logger.info(f"smoothing flows with sigma={flow3D_smooth}")
-                dP = gaussian_filter(dP, [0, *flow3D_smooth])
             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)
@@ -344,16 +334,20 @@ 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 > 0:
+            if isinstance(flow3D_smooth, int) or isinstance(flow3D_smooth, float):
+                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

cellpose/transforms.py

@@ -612,7 +612,7 @@ def convert_image(x, channel_axis=None, z_axis=None, do_3D=False):
         x_out[..., 0] = x
         x = x_out
         del x_out
-        transforms_logger.info(f'processing grayscale image with {x.shape[0], x.shape[1]} HW')
+        transforms_logger.debug(f'processing grayscale image with {x.shape[0], x.shape[1]} HW')
     elif ndim == 3:
         # assume 2d with channels
         # find dim with smaller size between first and last dims
@@ -629,7 +629,7 @@ def convert_image(x, channel_axis=None, z_axis=None, do_3D=False):
         x_out[..., :num_channels] = x[..., :num_channels]
         x = x_out
         del x_out
-        transforms_logger.info(f'processing image with {x.shape[0], x.shape[1]} HW, and {x.shape[2]} channels')
+        transforms_logger.debug(f'processing image with {x.shape[0], x.shape[1]} HW, and {x.shape[2]} channels')
     elif ndim == 4:
         # assume batch of 2d with channels
 
@@ -642,7 +642,7 @@ def convert_image(x, channel_axis=None, z_axis=None, do_3D=False):
         x_out[..., :num_channels] = x[..., :num_channels]
         x = x_out
         del x_out
-        transforms_logger.info(f'processing image batch with {x.shape[0]} images, {x.shape[1], x.shape[2]} HW, and {x.shape[3]} channels')
+        transforms_logger.debug(f'processing image batch with {x.shape[0]} images, {x.shape[1], x.shape[2]} HW, and {x.shape[3]} channels')
     else:
         # something is wrong: yell
         expected_shapes = "2D (H, W), 3D (H, W, C), or 4D (Z, H, W, C)"
@@ -776,10 +776,6 @@ def normalize_img(img, normalize=True, norm3D=True, invert=False, lowhigh=None,
             transforms_logger.critical(error_message)
             raise ValueError(error_message)
 
-    # Move channel axis back to the original position
-    if axis != -1 and axis != img_norm.ndim - 1:
-        img_norm = np.moveaxis(img_norm, -1, axis)
-
     # The transformer can get confused if a channel is all 1's instead of all 0's:
     for i, chan_did_normalize in enumerate(cgood):
         if not chan_did_normalize:
@@ -788,6 +784,10 @@ def normalize_img(img, normalize=True, norm3D=True, invert=False, lowhigh=None,
             if img_norm.ndim == 4:
                 img_norm[:, :, :, i] = 0
 
+    # Move channel axis back to the original position
+    if axis != -1 and axis != img_norm.ndim - 1:
+        img_norm = np.moveaxis(img_norm, -1, axis)
+
     return img_norm
 
 def resize_safe(img, Ly, Lx, interpolation=cv2.INTER_LINEAR):

docs/do3d.rst

@@ -25,12 +25,21 @@ then the GUI will automatically run 3D segmentation and display it in the GUI. W
 the command line for progress. It is recommended to use a GPU to speed up processing.
 
 In the CLI/notebook, you need to specify the ``z_axis`` and the ``channel_axis``
-parameters to specify the axis (0-based) of the image which corresponds to the image channels and to the z axis. 
-For example an image with 2 channels of shape (1024,1024,2,105,1) can be 
-specified with ``channel_axis=2`` and ``z_axis=3``. These parameters can be specified using the command line 
-with ``--channel_axis`` or ``--z_axis`` or as inputs to ``model.eval`` for 
+parameters to specify the axis (0-based) of the image which corresponds to the image channels and to the z axis.
+For example an image with 2 channels of shape (1024,1024,2,105,1) can be
+specified with ``channel_axis=2`` and ``z_axis=3``. These parameters can be specified using the command line
+with ``--channel_axis`` or ``--z_axis`` or as inputs to ``model.eval`` for
 the ``CellposeModel`` model.
 
+As a convenience, :func:`cellpose.io.imread_3D` will attempt to load a 3D image and
+automatically guess the axes. For grayscale images (3D array), axis 0 is assumed
+to be the Z axis (e.g., Z x Y x X). For multichannel images (4D array), the
+channel axis is assumed to be the smallest dimension, and the Z axis is assumed to
+be the first remaining axis after the channel axis is identified (e.g., for a
+Z x C x Y x X image, channel axis = 1 and z axis = 0). If your image does not
+follow these conventions, use ``cellpose.io.imread`` and specify ``z_axis`` and
+``channel_axis`` manually.
+
 Volumetric stacks do not always have the same sampling in XY as they do in Z. 
 Therefore you can set an ``anisotropy`` parameter in CLI/notebook to allow for differences in 
 sampling, e.g. set to 2.0 if Z is sampled half as dense as X or Y, and then in the algorithm 
@@ -48,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

@@ -39,25 +39,28 @@ def clear_output(data_dir, image_names):
             os.remove(npy_output)
 
 
-@pytest.mark.parametrize('compute_masks, resample, diameter', 
+@pytest.mark.parametrize('compute_masks, resample, diameter, flow3D_smooth', 
                          [
-                             (True, True, 40), 
-                             (True, True, None), 
-                             (False, True, None),
-                             (False, False, None),
-                             (True, False, None),
-                             (True, False, 40)
+                             (True, True, 40, None), 
+                             (True, True, None, None), 
+                             (False, True, None, None),
+                             (False, False, None, None),
+                             (True, False, None, None),
+                             (True, False, 40, None),
+                             (False, True, None, 2),
+                             (False, True, None, [2, 0, 0]),
+                             (False, False, None, 2),
                          ]
 )
-def test_class_2D_one_img(data_dir, image_names, cellposemodel_fixture_24layer, compute_masks, resample, diameter):
+def test_class_2D_one_img(data_dir, image_names, cellposemodel_fixture_24layer, compute_masks, resample, diameter, flow3D_smooth):
     clear_output(data_dir, image_names)
 
     img_file = data_dir / '2D' / image_names[0]
 
     img = io.imread_2D(img_file)
     # flowps = io.imread(img_file.parent / (img_file.stem + "_cp4_gt_flowps.tif"))
 
-    masks_pred, _, _ = cellposemodel_fixture_24layer.eval(img, normalize=True, compute_masks=compute_masks, resample=resample, diameter=diameter)
+    masks_pred, _, _ = cellposemodel_fixture_24layer.eval(img, normalize=True, compute_masks=compute_masks, resample=resample, diameter=diameter, flow3D_smooth=flow3D_smooth)
 
     if not compute_masks:
         # not compute_masks won't return masks so can't check