""" face_model/analyzer.py — InsightFace wrapper với RAM cache và fallback matching Chức năng: - Khởi tạo InsightFace model (singleton) - Detect khuôn mặt & trích xuất embedding - RAM cache: known_embeddings, names, person_ids, etc. - find_match_ram: cosine similarity fallback khi Supabase RPC thất bại """ import logging import os import cv2 import numpy as np from insightface.app import FaceAnalysis import config logger = logging.getLogger("face_analyzer") class FaceAnalyzer: """Singleton quản lý InsightFace model và RAM embedding cache.""" def __init__(self): self._analyzer: FaceAnalysis | None = None # RAM cache — parallel arrays (index-aligned) self.known_embeddings: list[np.ndarray] = [] self.known_names: list[str] = [] self.known_person_ids: list[str] = [] self.known_embedding_ids: list[str] = [] self.known_mongo_ids: list[str] = [] # ── Model Init ──────────────────────────────────────────────────────── def initialize(self) -> None: """Khởi tạo InsightFace model. Gọi một lần khi khởi động server.""" model_root = os.getenv("INSIGHTFACE_ROOT", "~/.insightface") logger.info( f"[Model] Loading InsightFace '{config.MODEL_NAME}' from root '{model_root}' " f"with provider '{config.MODEL_PROVIDER}'..." ) self._analyzer = FaceAnalysis( name=config.MODEL_NAME, root=model_root, providers=[config.MODEL_PROVIDER], ) self._analyzer.prepare(ctx_id=0, det_size=config.DET_SIZE) logger.info("[Model] InsightFace initialized successfully.") @property def is_ready(self) -> bool: return self._analyzer is not None # ── Face Detection & Embedding ──────────────────────────────────────── def get_faces(self, img: np.ndarray) -> list: """Detect tất cả khuôn mặt trong ảnh. Returns list of face objects.""" if not self.is_ready: raise RuntimeError("Face analyzer not initialized.") return self._analyzer.get(img) # ── RAM Cache ───────────────────────────────────────────────────────── def reload_cache(self, records: list[dict]) -> int: """ Tải lại RAM cache từ danh sách records. records: list[{person_id, full_name, embedding_id, embedding, face_crop_mongo_id}] Returns: số lượng embeddings đã load """ temp_embeddings: list[np.ndarray] = [] temp_names: list[str] = [] temp_person_ids: list[str] = [] temp_embedding_ids: list[str] = [] temp_mongo_ids: list[str] = [] for r in records: emb = r.get("embedding", []) if not emb: continue temp_embeddings.append(np.array(emb, dtype=np.float32)) temp_names.append(r.get("full_name", "")) temp_person_ids.append(r.get("person_id", "")) temp_embedding_ids.append(r.get("embedding_id", "")) temp_mongo_ids.append(r.get("face_crop_mongo_id", "")) self.known_embeddings = temp_embeddings self.known_names = temp_names self.known_person_ids = temp_person_ids self.known_embedding_ids = temp_embedding_ids self.known_mongo_ids = temp_mongo_ids logger.info(f"[Cache] Reloaded {len(self.known_embeddings)} embeddings into RAM.") return len(self.known_embeddings) def reload_from_local_folder(self, folder: str) -> int: """ Fallback: Load embeddings từ thư mục ảnh cục bộ (khi Supabase offline). """ if not os.path.exists(folder): return 0 self.known_embeddings = [] self.known_names = [] self.known_person_ids = [] self.known_embedding_ids = [] self.known_mongo_ids = [] files = sorted([ f for f in os.listdir(folder) if f.lower().endswith((".png", ".jpg", ".jpeg")) ]) for filename in files: img_path = os.path.join(folder, filename) img = cv2.imread(img_path) if img is None: continue faces = self._analyzer.get(img) if faces: self.known_embeddings.append(faces[0].normed_embedding) name = os.path.splitext(filename)[0] self.known_names.append(name) self.known_person_ids.append("") self.known_embedding_ids.append("") self.known_mongo_ids.append("") logger.info(f"[Cache Fallback] Loaded {len(self.known_embeddings)} from '{folder}'.") return len(self.known_embeddings) def add_to_cache( self, embedding: np.ndarray, name: str, person_id: str, embedding_id: str, mongo_id: str, ) -> None: """Thêm một embedding mới vào RAM cache.""" self.known_embeddings.append(embedding) self.known_names.append(name) self.known_person_ids.append(person_id) self.known_embedding_ids.append(embedding_id) self.known_mongo_ids.append(mongo_id) def update_name_in_cache(self, person_id: str, new_name: str) -> None: """Cập nhật tên trong RAM cache sau khi update trên DB.""" for i, pid in enumerate(self.known_person_ids): if pid == person_id: self.known_names[i] = new_name def remove_from_cache(self, person_id: str) -> None: """Xóa tất cả entries của person khỏi RAM cache.""" indices = [i for i, pid in enumerate(self.known_person_ids) if pid == person_id] for i in reversed(indices): self.known_embeddings.pop(i) self.known_names.pop(i) self.known_person_ids.pop(i) self.known_embedding_ids.pop(i) self.known_mongo_ids.pop(i) @property def total(self) -> int: return len(self.known_embeddings) # ── RAM-based Matching (fallback) ───────────────────────────────────── def find_match_ram( self, embedding: np.ndarray, threshold: float, ) -> tuple[float, int]: """ Tìm khuôn mặt khớp nhất trong RAM cache bằng cosine similarity. Returns: (max_similarity, best_index) — best_index = -1 nếu không khớp """ if not self.known_embeddings: return 0.0, -1 similarities = [float(np.dot(embedding, e)) for e in self.known_embeddings] max_sim = max(similarities) best_idx = similarities.index(max_sim) if max_sim >= threshold: return max_sim, best_idx return max_sim, -1 # ── Singleton ───────────────────────────────────────────────────────────── face_analyzer = FaceAnalyzer()