moved shared functions

2024-03-11 20:36:14 +01:00
parent 9fa3574a7d
commit ded50d8095
4 changed files with 13 additions and 10 deletions
@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+import logging
+from base64 import b64decode
+from base64 import b64encode
+from io import BytesIO
+from pathlib import Path
+
+from PIL import Image
+from pydantic import BaseModel
+from pydantic import field_serializer
+from pydantic import field_validator
+
+from src.model_experiential import (
+    VisualSyntaxModelOutput,
+)
+from src.model_interpersonal import AngleModelOutput
+from src.model_interpersonal import ContactModelOutput
+from src.model_interpersonal import DistanceModelOutput
+from src.model_interpersonal import ModalityColorModelOutput
+from src.model_interpersonal import ModalityDepthModelOutput
+from src.model_interpersonal import ModalityLightingModelOutput
+from src.model_interpersonal import PointOfViewModelOutput
+from src.model_textual import FramingModelOutput
+from src.model_textual import InformationValueModelOutput
+from src.model_textual import SalienceModelOutput
+
+
+class NoDocumentFoundException(Exception):
+    pass
+
+
+class ModelOutputs(BaseModel):
+    visual_syntax: VisualSyntaxModelOutput
+    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
+
+    @classmethod
+    def list_fields(cls) -> list[str]:
+        """List options that are stored as attributes."""
+        return list(cls.model_fields.keys())
+
+    @classmethod
+    def from_random(cls) -> ModelOutputs:
+        """Instantiate with random numbers."""
+        kwargs = {
+            field: field_info.annotation.from_random()  # type: ignore
+            for field, field_info
+            in cls.model_fields.items()
+        }
+        return cls(**kwargs)
+
+    @classmethod
+    def from_annotations(
+        cls,
+        visual_syntax: str,
+        contact: str,
+        angle: str,
+        point_of_view: str,
+        distance: str,
+        modality_lighting: str,
+        modality_color: str,
+        modality_depth: str,
+        information_value: str,
+        framing: str,
+        salience: str,
+    ) -> ModelOutputs:
+        """Instantiate from annotation."""
+        kwargs = {
+            'visual_syntax': VisualSyntaxModelOutput
+            .from_choice(visual_syntax),
+            'contact': ContactModelOutput
+            .from_choice(contact),
+            'angle': AngleModelOutput
+            .from_choice(angle),
+            'point_of_view': PointOfViewModelOutput
+            .from_choice(point_of_view),
+            'distance': DistanceModelOutput
+            .from_choice(distance),
+            'modality_lighting': ModalityLightingModelOutput
+            .from_choice(modality_lighting),
+            'modality_color': ModalityColorModelOutput
+            .from_choice(modality_color),
+            'modality_depth': ModalityDepthModelOutput
+            .from_choice(modality_depth),
+            'information_value': InformationValueModelOutput
+            .from_choice(information_value),
+            'framing': FramingModelOutput
+            .from_choice(framing),
+            'salience': SalienceModelOutput
+            .from_choice(salience),
+        }
+        return cls(**kwargs)
+
+
+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)
+
+    @classmethod
+    def decode_image(cls, content: str) -> Image.Image:
+        """Decode image."""
+        _, content_data = content.split(',')
+        return Image.open(BytesIO(b64decode(content_data)))
+
+    @field_serializer('image')
+    def serialize_image(image: Image.Image) -> bytes:   # type: ignore
+        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 f"data:image/png;base64, {img_enc}"
+
+    def generate_random_prediction(self, force: bool = False) -> None:
+        """Generate random prediction values."""
+        if not force and self.prediction is not None:
+            logging.warning('set force=True to overwrite existing values.')
+        self.prediction = ModelOutputs.from_random()