Ensure that transformed masks are contiguous arrays

albumentations/core/transforms_interface.py

@@ -21,7 +21,7 @@
     ColorType,
     Targets,
 )
-from .utils import format_args
+from .utils import ensure_contiguous_output, format_args
 
 __all__ = ["BasicTransform", "DualTransform", "ImageOnlyTransform", "NoOp"]
 
@@ -141,14 +141,9 @@ def apply_with_params(self, params: dict[str, Any], *args: Any, **kwargs: Any) -
         for key, arg in kwargs.items():
             if key in self._key2func and arg is not None:
                 target_function = self._key2func[key]
-                if isinstance(arg, np.ndarray):
-                    result = target_function(np.require(arg, requirements=["C_CONTIGUOUS"]), **params)
-                    if isinstance(result, np.ndarray):
-                        res[key] = np.require(result, requirements=["C_CONTIGUOUS"])
-                    else:
-                        res[key] = result
-                else:
-                    res[key] = target_function(arg, **params)
+                res[key] = ensure_contiguous_output(
+                    target_function(ensure_contiguous_output(arg), **params),
+                )
             else:
                 res[key] = arg
         return res

albumentations/core/utils.py

@@ -310,6 +310,14 @@ def ensure_int_output(
     return (int(min_val), int(max_val)) if isinstance(param, int) else (float(min_val), float(max_val))
 
 
+def ensure_contiguous_output(arg: np.ndarray | Sequence[np.ndarray]) -> np.ndarray | list[np.ndarray]:
+    if isinstance(arg, np.ndarray):
+        arg = np.require(arg, requirements=["C_CONTIGUOUS"])
+    elif isinstance(arg, Sequence):
+        arg = list(map(ensure_contiguous_output, arg))
+    return arg
+
+
 @overload
 def to_tuple(param: ScaleIntType, low: ScaleType | None = None, bias: ScalarType | None = None) -> tuple[int, int]: ...
 

tests/test_pytorch.py

@@ -277,6 +277,9 @@ def test_to_tensor_v2_on_non_contiguous_array_with_random_rotate90():
 
     img = np.random.randint(0, 256, (640, 480, 3)).astype(np.uint8)
     masks = [np.random.randint(0, 2, (640, 480)).astype(np.uint8) for _ in range(4)]
-    transformed = transforms(image=img, masks=masks)
-    assert transformed["image"].numpy().shape == (3, 640, 480)
-    assert transformed["masks"][0].shape == (640, 480)
+    for _ in range(10):
+        transformed = transforms(image=img, masks=masks)
+        assert isinstance(transformed["image"], torch.Tensor)
+        assert isinstance(transformed["masks"][0], torch.Tensor)
+        assert transformed["image"].numpy().shape in ((3, 640, 480), (3, 480, 640))
+        assert transformed["masks"][0].shape in ((640, 480), (480, 640))