move reordering axes to after image normalization

Author: mrariden

cellpose/transforms.py

@@ -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):