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,12 @@
+from __future__ import annotations
+
+from .classes import ModelOutputs
+from .classes import NoDocumentFoundException
+from .classes import VisualCommunication
+from .database import connect
+from .utils import get_visual_communication
+from .utils import total_annotated
+from .utils import total_documents
+from .utils import upsert_annotations
+from .utils import upsert_predictions
+from .utils import upsert_visual_communication
@@ -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()
@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+import logging
+import os
+
+from dotenv import load_dotenv
+from pymongo import MongoClient
+
+
+def connect():
+    """Connect to MongoDB."""
+    # load env vars
+    load_dotenv()
+    necessary_env_vars = [
+        'MONGO_HOST',
+        'MONGO_DB',
+        'MONGO_COLLECTION',
+    ]
+    for env_var in necessary_env_vars:
+        assert env_var in os.environ, f"{env_var} not found"
+    # connect to database
+    client = MongoClient(os.getenv('MONGO_HOST'))
+    db = client[os.getenv('MONGO_DB')]
+    collection = db[os.getenv('MONGO_COLLECTION')]
+    collection.create_index('name', unique=True)
+    logging.info('connected to database')
+    return collection, db, client
@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+import logging
+
+from pymongo.collection import Collection
+
+from .classes import ModelOutputs
+from .classes import NoDocumentFoundException
+from .classes import VisualCommunication
+
+
+def total_documents(
+    collection: Collection,
+) -> int:
+    """Get total number of documents in database."""
+    return collection.count_documents(filter={})
+
+
+def total_annotated(
+    collection: Collection,
+) -> int:
+    """Get total number of annotated documents in database."""
+    query = {
+        'annotation': {
+            '$ne': None,
+        },
+    }
+    return collection.count_documents(filter=query)
+
+
+def get_visual_communication(
+    collection: Collection,
+    with_annotation: bool = False,
+) -> VisualCommunication:
+    """Get a random visual communication from the database."""
+    query = {}
+    if with_annotation:
+        query['annotation'] = {'$ne': None}
+    else:
+        query['annotation'] = {'$eq': None}
+    data = collection.aggregate([
+        {
+            '$match': query,     # find using filters
+        },
+        {
+            '$sample': {
+                'size': 1,   # get one random
+            },
+        },
+    ])
+    data_list = list(data)   # read data from cursor object
+    if len(data_list) == 0:
+        logging.error('failed getting visual communication')
+        raise NoDocumentFoundException()
+    logging.info('finished')
+    return VisualCommunication.model_validate(data_list[0])
+
+
+def upsert_predictions(
+    collection: Collection,
+    vis_com_name: str,
+    predictions: ModelOutputs,
+) -> None:
+    """Upsert prediction data in the database."""
+    query = {
+        'name': vis_com_name,
+    }
+    update = {
+        '$set': {
+            'prediction': predictions.model_dump(),
+        },
+    }
+    res = collection.update_one(
+        filter=query,
+        update=update,
+        upsert=True,
+    )
+    logging.debug('upserted document: %s', res)
+    logging.info('finished')
+
+
+def upsert_annotations(
+    collection: Collection,
+    vis_com_name: str,
+    annotations: ModelOutputs,
+) -> None:
+    """Upserts annotation data in the database."""
+    query = {
+        'name': vis_com_name,
+    }
+    update = {
+        '$set': {
+            'annotation': annotations.model_dump(),
+        },
+    }
+    res = collection.update_one(
+        filter=query,
+        update=update,
+        upsert=True,
+    )
+    logging.info('upserted document: %s', res)
+    logging.info('finished')
+
+
+def upsert_visual_communication(
+    collection: Collection,
+    visual_communication_list: list[VisualCommunication],
+) -> bool:
+    """
+    Upsert VisualCommunication object in the database.
+    Returns bool stating success.
+    """
+    response = collection.insert_many(
+        [
+            vis_com.model_dump()
+            for vis_com
+            in visual_communication_list
+        ],
+    )
+    return response.acknowledged