Archived
79 lines
3.1 KiB
Python
79 lines
3.1 KiB
Python
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") |