76 lines
2.2 KiB
Python
76 lines
2.2 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 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 | 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_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}')"
|