adding z_axis in cellposedenoisemodel (#1099)

Author: carsen-stringer

cellpose/denoise.py

@@ -584,6 +584,7 @@ def eval(self, x, batch_size=8, channels=None, channel_axis=None, z_axis=None,
         diameter = self.dn.diam_mean if self.dn.ratio > 1 else diameter
         masks, flows, styles = self.cp.eval(
             img_restore, batch_size=batch_size, channels=channels_new, channel_axis=-1,
+            z_axis=0 if img_restore.ndim > 3 and img_restore.shape[0] > 0 else None,
             normalize=normalize_params, rescale=rescale, diameter=diameter,
             tile_overlap=tile_overlap, augment=augment, resample=resample,
             invert=invert, flow_threshold=flow_threshold,

cellpose/transforms.py

@@ -470,6 +470,7 @@ def convert_image(x, channels, channel_axis=None, z_axis=None, do_3D=False, ncha
     """
     # check if image is a torch array instead of numpy array
     # converts torch to numpy
+    ndim = x.ndim
     if torch.is_tensor(x):
         transforms_logger.warning("torch array used as input, converting to numpy")
         x = x.cpu().numpy()
@@ -516,7 +517,7 @@ def convert_image(x, channels, channel_axis=None, z_axis=None, do_3D=False, ncha
         x = x[:, :, np.newaxis]
 
     if do_3D:
-        if x.ndim < 3:
+        if ndim < 3:
             transforms_logger.critical("ERROR: cannot process 2D images in 3D mode")
             raise ValueError("ERROR: cannot process 2D images in 3D mode")
         elif x.ndim < 4: