chore: import upstream snapshot with attribution

python/ray/data/tests/preprocessors/test_torch.py

@@ -0,0 +1,161 @@
+import numpy as np
+import pytest
+import torch
+from torchvision import transforms
+
+import ray
+from ray.data.exceptions import UserCodeException
+from ray.data.preprocessors import TorchVisionPreprocessor
+
+
+class TestTorchVisionPreprocessor:
+    def test_repr(self):
+        class StubTransform:
+            def __call__(self, tensor):
+                return tensor
+
+            def __repr__(self):
+                return "StubTransform()"
+
+        preprocessor = TorchVisionPreprocessor(
+            columns=["spam"], transform=StubTransform()
+        )
+        assert repr(preprocessor) == (
+            "TorchVisionPreprocessor(columns=['spam'], "
+            "output_columns=['spam'], transform=StubTransform())"
+        )
+
+    @pytest.mark.parametrize(
+        "transform",
+        [
+            transforms.ToTensor(),  # `ToTensor` accepts an `np.ndarray` as input
+            transforms.Lambda(lambda tensor: tensor.permute(2, 0, 1)),
+        ],
+    )
+    def test_transform_images(self, transform):
+        dataset = ray.data.from_items(
+            [
+                {"image": np.zeros((32, 32, 3)), "label": 0},
+                {"image": np.zeros((32, 32, 3)), "label": 1},
+            ]
+        )
+        preprocessor = TorchVisionPreprocessor(columns=["image"], transform=transform)
+
+        transformed_dataset = preprocessor.transform(dataset)
+
+        assert transformed_dataset.schema().names == ["image", "label"]
+        transformed_images = [
+            record["image"] for record in transformed_dataset.take_all()
+        ]
+        assert all(image.shape == (3, 32, 32) for image in transformed_images)
+        assert all(image.dtype == np.double for image in transformed_images)
+        labels = {record["label"] for record in transformed_dataset.take_all()}
+        assert labels == {0, 1}
+
+    def test_batch_transform_images(self):
+        dataset = ray.data.from_items(
+            [
+                {"image": np.zeros((32, 32, 3)), "label": 0},
+                {"image": np.zeros((32, 32, 3)), "label": 1},
+            ]
+        )
+        transform = transforms.Compose(
+            [
+                transforms.Lambda(
+                    lambda batch: torch.as_tensor(batch).permute(0, 3, 1, 2)
+                ),
+                transforms.Resize(64),
+            ]
+        )
+        preprocessor = TorchVisionPreprocessor(
+            columns=["image"], transform=transform, batched=True
+        )
+
+        transformed_dataset = preprocessor.transform(dataset)
+
+        assert transformed_dataset.schema().names == ["image", "label"]
+        transformed_images = [
+            record["image"] for record in transformed_dataset.take_all()
+        ]
+        assert all(image.shape == (3, 64, 64) for image in transformed_images)
+        assert all(image.dtype == np.double for image in transformed_images)
+        labels = {record["label"] for record in transformed_dataset.take_all()}
+        assert labels == {0, 1}
+
+    def test_transform_ragged_images(self):
+        dataset = ray.data.from_items(
+            [
+                {"image": np.zeros((16, 16, 3)), "label": 0},
+                {"image": np.zeros((32, 32, 3)), "label": 1},
+            ]
+        )
+        transform = transforms.ToTensor()
+        preprocessor = TorchVisionPreprocessor(columns=["image"], transform=transform)
+
+        transformed_dataset = preprocessor.transform(dataset)
+
+        assert transformed_dataset.schema().names == ["image", "label"]
+        transformed_images = [
+            record["image"] for record in transformed_dataset.take_all()
+        ]
+        assert sorted(image.shape for image in transformed_images) == [
+            (3, 16, 16),
+            (3, 32, 32),
+        ]
+        assert all(image.dtype == np.double for image in transformed_images)
+        labels = {record["label"] for record in transformed_dataset.take_all()}
+        assert labels == {0, 1}
+
+    def test_invalid_transform_raises_value_error(self):
+        dataset = ray.data.from_items(
+            [
+                {"image": np.zeros((32, 32, 3)), "label": 0},
+                {"image": np.zeros((32, 32, 3)), "label": 1},
+            ]
+        )
+        transform = transforms.Lambda(lambda tensor: "BLAH BLAH INVALID")
+        preprocessor = TorchVisionPreprocessor(columns=["image"], transform=transform)
+
+        with pytest.raises((UserCodeException, ValueError)):
+            preprocessor.transform(dataset).materialize()
+
+
+def test_torchvision_preprocessor_serialization():
+    """Test TorchVisionPreprocessor serialization and deserialization functionality."""
+    from torchvision import transforms
+
+    from ray.data.preprocessor import SerializablePreprocessorBase
+
+    # Create preprocessor
+    transform = transforms.Compose([transforms.ToTensor()])
+    preprocessor = TorchVisionPreprocessor(columns=["image"], transform=transform)
+
+    # Serialize using CloudPickle
+    serialized = preprocessor.serialize()
+
+    # Verify it's binary CloudPickle format
+    assert isinstance(serialized, bytes)
+    assert serialized.startswith(SerializablePreprocessorBase.MAGIC_CLOUDPICKLE)
+
+    # Deserialize
+    deserialized = TorchVisionPreprocessor.deserialize(serialized)
+
+    # Verify type and field values
+    assert isinstance(deserialized, TorchVisionPreprocessor)
+    assert deserialized.columns == ["image"]
+    assert isinstance(deserialized.torchvision_transform, type(transform))
+
+    # Verify it works correctly
+    test_data = {"image": np.zeros((32, 32, 3), dtype=np.uint8)}
+    result = deserialized.transform_batch(test_data)
+
+    # Verify transformation was applied - ToTensor converts uint8 [0,255] to float [0.0, 1.0]
+    assert "image" in result
+    assert result["image"].dtype in (np.float32, np.float64)
+    assert result["image"].min() >= 0.0 and result["image"].max() <= 1.0
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-sv", __file__]))