wehub-resource-sync
25 lines
705 B
Python
25 lines
705 B
Python
import io
|
|
from typing import Any, Dict
|
|
|
|
import pyarrow
|
|
|
|
from ray.data.datasource.file_datasink import RowBasedFileDatasink
|
|
|
|
|
|
class ImageDatasink(RowBasedFileDatasink):
|
|
def __init__(
|
|
self, path: str, column: str, file_format: str, **file_datasink_kwargs
|
|
):
|
|
super().__init__(path, file_format=file_format, **file_datasink_kwargs)
|
|
|
|
self.column = column
|
|
self.file_format = file_format
|
|
|
|
def write_row_to_file(self, row: Dict[str, Any], file: "pyarrow.NativeFile"):
|
|
from PIL import Image
|
|
|
|
image = Image.fromarray(row[self.column])
|
|
buffer = io.BytesIO()
|
|
image.save(buffer, format=self.file_format)
|
|
file.write(buffer.getvalue())
|