Merge pull request #1183 from MouseLand/3d_shape_fix

Refactor image conversion functions for improved clarity and error ha…

Author: mrariden

cellpose/io.py

@@ -230,33 +230,7 @@ def imread_2D(img_file):
         img_out (numpy.ndarray): The 3-channel image data as a NumPy array.
     """
     img = imread(img_file)
-
-    if (img.ndim == 1) or (img.ndim == 4):
-        raise ValueError("img_file should have 2 or 3 dimensions, shape: %s" % img.shape)
-
-    # if image has no channel dimension, add one and return the image
-    if img.ndim == 2:
-        img_out = np.zeros((img.shape[0], img.shape[1], 3), dtype=img.dtype)
-        img_out[:, :, 0] = img
-        return img_out
-
-    # Otherwise, image will have a channel dimension, assume it's either first or last
-    # force it to be last (XYC): 
-    if img.shape[0] <  img.shape[-1]:
-        img = np.moveaxis(img, 0, -1)
-
-    nchan = img.shape[-1]
-
-    if nchan == 3:
-        # already has 3 channels
-        return img
-    
-    # ensure there are 3 channels
-    img_out = np.zeros((img.shape[0], img.shape[1], 3), dtype=img.dtype)
-    copy_chan = min(3, nchan)
-    img_out[:, :, :copy_chan] = img[:, :, :copy_chan]
-
-    return img_out
+    return transforms.convert_image(img, do_3D=False)
 
 
 def imread_3D(img_file):
@@ -271,20 +245,25 @@ def imread_3D(img_file):
     """
     img = imread(img_file)
 
+    dimension_lengths = list(img.shape)
+
+    # guess at channel axis:
+    channel_axis = np.argmin(img.shape)
+    del dimension_lengths[channel_axis]
+
+    # guess at z axis:
+    z_axis = np.argmin(dimension_lengths)
+
+    # grayscale images:
     if img.ndim == 3:
-        # add a channel dimension
-        img_out = np.zeros((img.shape[0], img.shape[1], img.shape[2], 3), dtype=img.dtype)
-        img_out[:, :, :, 0] = img
-    elif img.ndim == 4:
-        # assume it's opening as (z, c, y, x)
-        img_out = np.zeros((img.shape[0], img.shape[2], img.shape[3], 3), dtype=img.dtype)
-        img_out[:, :, :, :img.shape[1]] = img.transpose(0, 2, 3, 1)
-    else:
-        raise ValueError("Image should have 3 or 4 dimensions, shape: %s" % img.shape)
+        channel_axis = None
 
-    del img
-    return img_out
-
+    try:
+        return transforms.convert_image(img, channel_axis=channel_axis, z_axis=z_axis, do_3D=True)
+    except Exception as e:
+        io_logger.critical("ERROR: could not read file, %s" % e)
+        io_logger.critical("ERROR: Guessed z_axis: %s, channel_axis: %s" % (z_axis, channel_axis))
+        return None
 
 def remove_model(filename, delete=False):
     """ remove model from .cellpose custom model list """

cellpose/transforms.py

@@ -447,9 +447,12 @@ def update_axis(m_axis, to_squeeze, ndim):
     return m_axis
 
 
-def convert_image_3d(x, channel_axis=None, z_axis=None):
+def _convert_image_3d(x, channel_axis=None, z_axis=None):
     """
     Convert a 3D or 4D image array to have dimensions ordered as (Z, X, Y, C).
+
+    Arrays of ndim=3 are assumed to be grayscale and must be specified with z_axis. 
+    Arrays of ndim=4 must have both `channel_axis` and `z_axis` specified.
     
     Args:
         x (numpy.ndarray): Input image array. Must be either 3D (assumed to be grayscale 3D) or 4D. 
@@ -474,6 +477,12 @@ def convert_image_3d(x, channel_axis=None, z_axis=None):
             channels to ensure the output has exactly 3 channels.
     """
 
+    if x.ndim < 3:
+        raise ValueError(f"Input image must have at least 3 dimensions, input shape: {x.shape}, ndim={x.ndim}")
+    
+    if z_axis is not None and z_axis < 0:
+        z_axis += x.ndim
+
     # if image is ndim==3, assume it is greyscale 3D and use provided z_axis
     if x.ndim == 3 and z_axis is not None:
         # add in channel axis
@@ -484,7 +493,11 @@ def convert_image_3d(x, channel_axis=None, z_axis=None):
 
 
     if channel_axis is None or z_axis is None:
