updated database interface

src/database/__init__.py

@@ -0,0 +1,5 @@
+from .classes import (
+    ModelOutputs,
+    VisualCommunication
+)
+from .database import connect

src/database/classes.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+from pydantic import BaseModel, field_validator
+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: bytes
+    annotation: ModelOutputs | None = None
+    prediction: ModelOutputs | None = None
+    
+    @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)
+        return VisualCommunication(name=name, image=image)
+    
+    @field_validator("image", mode="before")
+    @classmethod
+    def convert_to_bytes(cls, raw: Image.Image | BytesIO | bytes) -> bytes:
+        if isinstance(raw, Image.Image):
+            raw = raw.tobytes()
+        if isinstance(raw, BytesIO):
+            raw = raw.read()
+        return raw
+    
+    def __repr__(self) -> str:
+        return f"{self.classname()}(name='{self.name}')"
+
+    @property
+    def image(self) -> Image.Image:
+        return Image.open(BytesIO(self.image))

src/database/database.py

@@ -0,0 +1,22 @@
+from pymongo import MongoClient
+from dotenv import load_dotenv
+import os
+
+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")]
+    return collection, db, client
+
+