Files
visual_critical_discourse_a…/src/database/classes.py
T

89 lines
2.5 KiB
Python

from __future__ import annotations
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 (
ContactModelOutput,
AngleModelOutput,
PointOfViewModelOutput,
DistanceModelOutput,
ModalityLightingModelOutput,
ModalityColorModelOutput,
ModalityDepthModelOutput
)
from src.model_textual import (
InformationValueModelOutput,
FramingModelOutput,
SalienceModelOutput
)
class ModelOutputs(BaseModel):
experiential: ExperientialModelOutput
contact: ContactModelOutput
angle: AngleModelOutput
point_of_view: PointOfViewModelOutput
distance: DistanceModelOutput
modality_lighting: ModalityLightingModelOutput
modality_color: ModalityColorModelOutput
modality_depth: ModalityDepthModelOutput
information_value: InformationValueModelOutput
framing: FramingModelOutput
salience: SalienceModelOutput
class VisualCommunication(BaseModel):
name: str
image: Image.Image
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_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}')"
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