""" storage/mongo_storage.py — MongoDB Atlas: quản lý 4 collections theo ERD Collections: raw_images — Ảnh gốc từ camera / upload face_crops — Ảnh khuôn mặt đã crop (có đầy đủ metadata) unknown_faces — Khuôn mặt không nhận dạng được recognition_payloads — Payload debug cho mỗi lần nhận diện """ import logging import time from datetime import datetime, timezone from bson.objectid import ObjectId from pymongo import MongoClient import config logger = logging.getLogger("mongo_storage") def _utcnow() -> datetime: return datetime.now(timezone.utc) class MongoStorage: """Singleton wrapper cho tất cả thao tác MongoDB.""" def __init__(self): self._client: MongoClient | None = None def _get_db(self): if self._client is None: self._client = MongoClient(config.MONGO_URI) return self._client[config.MONGO_DB_NAME] # ── raw_images ──────────────────────────────────────────────────────── def save_raw_image( self, img_bytes: bytes, source_type: str = "upload", width: int = 0, height: int = 0, fmt: str = "jpeg", ) -> str: """ Lưu ảnh gốc vào collection raw_images. Returns: str: ObjectId string của document vừa lưu (rỗng nếu lỗi) """ try: col = self._get_db()[config.MONGO_COL_RAW_IMAGES] doc = { "image_bytes": img_bytes, "source_type": source_type, # "webcam" | "upload" "camera_id": None, "width": width, "height": height, "format": fmt, "metadata": {}, "created_at": _utcnow(), } result = col.insert_one(doc) logger.debug(f"[Mongo] Saved raw_image: {result.inserted_id}") return str(result.inserted_id) except Exception as e: logger.error(f"[Mongo] save_raw_image failed: {e}") return "" def get_raw_image(self, raw_image_id: str) -> bytes | None: """Lấy bytes ảnh gốc theo ID.""" try: col = self._get_db()[config.MONGO_COL_RAW_IMAGES] doc = col.find_one({"_id": ObjectId(raw_image_id)}) return doc.get("image_bytes") if doc else None except Exception as e: logger.error(f"[Mongo] get_raw_image failed: {e}") return None # ── face_crops ──────────────────────────────────────────────────────── def save_face_crop( self, img_bytes: bytes, raw_image_id: str = "", person_id: str | None = None, visit_id: str | None = None, recognition_log_id: str | None = None, bbox: dict | None = None, image_type: str = "recognition", confidence: float = 0.0, quality_score: float = 0.0, model_name: str = "buffalo_l", model_version: str = "1.0.0", ) -> str: """ Lưu ảnh khuôn mặt crop vào collection face_crops (theo ERD). Returns: str: ObjectId string (rỗng nếu lỗi) """ try: col = self._get_db()[config.MONGO_COL_FACE_CROPS] doc = { # Foreign keys "person_id": person_id, "visit_id": visit_id, "recognition_log_id": recognition_log_id, "raw_image_id": raw_image_id, # Image data "image_bytes": img_bytes, "image_url": None, "image_base64": None, # Bounding box & detection info "bbox": bbox or {}, "image_type": image_type, # "recognition" | "registration" "confidence": confidence, "quality_score": quality_score, # Model info "model": { "name": model_name, "version": model_version, }, "created_at": _utcnow(), } result = col.insert_one(doc) logger.debug(f"[Mongo] Saved face_crop: {result.inserted_id}") return str(result.inserted_id) except Exception as e: logger.error(f"[Mongo] save_face_crop failed: {e}") return "" def update_face_crop_links( self, crop_id: str, person_id: str, visit_id: str, recognition_log_id: str, ) -> bool: """Cập nhật các liên kết sau khi nhận diện thành công.""" try: col = self._get_db()[config.MONGO_COL_FACE_CROPS] result = col.update_one( {"_id": ObjectId(crop_id)}, {"$set": { "person_id": person_id, "visit_id": visit_id, "recognition_log_id": recognition_log_id, }}, ) return result.modified_count > 0 except Exception as e: logger.error(f"[Mongo] update_face_crop_links failed: {e}") return False def get_face_crop_bytes(self, crop_id: str) -> bytes | None: """Lấy bytes ảnh crop theo ID.""" try: col = self._get_db()[config.MONGO_COL_FACE_CROPS] doc = col.find_one({"_id": ObjectId(crop_id)}) return doc.get("image_bytes") if doc else None except Exception as e: logger.error(f"[Mongo] get_face_crop_bytes failed: {e}") return None def delete_face_crop(self, crop_id: str) -> bool: """Xóa ảnh crop theo ID.""" try: col = self._get_db()[config.MONGO_COL_FACE_CROPS] result = col.delete_one({"_id": ObjectId(crop_id)}) return result.deleted_count > 0 except Exception as e: logger.error(f"[Mongo] delete_face_crop failed: {e}") return False def list_face_crops_by_person(self, person_id: str) -> list[str]: """Lấy danh sách crop_id của một person.""" try: col = self._get_db()[config.MONGO_COL_FACE_CROPS] docs = col.find({"person_id": person_id}, {"_id": 1}) return [str(d["_id"]) for d in docs] except Exception as e: logger.error(f"[Mongo] list_face_crops_by_person failed: {e}") return [] # ── unknown_faces ───────────────────────────────────────────────────── def save_unknown_face( self, face_crop_id: str, raw_image_id: str, recognition_log_id: str, reason: str, confidence: float, similarity_distance: float, ) -> str: """ Lưu thông tin khuôn mặt không nhận diện được vào unknown_faces. Returns: str: ObjectId string (rỗng nếu lỗi) """ try: col = self._get_db()[config.MONGO_COL_UNKNOWN_FACES] doc = { "face_crop_id": face_crop_id, "raw_image_id": raw_image_id, "recognition_log_id": recognition_log_id, "reason": reason, # "no_match" | "low_confidence" | "no_db" "confidence": confidence, "similarity_distance": similarity_distance, "is_reviewed": False, "review_result": None, "reviewed_by": None, "reviewed_at": None, "review_note": None, "created_at": _utcnow(), } result = col.insert_one(doc) logger.debug(f"[Mongo] Saved unknown_face: {result.inserted_id}") return str(result.inserted_id) except Exception as e: logger.error(f"[Mongo] save_unknown_face failed: {e}") return "" def list_unknown_faces(self, limit: int = 50) -> list[dict]: """Lấy danh sách unknown faces gần nhất.""" try: col = self._get_db()[config.MONGO_COL_UNKNOWN_FACES] docs = col.find({}, {"image_bytes": 0}).sort("created_at", -1).limit(limit) results = [] for d in docs: d["id"] = str(d.pop("_id")) if "created_at" in d: d["created_at"] = d["created_at"].isoformat() results.append(d) return results except Exception as e: logger.error(f"[Mongo] list_unknown_faces failed: {e}") return [] # ── recognition_payloads ────────────────────────────────────────────── def save_recognition_payload( self, recognition_log_id: str, raw_image_id: str, face_crop_id: str, model_name: str, model_version: str, request_payload: dict, response_payload: dict, runtime_ms: float, ) -> str: """ Lưu payload nhận diện đầy đủ vào recognition_payloads (để debug). Returns: str: ObjectId string (rỗng nếu lỗi) """ try: col = self._get_db()[config.MONGO_COL_PAYLOADS] doc = { "recognition_log_id": recognition_log_id, "raw_image_id": raw_image_id, "face_crop_id": face_crop_id, "model_name": model_name, "model_version": model_version, "request_payload": request_payload, "response_payload": response_payload, "runtime": {"total_ms": runtime_ms}, "created_at": _utcnow(), } result = col.insert_one(doc) logger.debug(f"[Mongo] Saved recognition_payload: {result.inserted_id}") return str(result.inserted_id) except Exception as e: logger.error(f"[Mongo] save_recognition_payload failed: {e}") return "" # ── Singleton ───────────────────────────────────────────────────────────── mongo_storage = MongoStorage()