Speedup in ChannelShuffle

albumentations/augmentations/functional.py

@@ -1145,7 +1145,12 @@ def invert(img: np.ndarray) -> np.ndarray:
 
 
 def channel_shuffle(img: np.ndarray, channels_shuffled: np.ndarray) -> np.ndarray:
-    return img[..., channels_shuffled]
+    img = img.copy()
+    from_to = []
+    for i, j in enumerate(channels_shuffled):
+        from_to.extend([i, j])
+    cv2.mixChannels([img], [img], from_to)
+    return img
 
 
 def gamma_transform(img: np.ndarray, gamma: float) -> np.ndarray:

albumentations/augmentations/transforms.py

@@ -2489,6 +2489,9 @@
     Targets:
         image
 
+    Number of channels:
+        Any
+
     Image types:
         uint8, float32
 
@@ -2497,18 +2500,23 @@
     def apply(
         self,
         img: np.ndarray,
-        channels_shuffled: tuple[int, ...],
+        channels_shuffled: tuple[int, ...] | None,
         **params: Any,
     ) -> np.ndarray:
+        if channels_shuffled is None:
+            return img
         return fmain.channel_shuffle(img, channels_shuffled)
 
     def get_params_dependent_on_data(
         self,
         params: dict[str, Any],
         data: dict[str, Any],
     ) -> dict[str, Any]:
-        ch_arr = list(range(params["shape"][2]))
-        self.random_generator.shuffle(ch_arr)
+        shape = params["shape"]
+        if len(shape) == 2 or shape[-1] == 1:
+            return {"channels_shuffled": None}
+        ch_arr = list(range(shape[-1]))
+        self.py_random.shuffle(ch_arr)
         return {"channels_shuffled": ch_arr}
 
     def get_transform_init_args_names(self) -> tuple[()]: