wehub-resource-sync
80 lines
2.3 KiB
Python
80 lines
2.3 KiB
Python
# flake8: noqa
|
|
# fmt: off
|
|
|
|
from typing import Iterator, Union, List
|
|
|
|
import pyarrow
|
|
|
|
from ray.data.block import Block
|
|
|
|
# __datasource_constructor_start__
|
|
from ray.data.datasource import FileBasedDatasource
|
|
|
|
class ImageDatasource(FileBasedDatasource):
|
|
def __init__(self, paths: Union[str, List[str]], *, mode: str):
|
|
super().__init__(
|
|
paths,
|
|
file_extensions=["png", "jpg", "jpeg", "bmp", "gif", "tiff"],
|
|
)
|
|
|
|
self.mode = mode # Specify read options in the constructor
|
|
# __datasource_constructor_end__
|
|
|
|
# __read_stream_start__
|
|
def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
|
|
import io
|
|
import numpy as np
|
|
from PIL import Image
|
|
from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
|
|
|
|
data = f.readall()
|
|
image = Image.open(io.BytesIO(data))
|
|
image = image.convert(self.mode)
|
|
|
|
# Each block contains one row
|
|
builder = DelegatingBlockBuilder()
|
|
array = np.asarray(image)
|
|
item = {"image": array}
|
|
builder.add(item)
|
|
yield builder.build()
|
|
# __read_stream_end__
|
|
|
|
# __read_datasource_start__
|
|
import ray
|
|
|
|
ds = ray.data.read_datasource(
|
|
ImageDatasource("s3://anonymous@ray-example-data/batoidea", mode="RGB")
|
|
)
|
|
# __read_datasource_end__
|
|
|
|
|
|
from typing import Any, Dict
|
|
import pyarrow
|
|
|
|
# __datasink_constructor_start__
|
|
from ray.data.datasource import RowBasedFileDatasink
|
|
|
|
class ImageDatasink(RowBasedFileDatasink):
|
|
def __init__(self, path: str, column: str, file_format: str):
|
|
super().__init__(path, file_format=file_format)
|
|
|
|
self.column = column
|
|
self.file_format = file_format # Specify write options in the constructor
|
|
# __datasink_constructor_end__
|
|
|
|
# __write_row_to_file_start__
|
|
def write_row_to_file(self, row: Dict[str, Any], file: pyarrow.NativeFile):
|
|
import io
|
|
from PIL import Image
|
|
|
|
# PIL can't write to a NativeFile, so we have to write to a buffer first.
|
|
image = Image.fromarray(row[self.column])
|
|
buffer = io.BytesIO()
|
|
image.save(buffer, format=self.file_format)
|
|
file.write(buffer.getvalue())
|
|
# __write_row_to_file_end__
|
|
|
|
# __write_datasink_start__
|
|
ds.write_datasink(ImageDatasink("/tmp/results", column="image", file_format="png"))
|
|
# __write_datasink_end__
|