Дипломная работа

Дипломная работа/app/workers/mongo_worker.py

@@ -0,0 +1,79 @@
+import os
+
+import gridfs
+import numpy as np
+import cv2
+
+from pymongo import MongoClient
+from dotenv import load_dotenv
+
+from schemas.base_schemas import *
+from schemas.api_schemas import *
+
+class MongoWorker():
+    def __init__(self):
+        load_dotenv()
+        self.client = MongoClient(host=os.getenv('MONGO_HOST'),
+                                  port=int(os.getenv('MONGO_PORT')),
+                                  username=os.getenv('MONGO_USER'),
+                                  password=os.getenv('MONGO_PASS'))
+        self.db = self.client["Face_recognition"]
+        self.persons = self.db["persons"]
+        self.counters = self.db["counters"]
+        self.gfs = gridfs.GridFS(self.db)
+
+    def count_persons(self):
+        docs_count = self.persons.count_documents({})
+        return docs_count
+
+    def get_and_update_counter(self, counter_name: str) -> int:
+        counter = self.counters.find_one_and_update(
+            {"counter_name": counter_name},
+            {"$inc": {"counter": 1}},
+            upsert=True,
+            return_document=True)
+        return counter["counter"]
+
+    def add_person(self, person_name:str, encodings:ndarray[Any, dtype], image_name:str, face_img:np.ndarray) -> BaseResponse:
+        try:
+            new_person_id = self.get_and_update_counter("persons")
+            new_person = Person(
+                person_id = new_person_id,
+                person_name = person_name,
+                encodings = encodings.tolist(),
+                image_name = f"{new_person_id}_{image_name}")
+
+            file_id = self.save_image_to_gridfs(face_img, image_name)
+            mongo_result = self.persons.insert_one(new_person.model_dump())
+            return BaseResponse(result={"file_id": file_id, "mongo_result": mongo_result})
+        except Exception as exception:
+            print(exception)
+            return BaseResponse(result=exception, error=True)
+
+    def save_image_to_gridfs(self, image_array, filename):
+        success, encoded_image = cv2.imencode('.png', image_array)
+        if not success:
+            raise ValueError("Не удалось закодировать изображение")
+        file_id = self.gfs.put(encoded_image.tobytes(), filename=filename)
+        return file_id
+
+    def get_all_encodings(self):
+        result = {}
+        projection = {"person_id": 1, "encodings": 1, "_id": 0}
+
+        for document in self.persons.find({}, projection):
+            entry = {document["person_id"]: document["encodings"]}
+            result.update(entry)
+
+        return result
+
+    def retrieve_image_from_gridfs(self, file_id):
+        try:
+            gridfs_file = self.gfs.get(file_id)
+            image_bytes = gridfs_file.read()
+            image_array = cv2.imdecode(np.frombuffer(image_bytes, np.uint8), cv2.IMREAD_COLOR)
+            if image_array is None:
+                raise ValueError("Не удалось декодировать изображение")
+            return image_array
+        except gridfs.errors.NoFile:
+            raise FileNotFoundError(f"Файл с ID {file_id} не найден в GridFS")