315 lines
10 KiB
Python
315 lines
10 KiB
Python
# built-in dependencies
|
|
import os
|
|
import json
|
|
import hashlib
|
|
import struct
|
|
from datetime import datetime, timezone
|
|
from typing import Any, Dict, List, Optional, Union
|
|
|
|
# 3rd party dependencies
|
|
import numpy as np
|
|
|
|
# project dependencies
|
|
from deepface.modules.database.types import Database
|
|
from deepface.modules.exceptions import DuplicateEntryError
|
|
from deepface.commons.logger import Logger
|
|
|
|
logger = Logger()
|
|
|
|
|
|
# pylint: disable=too-many-positional-arguments, too-many-instance-attributes
|
|
class MongoDbClient(Database):
|
|
"""
|
|
MongoDB equivalent of PostgresClient for DeepFace embeddings storage.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
connection_details: Optional[Union[str, Dict[str, Any]]] = None,
|
|
connection: Any = None,
|
|
db_name: str = "deepface",
|
|
) -> None:
|
|
try:
|
|
from pymongo import MongoClient, ASCENDING
|
|
from pymongo.errors import DuplicateKeyError, BulkWriteError
|
|
from bson import Binary
|
|
except (ModuleNotFoundError, ImportError) as e:
|
|
raise ValueError(
|
|
"pymongo is an optional dependency. Please install it as `pip install pymongo`"
|
|
) from e
|
|
self.MongoClient = MongoClient
|
|
self.ASCENDING = ASCENDING
|
|
self.DuplicateKeyError = DuplicateKeyError
|
|
self.BulkWriteError = BulkWriteError
|
|
self.Binary = Binary
|
|
|
|
if connection is not None:
|
|
self.client = connection
|
|
else:
|
|
self.conn_details = connection_details or os.environ.get("DEEPFACE_MONGO_URI")
|
|
if not self.conn_details:
|
|
raise ValueError(
|
|
"MongoDB connection information not found. "
|
|
"Please provide connection_details or set DEEPFACE_MONGO_URI"
|
|
)
|
|
|
|
if isinstance(self.conn_details, str):
|
|
self.client = MongoClient(self.conn_details)
|
|
else:
|
|
self.client = MongoClient(**self.conn_details)
|
|
|
|
self.db = self.client[db_name]
|
|
self.embeddings = self.db.embeddings
|
|
self.embeddings_index = self.db.embeddings_index
|
|
self.counters = self.db.counters
|
|
self.initialize_database()
|
|
|
|
def close(self) -> None:
|
|
"""Close MongoDB connection."""
|
|
self.client.close()
|
|
|
|
def initialize_database(self, **kwargs: Any) -> None:
|
|
"""
|
|
Ensure required MongoDB indexes exist.
|
|
"""
|
|
|
|
# Unique constraint for embeddings
|
|
self.embeddings.create_index(
|
|
[("face_hash", self.ASCENDING), ("embedding_hash", self.ASCENDING)],
|
|
unique=True,
|
|
name="uniq_face_embedding",
|
|
)
|
|
|
|
# Unique constraint for embeddings_index
|
|
self.embeddings_index.create_index(
|
|
[
|
|
("model_name", self.ASCENDING),
|
|
("detector_backend", self.ASCENDING),
|
|
("align", self.ASCENDING),
|
|
("l2_normalized", self.ASCENDING),
|
|
],
|
|
unique=True,
|
|
name="uniq_index_config",
|
|
)
|
|
|
|
# counters collection for auto-incrementing IDs
|
|
if not self.counters.find_one({"_id": "embedding_id"}):
|
|
self.counters.insert_one({"_id": "embedding_id", "seq": 0})
|
|
|
|
logger.debug("MongoDB indexes ensured.")
|
|
|
|
def upsert_embeddings_index(
|
|
self,
|
|
model_name: str,
|
|
detector_backend: str,
|
|
aligned: bool,
|
|
l2_normalized: bool,
|
|
index_data: bytes,
|
|
) -> None:
|
|
"""
|
|
Upsert embeddings index into MongoDB.
|
|
Args:
|
|
model_name (str): Name of the model.
|
|
detector_backend (str): Name of the detector backend.
|
|
aligned (bool): Whether the embeddings are aligned.
|
|
l2_normalized (bool): Whether the embeddings are L2 normalized.
|
|
index_data (bytes): Serialized index data.
|
|
"""
|
|
self.embeddings_index.update_one(
|
|
{
|
|
"model_name": model_name,
|
|
"detector_backend": detector_backend,
|
|
"align": aligned,
|
|
"l2_normalized": l2_normalized,
|
|
},
|
|
{
|
|
"$set": {
|
|
"index_data": self.Binary(index_data),
|
|
"updated_at": datetime.now(timezone.utc),
|
|
},
|
|
"$setOnInsert": {
|
|
"created_at": datetime.now(timezone.utc),
|
|
},
|
|
},
|
|
upsert=True,
|
|
)
|
|
|
|
def get_embeddings_index(
|
|
self,
|
|
model_name: str,
|
|
detector_backend: str,
|
|
aligned: bool,
|
|
l2_normalized: bool,
|
|
) -> bytes:
|
|
"""
|
|
Retrieve embeddings index from MongoDB.
|
|
Args:
|
|
model_name (str): Name of the model.
|
|
detector_backend (str): Name of the detector backend.
|
|
aligned (bool): Whether the embeddings are aligned.
|
|
l2_normalized (bool): Whether the embeddings are L2 normalized.
|
|
Returns:
|
|
bytes: Serialized index data.
|
|
"""
|
|
doc = self.embeddings_index.find_one(
|
|
{
|
|
"model_name": model_name,
|
|
"detector_backend": detector_backend,
|
|
"align": aligned,
|
|
"l2_normalized": l2_normalized,
|
|
},
|
|
{"index_data": 1},
|
|
)
|
|
|
|
if not doc:
|
|
raise ValueError(
|
|
"No Embeddings index found for the specified parameters "
|
|
f"{model_name=}, {detector_backend=}, "
|
|
f"{aligned=}, {l2_normalized=}. "
|
|
"You must run build_index first."
|
|
)
|
|
|
|
return bytes(doc["index_data"])
|
|
|
|
def insert_embeddings(self, embeddings: List[Dict[str, Any]], batch_size: int = 100) -> int:
|
|
"""
|
|
Insert embeddings into MongoDB.
|
|
Args:
|
|
embeddings (List[Dict[str, Any]]): List of embedding records to insert.
|
|
batch_size (int): Number of records to insert in each batch.
|
|
Returns:
|
|
int: Number of embeddings successfully inserted.
|
|
"""
|
|
if not embeddings:
|
|
raise ValueError("No embeddings to insert.")
|
|
|
|
docs: List[Dict[str, Any]] = []
|
|
|
|
for e in embeddings:
|
|
face = e["face"]
|
|
face_shape = list(face.shape)
|
|
|
|
binary_face_data = self.Binary(face.astype(np.float32).tobytes())
|
|
|
|
embedding_bytes = struct.pack(f'{len(e["embedding"])}d', *e["embedding"])
|
|
|
|
face_hash = hashlib.sha256(json.dumps(face.tolist()).encode()).hexdigest()
|
|
embedding_hash = hashlib.sha256(embedding_bytes).hexdigest()
|
|
|
|
int_id = self.counters.find_one_and_update(
|
|
{"_id": "embedding_id"}, {"$inc": {"seq": 1}}, upsert=True, return_document=True
|
|
)["seq"]
|
|
|
|
docs.append(
|
|
{
|
|
"sequence": int_id,
|
|
"img_name": e["img_name"],
|
|
"face": binary_face_data,
|
|
"face_shape": face_shape,
|
|
"model_name": e["model_name"],
|
|
"detector_backend": e["detector_backend"],
|
|
"aligned": e["aligned"],
|
|
"l2_normalized": e["l2_normalized"],
|
|
"embedding": e["embedding"],
|
|
"face_hash": face_hash,
|
|
"embedding_hash": embedding_hash,
|
|
"created_at": datetime.now(timezone.utc),
|
|
}
|
|
)
|
|
|
|
inserted = 0
|
|
try:
|
|
for i in range(0, len(docs), batch_size):
|
|
result = self.embeddings.insert_many(docs[i : i + batch_size], ordered=False)
|
|
inserted += len(result.inserted_ids)
|
|
except (self.DuplicateKeyError, self.BulkWriteError) as e:
|
|
if len(docs) == 1:
|
|
logger.warn("Duplicate detected for extracted face and embedding.")
|
|
return inserted
|
|
raise DuplicateEntryError(
|
|
f"Duplicate detected for extracted face and embedding in {i}-th batch"
|
|
) from e
|
|
|
|
return inserted
|
|
|
|
def fetch_all_embeddings(
|
|
self,
|
|
model_name: str,
|
|
detector_backend: str,
|
|
aligned: bool,
|
|
l2_normalized: bool,
|
|
batch_size: int = 1000,
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
Fetch all embeddings from MongoDB based on specified parameters.
|
|
Args:
|
|
model_name (str): Name of the model.
|
|
detector_backend (str): Name of the detector backend.
|
|
aligned (bool): Whether the embeddings are aligned.
|
|
l2_normalized (bool): Whether the embeddings are L2 normalized.
|
|
batch_size (int): Number of records to fetch in each batch.
|
|
Returns:
|
|
List[Dict[str, Any]]: List of embedding records.
|
|
"""
|
|
|
|
cursor = self.embeddings.find(
|
|
{
|
|
"model_name": model_name,
|
|
"detector_backend": detector_backend,
|
|
"aligned": aligned,
|
|
"l2_normalized": l2_normalized,
|
|
},
|
|
{
|
|
"_id": 1,
|
|
"sequence": 1,
|
|
"img_name": 1,
|
|
"embedding": 1,
|
|
},
|
|
batch_size=batch_size,
|
|
).sort("sequence", self.ASCENDING)
|
|
|
|
results: List[Dict[str, Any]] = []
|
|
for doc in cursor:
|
|
results.append(
|
|
{
|
|
"_id": str(doc["_id"]),
|
|
"id": doc["sequence"],
|
|
"img_name": doc["img_name"],
|
|
"embedding": doc["embedding"],
|
|
"model_name": model_name,
|
|
"detector_backend": detector_backend,
|
|
"aligned": aligned,
|
|
"l2_normalized": l2_normalized,
|
|
}
|
|
)
|
|
|
|
return results
|
|
|
|
def search_by_id(
|
|
self,
|
|
ids: Union[List[str], List[int]],
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
Search records by their IDs.
|
|
"""
|
|
cursor = self.embeddings.find(
|
|
{"sequence": {"$in": ids}},
|
|
{
|
|
"_id": 1,
|
|
"sequence": 1,
|
|
"img_name": 1,
|
|
},
|
|
)
|
|
|
|
results: List[Dict[str, Any]] = []
|
|
for doc in cursor:
|
|
results.append(
|
|
{
|
|
"_id": str(doc["_id"]),
|
|
"id": doc["sequence"],
|
|
"img_name": doc["img_name"],
|
|
}
|
|
)
|
|
|
|
return results
|