Add image transformer

flytekit/core/type_engine.py

@@ -874,6 +874,8 @@
             register_bigquery_handlers()
         if is_imported("numpy"):
             from flytekit.types import numpy  # noqa: F401
+        if is_imported("PIL"):
+            from flytekit.types.file import image  # noqa: F401
 
     @classmethod
     def to_literal_type(cls, python_type: Type) -> LiteralType:

flytekit/types/file/image.py

@@ -0,0 +1,82 @@
+import pathlib
+import typing
+from typing import Type
+
+import PIL.Image
+
+from flytekit.core.context_manager import FlyteContext
+from flytekit.core.type_engine import TypeEngine, TypeTransformer, TypeTransformerFailedError
+from flytekit.models.core import types as _core_types
+from flytekit.models.literals import Blob, BlobMetadata, Literal, Scalar
+from flytekit.models.types import LiteralType
+
+
+T = typing.TypeVar("T")
+
+
+class PILImageTransformer(TypeTransformer[T]):
+    """
+    TypeTransformer that supports np.ndarray as a native type.
+    """
+
+    FILE_FORMAT = "image"
+
+    def __init__(self):
+        super().__init__(name="PIL.Image", t=PIL.Image.Image)
+
+    def get_literal_type(self, t: Type[T]) -> LiteralType:
+        return LiteralType(
+            blob=_core_types.BlobType(
+                format=self.FILE_FORMAT, dimensionality=_core_types.BlobType.BlobDimensionality.SINGLE
+            )
+        )
+
+    def to_literal(
+        self, ctx: FlyteContext, python_val: PIL.Image.Image, python_type: Type[T], expected: LiteralType
+    ) -> Literal:
+
+        meta = BlobMetadata(
+            type=_core_types.BlobType(
+                format=self.FILE_FORMAT, dimensionality=_core_types.BlobType.BlobDimensionality.SINGLE
+            )
+        )
+
+        local_path = ctx.file_access.get_random_local_path() + ".png"
+        pathlib.Path(local_path).parent.mkdir(parents=True, exist_ok=True)
+        python_val.save(local_path)
+
+        remote_path = ctx.file_access.get_random_remote_path(local_path)
+        ctx.file_access.put_data(local_path, remote_path, is_multipart=False)
+        return Literal(scalar=Scalar(blob=Blob(metadata=meta, uri=remote_path)))
+
+    def to_python_value(self, ctx: FlyteContext, lv: Literal, expected_python_type: Type[T]) -> PIL.Image.Image:
+        try:
+            uri = lv.scalar.blob.uri
+        except AttributeError:
+            raise TypeTransformerFailedError(f"Cannot convert from {lv} to {expected_python_type}")
+
+        local_path = ctx.file_access.get_random_local_path()
+        ctx.file_access.get_data(uri, local_path, is_multipart=False)
+
+        return PIL.Image.open(local_path)
+
+    def guess_python_type(self, literal_type: LiteralType) -> Type[T]:
+        if (
+            literal_type.blob is not None
+            and literal_type.blob.dimensionality == _core_types.BlobType.BlobDimensionality.SINGLE
+            and literal_type.blob.format == self.FILE_FORMAT
+        ):
+            return PIL.Image.Image
+
+        raise ValueError(f"Transformer {self} cannot reverse {literal_type}")
+
+    def to_html(self, ctx: FlyteContext, python_val: PIL.Image.Image, expected_python_type: Type[T]) -> str:
+        import base64
+        from io import BytesIO
+        buffered = BytesIO()
+        python_val.save(buffered, format="PNG")
+        img_base64 = base64.b64encode(buffered.getvalue()).decode()
+        return f'<img src="data:image/png;base64,{img_base64}" alt="Rendered Image" />'
+
+
+TypeEngine.register(PILImageTransformer())

plugins/flytekit-deck-standard/flytekitplugins/deck/renderer.py

@@ -7,6 +7,7 @@
     import markdown
     import pandas as pd
     import PIL
+    import PIL.Image
     import plotly.express as px
 else:
     pd = lazy_module("pandas")

tests/flytekit/unit/types/file/test_image.py

@@ -0,0 +1,29 @@
+import io
+
+import requests
+
+from flytekit import task, workflow
+import PIL.Image
+
+
+@task(disable_deck=False)
+def t1() -> PIL.Image.Image:
+    url = "https://miro.medium.com/v2/resize:fit:1400/1*0T9PjBnJB9H0Y4qrllkJtQ.png"
+    response = requests.get(url)
+    image_bytes = io.BytesIO(response.content)
+    im = PIL.Image.open(image_bytes)
+    return im
+
+
+@task
+def t2(im: PIL.Image.Image) -> PIL.Image.Image:
+    return im
+
+
+@workflow
+def wf():
+    t2(im=t1())
+
+
+def test_image_transformer():
+    wf()