Fix channel_axis bug in transforms.py

MouseLand · Apr 6, 2023 · 70e66c6 · 70e66c6
1 parent e559dae
commit 70e66c6
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/cellpose/transforms.py b/cellpose/transforms.py
@@ -258,7 +258,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
@@ -770,4 +770,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