Fix new skimage channel_axis (#1421)

Co-authored-by: Mikhail Druzhinin <dipet@gmail.com>
albumentations-team · Mar 8, 2023 · 227bd74 · 227bd74
1 parent cb37273
commit 227bd74
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/albumentations/augmentations/functional.py b/albumentations/augmentations/functional.py
@@ -1272,9 +1272,9 @@ def superpixels(
             resize_fn = _maybe_process_in_chunks(cv2.resize, dsize=(new_width, new_height), interpolation=interpolation)
             image = resize_fn(image)
 
-    from skimage.segmentation import slic
-
-    segments = skimage.segmentation.slic(image, n_segments=n_segments, compactness=10)
+    segments = skimage.segmentation.slic(
+        image, n_segments=n_segments, compactness=10, channel_axis=-1 if image.ndim > 2 else None
+    )
 
     min_value = 0
     max_value = MAX_VALUES_BY_DTYPE[image.dtype]