updated and added tests

src/database/classes.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from pydantic import BaseModel, field_validator
+from pydantic import BaseModel, field_validator, field_serializer
 from PIL import Image
 from io import BytesIO
 from pathlib import Path
@@ -36,10 +36,13 @@ class ModelOutputs(BaseModel):
 
 class VisualCommunication(BaseModel):
     name: str
-    image: bytes
+    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."""
@@ -50,20 +53,23 @@ class VisualCommunication(BaseModel):
         """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_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 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}')"
-
-    @property
-    def image(self) -> Image.Image:
-        return Image.open(BytesIO(self.image))

src/database/database.py

@@ -17,6 +17,5 @@ def connect():
     client = MongoClient(os.getenv("MONGO_HOST"))
     db = client[os.getenv("MONGO_DB")]
     collection = db[os.getenv("MONGO_COLLECTION")]
+    collection.create_index("name", unique=True)
     return collection, db, client
-
-

tests/test_image_download.py

@@ -0,0 +1,27 @@
+from pathlib import Path
+from dotenv import load_dotenv
+from pymongo import MongoClient
+from typing import List
+import os
+
+from src.database import VisualCommunication, connect
+
+if __name__ == "__main__":
+    # prepare env vars
+    env_path = Path(__file__).parent.parent / "local.env"
+    assert env_path.exists()
+    load_dotenv(env_path)
+    os.environ["MONGO_HOST"] = "localhost"
+    # connect to database
+    collection, db, client = connect()
+    print(client.server_info())
+    # download images
+    data = None
+    for data in collection.find().limit(3):
+        if data is None:
+            print("no document found")
+            break
+        vis_com: VisualCommunication = VisualCommunication.model_validate(data)
+        print(repr(vis_com))
+    if data is not None:
+        print(vis_com.image)

tests/test_image_upload.py

@@ -1,6 +1,6 @@
 from pathlib import Path
 from dotenv import load_dotenv
-from pymongo import MongoClient
+from pymongo.errors import DuplicateKeyError
 import os
 
 from src.database import VisualCommunication, connect
@@ -15,12 +15,19 @@ if __name__ == "__main__":
     vis_com_list = [VisualCommunication.from_file(path) for path in img_path_list]
     for vis_com in vis_com_list:
         print(repr(vis_com))
-    # upload images to database
-    env_path = test_dir.parent / "mongodb.env"
+    # prepare env vars
+    env_path = test_dir.parent / "local.env"
     assert env_path.exists()
     load_dotenv(env_path)
+    os.environ["MONGO_HOST"] = "localhost"
+    # connect to database
     collection, db, client = connect()
     print(client.server_info())
+    # upload images
     for vis_com in vis_com_list:
-        result = collection.insert_one(vis_com.model_dump_json())
+        try:
+            result = collection.insert_one(vis_com.model_dump())
+        except DuplicateKeyError as exc:
+            print("ignoring:\n", exc)
+        else:
             print(f"inserted document: {result}")