defined serialization of images

This commit is contained in:
Brian Bjarke Jensen
2024-02-21 20:05:27 +01:00
parent a3b95f88ba
commit 68417a7afc
+18 -12
View File
@@ -1,5 +1,5 @@
from __future__ import annotations from __future__ import annotations
from pydantic import BaseModel, field_validator from pydantic import BaseModel, field_validator, field_serializer
from PIL import Image from PIL import Image
from io import BytesIO from io import BytesIO
from pathlib import Path from pathlib import Path
@@ -36,10 +36,13 @@ class ModelOutputs(BaseModel):
class VisualCommunication(BaseModel): class VisualCommunication(BaseModel):
name: str name: str
image: bytes image: Image.Image | BytesIO | bytes
annotation: ModelOutputs | None = None annotation: ModelOutputs | None = None
prediction: ModelOutputs | None = None prediction: ModelOutputs | None = None
class Config:
arbitrary_types_allowed = True
@classmethod @classmethod
def classname(cls) -> str: def classname(cls) -> str:
"""Return classname.""" """Return classname."""
@@ -50,20 +53,23 @@ class VisualCommunication(BaseModel):
"""Instantiate from file.""" """Instantiate from file."""
name = path.stem name = path.stem
image = Image.open(path) image = Image.open(path)
image.load()
return VisualCommunication(name=name, image=image) return VisualCommunication(name=name, image=image)
@field_serializer("image")
def serialize_image(image: Image.Image) -> bytes:
buffer = BytesIO()
image.save(buffer, format="JPEG")
return buffer.getvalue()
@field_validator("image", mode="before") @field_validator("image", mode="before")
@classmethod @classmethod
def convert_to_bytes(cls, raw: Image.Image | BytesIO | bytes) -> bytes: def convert_to_image(cls, image: Image.Image | BytesIO | bytes) -> Image.Image:
if isinstance(raw, Image.Image): if isinstance(image, bytes):
raw = raw.tobytes() image = BytesIO(image)
if isinstance(raw, BytesIO): if isinstance(image, BytesIO):
raw = raw.read() image = Image.open(image)
return raw return image
def __repr__(self) -> str: def __repr__(self) -> str:
return f"{self.classname()}(name='{self.name}')" return f"{self.classname()}(name='{self.name}')"
@property
def image(self) -> Image.Image:
return Image.open(BytesIO(self.image))