Merge pull request #688 from nickeener/main

Fix channel_axis bug in transforms.py
MouseLand · Feb 13, 2024 · 28ddc24 · 28ddc24
2 parents fbab6fd + 70e66c6
commit 28ddc24
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/cellpose/transforms.py b/cellpose/transforms.py
@@ -262,7 +262,7 @@ def convert_image(x, channels, channel_axis=None, z_axis=None,
         x = x.squeeze()
 
     # put z axis first
-    if z_axis is not None and x.ndim > 2:
+    if z_axis is not None and x.ndim > 2 and z_axis != 0:
         x = move_axis(x, m_axis=z_axis, first=True)
         if channel_axis is not None:
             channel_axis += 1
@@ -774,4 +774,4 @@ def random_rotate_and_resize(X, Y=None, scale_range=1., xy = (224,224),
                 lbl[n,1] = (-v1 * np.sin(-theta) + v2*np.cos(-theta))
                 lbl[n,2] = (v1 * np.cos(-theta) + v2*np.sin(-theta))
 
-    return imgi, lbl, scale
+    return imgi, lbl, scale