Automatic Serialization/deserialization of images (#361)

* image input/output

* format

Author: aniketmaurya

src/litserve/schema/__init__.py

@@ -0,0 +1,69 @@
+import base64
+from io import BytesIO
+from typing import TYPE_CHECKING, Any
+
+from pydantic import BaseModel, field_serializer, field_validator
+
+if TYPE_CHECKING:
+    from PIL import Image
+
+
+class ImageInput(BaseModel):
+    image_data: str
+
+    @field_validator("image_data")
+    def validate_base64(cls, value: str) -> str:  # noqa
+        """Ensure the string is a valid Base64."""
+        try:
+            base64.b64decode(value)
+        except base64.binascii.Error:
+            raise ValueError("Invalid Base64 string.")
+        return value
+
+    def get_image(self) -> "Image.Image":
+        """Decode the Base64 string and return a PIL Image object."""
+        try:
+            from PIL import Image, UnidentifiedImageError
+        except ImportError:
+            raise ImportError("Pillow is required to use the ImageInput schema. Install it with `pip install Pillow`.")
+        try:
+            decoded_data = base64.b64decode(self.image_data)
+            return Image.open(BytesIO(decoded_data))
+        except base64.binascii.Error as e:
+            raise ValueError(f"Error decoding Base64 string: {e}")
+        except UnidentifiedImageError as e:
+            raise ValueError(f"Error loading image from decoded data: {e}")
+
+
+class ImageOutput(BaseModel):
+    image: Any
+
+    @field_serializer("image")
+    def serialize_image(self, image: Any, _info):
+        """
+        Serialize a PIL Image into a base64 string.
+        Args:
+            image (Any): The image object to serialize.
+            _info: Metadata passed during serialization (not used here).
+
+        Returns:
+            str: Base64-encoded image string.
+        """
+        try:
+            from PIL import Image
+        except ImportError:
+            raise ImportError("Pillow is required to use the ImageOutput schema. Install it with `pip install Pillow`.")
+
+        if not isinstance(image, Image.Image):
+            raise TypeError(f"Expected a PIL Image, got {type(image)}")
+
+        # Save the image to a BytesIO buffer
+        buffer = BytesIO()
+        image.save(buffer, format="PNG")  # Default format is PNG
+        buffer.seek(0)
+
+        # Encode the buffer content to base64
+        base64_bytes = base64.b64encode(buffer.read())
+
+        # Decode to string for JSON serialization
+        return base64_bytes.decode("utf-8")

src/litserve/schema/image.py

@@ -0,0 +1,41 @@
+import base64
+import io
+import os
+
+import numpy as np
+from fastapi.testclient import TestClient
+from PIL import Image
+
+import litserve as ls
+from litserve.schema.image import ImageInput, ImageOutput
+from litserve.utils import wrap_litserve_start
+
+
+class ImageAPI(ls.LitAPI):
+    def setup(self, device):
+        self.model = lambda x: np.array(x) * 2
+
+    def decode_request(self, request: ImageInput):
+        return request.get_image()
+
+    def predict(self, x):
+        return self.model(x)
+
+    def encode_response(self, numpy_image) -> ImageOutput:
+        output = Image.fromarray(np.uint8(numpy_image)).convert("RGB")
+        return ImageOutput(image=output)
+
+
+def test_image_input_output(tmpdir):
+    path = os.path.join(tmpdir, "test.png")
+    server = ls.LitServer(ImageAPI(), accelerator="cpu", devices=1, workers_per_device=1)
+    with wrap_litserve_start(server) as server, TestClient(server.app) as client:
+        Image.new("RGB", (32, 32)).save(path)
+        with open(path, "rb") as image_file:
+            encoded_string = base64.b64encode(image_file.read()).decode("utf-8")
+        response = client.post("/predict", json={"image_data": encoded_string})
+
+        assert response.status_code == 200, f"Unexpected status code: {response.status_code}"
+        image_data = response.json()["image"]
+        image = Image.open(io.BytesIO(base64.b64decode(image_data)))
+        assert image.size == (32, 32), f"Unexpected image size: {image.size}"