shared/mongodb/src/classes/visual_communication.py

@@ -12,7 +12,7 @@ from PIL import Image
 from pydantic import BaseModel, ConfigDict
 from pymongo.collection import Collection
 
-from shared.datastore import get_image, put_image
+from shared.datastore import Datastore
 from shared.mongodb.src.classes import ModelData
 
 
@@ -39,10 +39,10 @@ class VisualCommunication(BaseModel):
         """Upload image to MinIO and return MD5 checksum of hashed image."""
         assert isinstance(image, Image.Image)
         assert isinstance(minio_client, Minio)
-        object_name = put_image(
-            client=minio_client,
-            image=image,
-        )
+        with Datastore() as ds:
+            object_name = ds.put_image(
+                image=image,
+            )
         return object_name
 
     @classmethod
@@ -91,10 +91,10 @@ class VisualCommunication(BaseModel):
         """Load image data from minio."""
         assert isinstance(minio_client, Minio)
         # get image from minio
-        image = get_image(
-            client=minio_client,
-            object_name=self.object_name,
-        )
+        with Datastore() as ds:
+            image = ds.get_image(
+                object_name=self.object_name,
+            )
         return image
 
     def save_to_mongo(self, collection: Collection) -> None: