File size: 2,179 Bytes
69def8e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | import numpy as np
from typing import List, Dict, Any, Optional
import os
class FaceGallery:
"""In-memory Face Identification Gallery for 1:N face searching."""
def __init__(self):
self.identities: List[Dict[str, Any]] = []
def enroll(self, name: str, embedding: np.ndarray, image_path: Optional[str] = None, metadata: Optional[Dict] = None) -> None:
"""Enroll a new identity with name, face embedding vector, and image reference."""
# Normalize embedding for cosine similarity
norm = np.linalg.norm(embedding)
norm_emb = embedding / (norm + 1e-10) if norm > 0 else embedding
self.identities.append({
"name": name,
"embedding": norm_emb,
"image_path": image_path,
"metadata": metadata or {}
})
def search(self, query_embedding: np.ndarray, top_k: int = 3, threshold: float = 0.35) -> List[Dict[str, Any]]:
"""Search query embedding against all enrolled identities.
Returns top_k results sorted by cosine similarity descending.
"""
if not self.identities:
return []
norm = np.linalg.norm(query_embedding)
q_norm = query_embedding / (norm + 1e-10) if norm > 0 else query_embedding
gallery_embs = np.vstack([id_dict["embedding"] for id_dict in self.identities]) # (N, D)
similarities = np.dot(gallery_embs, q_norm) # (N,)
results = []
for idx, score in enumerate(similarities):
item = self.identities[idx]
sim_score = float(score)
results.append({
"name": item["name"],
"similarity": sim_score,
"is_match": bool(sim_score >= threshold),
"image_path": item["image_path"],
"metadata": item.get("metadata", {})
})
# Sort descending by similarity score
results.sort(key=lambda x: x["similarity"], reverse=True)
return results[:top_k]
def clear(self) -> None:
self.identities.clear()
def count(self) -> int:
return len(self.identities)
|