src/database/classes.py

@@ -3,6 +3,7 @@ from pydantic import BaseModel, field_validator, field_serializer
 from PIL import Image
 from io import BytesIO
 from pathlib import Path
+from base64 import b64encode
 
 from src.model_experiential import ExperientialModelOutput
 from src.model_interpersonal import (
@@ -36,7 +37,7 @@ class ModelOutputs(BaseModel):
 
 class VisualCommunication(BaseModel):
     name: str
-    image: Image.Image | BytesIO | bytes
+    image: Image.Image
     annotation: ModelOutputs | None = None
     prediction: ModelOutputs | None = None
 
@@ -73,3 +74,15 @@ class VisualCommunication(BaseModel):
 
     def __repr__(self) -> str:
         return f"{self.classname()}(name='{self.name}')"
+    
+    def webencoded_image(self) -> str:
+        """Convert image to be displayed on webpage."""
+        # convert images to bytes string
+        buffer = BytesIO()
+        self.image.save(buffer, format="png")
+        img_enc = b64encode(buffer.getvalue()).decode("utf-8")
+        return img_enc
+
+
+class NoDocumentFoundException(Exception):
+    pass