""" Database & Index Manager for VideoDetection. Quản lý vector index (FAISS / Numpy) và catalog video gốc của chủ sở hữu trong thư mục Data/. """ from __future__ import annotations import json import logging import pickle from dataclasses import asdict, dataclass from pathlib import Path from typing import Any, Dict, List, Optional, Tuple import numpy as np from ..core.paths import DATA_DIR logger = logging.getLogger(__name__) @dataclass class VideoRecord: video_id: str filename: str rel_path: str duration_sec: float visual_clips_count: int audio_chunks_count: int indexed_at: str class VectorIndexWrapper: """ Bọc FAISS Index với fallback Numpy Cosine Similarity tự động nếu môi trường chưa cài faiss-cpu. """ def __init__(self, dim: int, use_faiss: bool = True) -> None: self.dim = dim self.use_faiss = use_faiss self.index = None self.vectors: List[np.ndarray] = [] if self.use_faiss: try: import faiss # Dùng IndexFlatIP (Inner Product) cho vector L2-normalized = Cosine Similarity self.index = faiss.IndexFlatIP(self.dim) except ImportError: logger.warning("faiss-cpu chưa sẵn sàng, chuyển sang chế độ Numpy Vector Search.") self.index = None def add(self, vectors: np.ndarray) -> None: """vectors: shape (N, dim) float32 L2-normalized.""" if len(vectors) == 0: return vectors = vectors.astype(np.float32) if self.index is not None: self.index.add(vectors) else: self.vectors.append(vectors) def search(self, queries: np.ndarray, top_k: int = 5) -> Tuple[np.ndarray, np.ndarray]: """ Tìm kiếm vector tương đồng nhất. Trả về: (distances [N, top_k], indices [N, top_k]) """ if len(queries) == 0: return np.empty((0, top_k), dtype=np.float32), np.empty((0, top_k), dtype=np.int64) queries = queries.astype(np.float32) if self.index is not None and self.index.ntotal > 0: actual_k = min(top_k, self.index.ntotal) distances, indices = self.index.search(queries, actual_k) return distances, indices # Fallback Numpy if not self.vectors: return np.empty((len(queries), 0), dtype=np.float32), np.empty((len(queries), 0), dtype=np.int64) all_vecs = np.vstack(self.vectors) if len(all_vecs) == 0: return np.empty((len(queries), 0), dtype=np.float32), np.empty((len(queries), 0), dtype=np.int64) # Dot product (Cosine sim) sim_matrix = np.dot(queries, all_vecs.T) # [Q, N] actual_k = min(top_k, all_vecs.shape[0]) # Lấy top k lớn nhất indices = np.argsort(-sim_matrix, axis=1)[:, :actual_k] distances = np.take_along_axis(sim_matrix, indices, axis=1) return distances.astype(np.float32), indices.astype(np.int64) @property def total(self) -> int: if self.index is not None: return self.index.ntotal if self.vectors: return sum(len(v) for v in self.vectors) return 0 class OwnerDatabase: """ Quản lý kho dữ liệu so sánh trong Data/. Bao gồm catalog, visual vector index và audio fingerprint index. """ def __init__(self, data_dir: Path = DATA_DIR) -> None: self.data_dir = data_dir self.catalog_file = self.data_dir / "catalog.json" self.visual_index_file = self.data_dir / "visual.index" self.visual_meta_file = self.data_dir / "visual_meta.json" self.audio_index_file = self.data_dir / "audio.index" self.audio_meta_file = self.data_dir / "audio_meta.json" self.catalog: Dict[str, VideoRecord] = {} # metadata maps: index_id -> {"video_id": ..., "timestamp": ...} self.visual_meta: List[Dict[str, Any]] = [] self.audio_meta: List[Dict[str, Any]] = [] self.visual_index = VectorIndexWrapper(dim=1024) self.audio_index = VectorIndexWrapper(dim=128) self.load() def add_reference( self, record: VideoRecord, visual_features: List[Tuple[float, np.ndarray]], # [(timestamp, emb_1024)] audio_features: List[Tuple[float, np.ndarray]], # [(timestamp, fp_128)] ) -> None: """Đăng ký một video vào kho bản quyền chủ sở hữu.""" self.catalog[record.video_id] = record # Thêm visual if visual_features: vis_vecs = np.stack([feat[1] for feat in visual_features], axis=0) self.visual_index.add(vis_vecs) for feat in visual_features: self.visual_meta.append({ "video_id": record.video_id, "timestamp": float(feat[0]), }) # Thêm audio if audio_features: aud_vecs = np.stack([feat[1] for feat in audio_features], axis=0) self.audio_index.add(aud_vecs) for feat in audio_features: self.audio_meta.append({ "video_id": record.video_id, "timestamp": float(feat[0]), }) def save(self) -> None: """Lưu toàn bộ index và metadata xuống thư mục Data/.""" self.data_dir.mkdir(parents=True, exist_ok=True) # 1. Lưu catalog catalog_dict = {vid: asdict(rec) for vid, rec in self.catalog.items()} with open(self.catalog_file, "w", encoding="utf-8") as f: json.dump(catalog_dict, f, indent=2, ensure_ascii=False) # 2. Lưu metadata with open(self.visual_meta_file, "w", encoding="utf-8") as f: json.dump(self.visual_meta, f, indent=2) with open(self.audio_meta_file, "w", encoding="utf-8") as f: json.dump(self.audio_meta, f, indent=2) # 3. Lưu index FAISS hoặc Numpy try: import faiss if self.visual_index.index is not None: faiss.write_index(self.visual_index.index, str(self.visual_index_file)) if self.audio_index.index is not None: faiss.write_index(self.audio_index.index, str(self.audio_index_file)) except Exception: # Fallback lưu pickle nếu là numpy with open(self.data_dir / "visual_numpy.pkl", "wb") as f: pickle.dump(self.visual_index.vectors, f) with open(self.data_dir / "audio_numpy.pkl", "wb") as f: pickle.dump(self.audio_index.vectors, f) def load(self) -> None: """Nạp dữ liệu từ Data/ nếu đã tồn tại.""" if self.catalog_file.exists(): try: with open(self.catalog_file, "r", encoding="utf-8") as f: data = json.load(f) self.catalog = {vid: VideoRecord(**rec) for vid, rec in data.items()} except Exception as e: logger.error(f"Lỗi đọc catalog.json: {e}") if self.visual_meta_file.exists(): try: with open(self.visual_meta_file, "r", encoding="utf-8") as f: self.visual_meta = json.load(f) except Exception as e: logger.error(f"Lỗi đọc visual_meta.json: {e}") if self.audio_meta_file.exists(): try: with open(self.audio_meta_file, "r", encoding="utf-8") as f: self.audio_meta = json.load(f) except Exception as e: logger.error(f"Lỗi đọc audio_meta.json: {e}") # Nạp index try: import faiss if self.visual_index_file.exists(): self.visual_index.index = faiss.read_index(str(self.visual_index_file)) if self.audio_index_file.exists(): self.audio_index.index = faiss.read_index(str(self.audio_index_file)) except Exception: pkl_vis = self.data_dir / "visual_numpy.pkl" if pkl_vis.exists(): with open(pkl_vis, "rb") as f: self.visual_index.vectors = pickle.load(f) pkl_aud = self.data_dir / "audio_numpy.pkl" if pkl_aud.exists(): with open(pkl_aud, "rb") as f: self.audio_index.vectors = pickle.load(f)