chore: import upstream snapshot with attribution

python/ray/data/tests/datasource/test_image.py

@@ -0,0 +1,198 @@
+import os
+from typing import Dict
+
+import numpy as np
+import pytest
+from fsspec.implementations.local import LocalFileSystem
+from PIL import Image
+
+import ray
+from ray.data._internal.datasource.image_datasource import (
+    ImageDatasource,
+    ImageFileMetadataProvider,
+)
+from ray.data._internal.tensor_extensions.arrow import (
+    get_arrow_extension_fixed_shape_tensor_types,
+)
+from ray.data.tests.conftest import *  # noqa
+from ray.tests.conftest import *  # noqa
+
+
+class TestReadImages:
+    def test_basic(self, ray_start_regular_shared):
+        # "simple" contains three 32x32 RGB images.
+        ds = ray.data.read_images("example://image-datasets/simple")
+        assert ds.schema().names == ["image"]
+        column_type = ds.schema().types[0]
+        assert isinstance(column_type, get_arrow_extension_fixed_shape_tensor_types())
+        assert all(record["image"].shape == (32, 32, 3) for record in ds.take())
+
+    @pytest.mark.parametrize("num_threads", [-1, 0, 1, 2, 4])
+    def test_multi_threading(self, ray_start_regular_shared, num_threads, monkeypatch):
+        monkeypatch.setattr(
+            ray.data._internal.datasource.image_datasource.ImageDatasource,
+            "_NUM_THREADS_PER_TASK",
+            num_threads,
+        )
+        ds = ray.data.read_images(
+            "example://image-datasets/simple",
+            override_num_blocks=1,
+            include_paths=True,
+        )
+        paths = [item["path"][-len("image1.jpg") :] for item in ds.take_all()]
+        if num_threads > 1:
+            # If there are more than 1 threads, the order is not guaranteed.
+            paths = sorted(paths)
+        expected_paths = ["image1.jpg", "image2.jpg", "image3.jpg"]
+        assert paths == expected_paths
+
+    def test_size(self, ray_start_regular_shared):
+        # "different-sizes" contains RGB images with different heights and widths.
+        ds = ray.data.read_images(
+            "example://image-datasets/different-sizes", size=(32, 32)
+        )
+        assert all(record["image"].shape == (32, 32, 3) for record in ds.take())
+
+    def test_different_sizes(self, ray_start_regular_shared):
+        ds = ray.data.read_images("example://image-datasets/different-sizes")
+        assert sorted(record["image"].shape for record in ds.take()) == [
+            (16, 16, 3),
+            (32, 32, 3),
+            (64, 64, 3),
+        ]
+
+    @pytest.mark.parametrize("size", [(-32, 32), (32, -32), (-32, -32)])
+    def test_invalid_size(self, ray_start_regular_shared, size):
+        with pytest.raises(ValueError):
+            ray.data.read_images("example://image-datasets/simple", size=size)
+
+    @pytest.mark.parametrize(
+        "mode, expected_shape", [("L", (32, 32)), ("RGB", (32, 32, 3))]
+    )
+    def test_mode(
+        self,
+        mode,
+        expected_shape,
+        ray_start_regular_shared,
+    ):
+        # "different-modes" contains 32x32 images with modes "CMYK", "L", and "RGB"
+        ds = ray.data.read_images("example://image-datasets/different-modes", mode=mode)
+        assert all([record["image"].shape == expected_shape for record in ds.take()])
+
+    def test_e2e_prediction(self, shutdown_only):
+        import torch
+        from torchvision import transforms
+        from torchvision.models import resnet18
+
+        ray.shutdown()
+        ray.init(num_cpus=2)
+
+        dataset = ray.data.read_images("example://image-datasets/simple")
+        transform = transforms.ToTensor()
+
+        def preprocess(batch: Dict[str, np.ndarray]):
+            return {"out": np.stack([transform(image) for image in batch["image"]])}
+
+        dataset = dataset.map_batches(preprocess, batch_format="numpy")
+
+        class Predictor:
+            def __init__(self):
+                self.model = resnet18(pretrained=True)
+
+            def __call__(self, batch: Dict[str, np.ndarray]):
+                with torch.inference_mode():
+                    torch_tensor = torch.as_tensor(batch["out"])
+                    return {"prediction": self.model(torch_tensor)}
+
+        predictions = dataset.map_batches(
+            Predictor, compute=ray.data.ActorPoolStrategy(min_size=1), batch_size=4096
+        )
+
+        for _ in predictions.iter_batches():
+            pass
+
+    @pytest.mark.parametrize(
+        "image_size,image_mode,expected_size,expected_ratio",
+        [(64, "RGB", 30000, 4), (32, "L", 3500, 0.5), (256, "RGBA", 750000, 85)],
+    )
+    def test_data_size_estimate(
+        self,
+        ray_start_regular_shared,
+        image_size,
+        image_mode,
+        expected_size,
+        expected_ratio,
+    ):
+        root = "example://image-datasets/different-sizes"
+        ds = ray.data.read_images(
+            root, size=(image_size, image_size), mode=image_mode, override_num_blocks=1
+        )
+
+        data_size = ds.size_bytes()
+        assert data_size >= 0, "estimated data size is out of expected bound"
+        data_size = ds.materialize().size_bytes()
+        assert data_size >= 0, "actual data size is out of expected bound"
+
+        datasource = ImageDatasource(
+            paths=[root],
+            size=(image_size, image_size),
+            mode=image_mode,
+            filesystem=LocalFileSystem(),
+            partitioning=None,
+            meta_provider=ImageFileMetadataProvider(),
+        )
+        assert (
+            datasource._encoding_ratio >= expected_ratio
+            and datasource._encoding_ratio <= expected_ratio * 1.5
+        ), "encoding ratio is out of expected bound"
+        data_size = datasource.estimate_inmemory_data_size()
+        assert data_size >= 0, "estimated data size is out of expected bound"
+
+    def test_dynamic_block_split(ray_start_regular_shared):
+        ctx = ray.data.context.DataContext.get_current()
+        target_max_block_size = ctx.target_max_block_size
+        # Reduce target max block size to trigger block splitting on small input.
+        # Otherwise we have to generate big input files, which is unnecessary.
+        ctx.target_max_block_size = 1
+        try:
+            root = "example://image-datasets/simple"
+            ds = ray.data.read_images(root, override_num_blocks=1)
+            assert ds._logical_plan.initial_num_blocks() == 1
+            ds = ds.materialize()
+            # Verify dynamic block splitting taking effect to generate more blocks.
+            assert ds._logical_plan.initial_num_blocks() == 3
+
+            # Test union of same datasets
+            union_ds = ds.union(ds, ds, ds).materialize()
+            assert union_ds._logical_plan.initial_num_blocks() == 12
+        finally:
+            ctx.target_max_block_size = target_max_block_size
+
+    def test_unidentified_image_error(ray_start_regular_shared, tmp_path):
+        path = str(tmp_path / "invalid.png")
+        with open(path, "wb") as file:
+            file.write(b"spam")  # Invalid bytes for a PNG file
+
+        with pytest.raises(ValueError):
+            ray.data.read_images(paths=file.name).materialize()
+
+
+class TestWriteImages:
+    def test_write_images(ray_start_regular_shared, tmp_path):
+        ds = ray.data.read_images("example://image-datasets/simple")
+        ds.write_images(
+            path=tmp_path,
+            column="image",
+        )
+
+        assert len(os.listdir(tmp_path)) == ds.count()
+
+        for filename in os.listdir(tmp_path):
+            path = os.path.join(tmp_path, filename)
+            Image.open(path)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))