-        raise ValueError("both channel_axis and z_axis must be specified when segmenting 3D images of ndim=4")
+        raise ValueError("For 4D images, both `channel_axis` and `z_axis` must be explicitly specified. Please provide values for both parameters.")
+    if channel_axis is not None and channel_axis < 0:
+        channel_axis += x.ndim
+    if channel_axis is None or channel_axis >= x.ndim:
+        raise IndexError(f"channel_axis {channel_axis} is out of bounds for input array with {x.ndim} dimensions")
     assert x.ndim == 4, f"input image must have ndim == 4, ndim={x.ndim}"
 
     x_dim_shapes = list(x.shape)
@@ -519,23 +532,26 @@ def convert_image_3d(x, channel_axis=None, z_axis=None):
             x = x[..., :x_chans_to_copy]
         else: 
             # less than 3 channels: pad up to 
-            x_out = np.zeros((num_z_layers, x_dim_shapes[0], x_dim_shapes[1], 3), dtype=x.dtype)
-            x_out[..., :x_chans_to_copy] = x[...]
-            x = x_out
-            del x_out
+            pad_width = [(0, 0), (0, 0), (0, 0), (0, 3 - x_chans_to_copy)]
+            x = np.pad(x, pad_width, mode='constant', constant_values=0)
 
     return x
 
 
 def convert_image(x, channel_axis=None, z_axis=None, do_3D=False):
-    """Converts the image to have the z-axis first, channels last.
+    """Converts the image to have the z-axis first, channels last. Image will be converted to 3 channels if it is not already.
+    If more than 3 channels are provided, only the first 3 channels will be used. 
+
+    Accepts: 
+        - 2D images with no channel dimension: `z_axis` and `channel_axis` must be `None`
+        - 2D images with channel dimension: `channel_axis` will be guessed between first or last axis, can also specify `channel_axis`. `z_axis` must be `None`
+        - 3D images with or without channels: 
 
     Args:
         x (numpy.ndarray or torch.Tensor): The input image.
         channel_axis (int or None): The axis of the channels in the input image. If None, the axis is determined automatically.
         z_axis (int or None): The axis of the z-dimension in the input image. If None, the axis is determined automatically.
         do_3D (bool): Whether to process the image in 3D mode. Defaults to False.
-        nchan (int): The number of channels to keep if the input image has more than nchan channels.
 
     Returns:
         numpy.ndarray: The converted image.
@@ -551,18 +567,23 @@ def convert_image(x, channel_axis=None, z_axis=None, do_3D=False):
         transforms_logger.warning("torch array used as input, converting to numpy")
         x = x.cpu().numpy()
 
+    # should be 2D
+    if z_axis is not None and not do_3D:
+        raise ValueError("2D image provided, but z_axis is not None. Set z_axis=None to process 2D images of ndim=2 or 3.")
+
+    # make sure that channel_axis and z_axis are specified if 3D
+    if ndim == 4 and not do_3D:
+        raise ValueError("3D input image provided, but do_3D is False. Set do_3D=True to process 3D images. ndims=4")
+
     # make sure that channel_axis and z_axis are specified if 3D
     if do_3D:
-        return convert_image_3d(x, channel_axis=channel_axis, z_axis=z_axis)
+        return _convert_image_3d(x, channel_axis=channel_axis, z_axis=z_axis)
 
-    if ndim == 4:
-        raise ValueError("3D input image provided, but do_3D is False. Set do_3D=True to process 3D images.")
-
     ######################## 2D reshaping ########################
     # if user specifies channel axis, return early
     if channel_axis is not None:
         if ndim == 2:
-            raise ValueError("2D image provided, but channel_axis is not None. Set channel_axis=None to process 2D images.")
+            raise ValueError("2D image provided, but channel_axis is not None. Set channel_axis=None to process 2D images of ndim=2.")
 
         # Put channel axis last:
         # Find the indices of the dims that need to be put in dim 0 and 1
@@ -613,8 +634,9 @@ def convert_image(x, channel_axis=None, z_axis=None, do_3D=False):
         del x_out
     else:
         # something is wrong: yell
-        transforms_logger.critical(f"ERROR: Unexpected image shape: {str(x.shape)}")
-        raise ValueError(f"ERROR: Unexpected image shape: {str(x.shape)}")
+        expected_shapes = "2D (H, W), 3D (H, W, C), or 4D (Z, H, W, C)"
+        transforms_logger.critical(f"ERROR: Unexpected image shape: {str(x.shape)}. Expected shapes: {expected_shapes}")
+        raise ValueError(f"ERROR: Unexpected image shape: {str(x.shape)}. Expected shapes: {expected_shapes}")
 
     return x
 

tests/test_output.py

@@ -94,19 +94,17 @@ def test_cyto2_to_seg(data_dir, image_names, cellposemodel_fixture_24layer):
     clear_output(data_dir, image_names)
 
 
-def test_class_3D(data_dir, image_names_3d, cellposemodel_fixture_2layer):
+def test_class_3D_one_img(data_dir, image_names_3d, cellposemodel_fixture_2layer):
     clear_output(data_dir, image_names_3d)
 
-    for image_name in image_names_3d:
-        img_file = data_dir / '3D' / image_name
-        img = io.imread_3D(img_file)
-        masks_pred, flows_pred, _ = cellposemodel_fixture_2layer.eval(img, do_3D=True, channel_axis=-1, z_axis=0)
-        # io.imsave(data_dir / "3D" / (img_file.stem + "_cp_masks.tif"), masks)
+    img_file = data_dir / '3D' / image_names_3d[0]
+    img = io.imread_3D(img_file)
+    masks_pred, flows_pred, _ = cellposemodel_fixture_2layer.eval(img, do_3D=True, channel_axis=-1, z_axis=0)
 
-        assert img.shape[:-1] == masks_pred.shape, f'mask incorrect shape for {image_name}, {masks_pred.shape=}'
-        assert img.shape[:-1] == flows_pred[1].shape[1:], f'flows incorrect shape for {image_name}, {flows_pred.shape=}'
+    assert img.shape[:-1] == masks_pred.shape, f'mask incorrect shape for {image_name}, {masks_pred.shape=}'
+    assert img.shape[:-1] == flows_pred[1].shape[1:], f'flows incorrect shape for {image_name}, {flows_pred.shape=}'
 
-        break # Just test one img for now
+    # just compare shapes for now
     # compare_masks_cp4(data_dir, image_names_3d, "3D")
     clear_output(data_dir, image_names_3d)
 

tests/test_transforms.py

@@ -1,6 +1,7 @@
 import numpy as np
 import pytest
 
+from cellpose import transforms
 from cellpose.io import imread
 from cellpose.transforms import normalize_img, random_rotate_and_resize, resize_image
 
@@ -90,3 +91,41 @@ def test_resize(img_2d):
     img32 = resize_image(img_2d.astype("uint32"), Lx=Lx, Ly=Ly)
     assert img32.shape == (Ly, Lx, 3)
     assert img32.dtype == np.uint32
+
+
+@pytest.mark.parametrize(
+        "input_shape, channel_axis, z_axis, do_3D, expected_shape, raises_error",
+        [   # passing:
+            # 2D:
+            ((100, 120), None, None, False, (100, 120, 3), False),  # 2D grayscale image
+            ((100, 120, 3), None, None, False, (100, 120, 3), False),  # 2D RGB image
+            ((3, 100, 120), 0, None, False, (100, 120, 3), False),  # 2D RGB image with channels first
+            ((3, 100, 120), None, None, False, (100, 120, 3), False),  # 2D RGB image with channels first
+
+            # 3D:
+            ((100, 120, 5), None, -1, True, (5, 100, 120, 3), False),  # 3D grayscale image
+            ((5, 100, 120), None, 0, True, (5, 100, 120, 3), False),  # 3D grayscale image
+            ((100, 5, 120, 5), 1, 3, True, (5, 100, 120, 3), False),  # 3D 5chan image
+            ((10, 100, 120, 3), -1, 0, True, (10, 100, 120, 3), False),  # 3D 5chan image
+            
+            # failing: 
+            # 2D:
+            ((100, 120), None, 0, False, (100, 120, 3), True),  # 2D grayscale image
+            ((100, 120, 3), None, None, True, (100, 120, 3), True),  # 2D RGB image
+            ((3, 100, 120), -1, 2, False, (100, 120, 3), True),  # 2D RGB image with channels first
+            ((3, 100, 120), None, None, True, (100, 120, 3), True),  # 2D RGB image with channels first
+
+            # 3D:
+            ((5, 100, 120), None, None, True, (5, 100, 120, 3), True),  # 3D grayscale image
+            ((10, 100, 120, 3), -1, 0, False, (10, 100, 120, 3), True),  # 3D rgb image
+        ],
+    )
+def test_convert_image(input_shape, channel_axis, z_axis, do_3D, expected_shape, raises_error):
+    """Test the convert_image function with various input shapes and parameters."""
+    img = np.random.rand(*input_shape).astype(np.float32)
+    if raises_error:
+        with pytest.raises(ValueError):
+            transforms.convert_image(img, channel_axis=channel_axis, z_axis=z_axis, do_3D=do_3D)
+    else:
+        converted_img = transforms.convert_image(img, channel_axis=channel_axis, z_axis=z_axis, do_3D=do_3D)
+        assert converted_img.shape == expected_shape, f"Expected shape {expected_shape}, but got {converted_img.shape}"