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))