diff --git a/src/database/classes.py b/src/database/classes.py index 620142a..39aad3a 100644 --- a/src/database/classes.py +++ b/src/database/classes.py @@ -1,5 +1,5 @@ from __future__ import annotations -from pydantic import BaseModel, field_validator +from pydantic import BaseModel, field_validator, field_serializer from PIL import Image from io import BytesIO from pathlib import Path @@ -36,34 +36,40 @@ class ModelOutputs(BaseModel): class VisualCommunication(BaseModel): name: str - image: bytes + image: Image.Image | BytesIO | bytes annotation: ModelOutputs | None = None prediction: ModelOutputs | None = None - + + class Config: + arbitrary_types_allowed = True + @classmethod def classname(cls) -> str: """Return classname.""" return cls.__name__ - + @classmethod def from_file(cls, path: Path) -> VisualCommunication: """Instantiate from file.""" name = path.stem image = Image.open(path) + image.load() 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") @classmethod - def convert_to_bytes(cls, raw: Image.Image | BytesIO | bytes) -> bytes: - if isinstance(raw, Image.Image): - raw = raw.tobytes() - if isinstance(raw, BytesIO): - raw = raw.read() - return raw - + def convert_to_image(cls, image: Image.Image | BytesIO | bytes) -> Image.Image: + if isinstance(image, bytes): + image = BytesIO(image) + if isinstance(image, BytesIO): + image = Image.open(image) + return image + def __repr__(self) -> str: return f"{self.classname()}(name='{self.name}')" - - @property - def image(self) -> Image.Image: - return Image.open(BytesIO(self.image))