wehub-resource-sync
199 lines
7.3 KiB
Python
199 lines
7.3 KiB
Python
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__]))
|