""" FAISS Vector Database for Multi-Modal Profile Search ====================================================== Maintains separate FAISS indices for each biometric modality: - gait_index: 256-d (gait embeddings) - biomech_index: 64-d (biomechanical features) - appearance_index: 512-d (OSNet appearance embeddings) - face_index: 512-d (ArcFace embeddings, optional) Uses IndexFlatIP (inner product) on L2-normalized vectors = cosine similarity. """ import os import json import numpy as np # Try FAISS GPU, then CPU, then fallback to brute force FAISS_AVAILABLE = False FAISS_GPU = False try: import faiss FAISS_AVAILABLE = True # Check for GPU support try: if faiss.get_num_gpus() > 0: FAISS_GPU = True print("[VectorDB] FAISS-GPU available") else: print("[VectorDB] FAISS-CPU available") except Exception: print("[VectorDB] FAISS-CPU available") except ImportError: print("[VectorDB] FAISS not found — using brute-force numpy fallback") class _NumpyFallbackIndex: """Brute-force numpy fallback when FAISS is not available.""" def __init__(self, dim): self.dim = dim self.vectors = np.zeros((0, dim), dtype=np.float32) @property def ntotal(self): return self.vectors.shape[0] def add(self, vectors): self.vectors = np.vstack([self.vectors, vectors.astype(np.float32)]) def search(self, query, k): if self.ntotal == 0: return np.array([[-1.0] * k]), np.array([[-1] * k]) k = min(k, self.ntotal) # Inner product (= cosine sim on normalized vectors) sims = np.dot(self.vectors, query.T).flatten() top_k = np.argsort(sims)[::-1][:k] scores = sims[top_k].reshape(1, -1) indices = top_k.reshape(1, -1) return scores, indices def reset(self): self.vectors = np.zeros((0, self.dim), dtype=np.float32) # Modality definitions MODALITIES = { "gait": {"dim": 256, "weight": 0.35}, "biomech": {"dim": 64, "weight": 0.25}, "appearance": {"dim": 512, "weight": 0.25}, "height": {"dim": 1, "weight": 0.15}, "face": {"dim": 512, "weight": 0.30}, # optional, applied differently } class ProfileVectorDB: """ Multi-modal FAISS vector database for person profiles. Each profile is stored across multiple indices (one per modality). All vectors must be L2-normalized before insertion. """ def __init__(self, db_dir="saved_profiles"): self.db_dir = db_dir os.makedirs(db_dir, exist_ok=True) self.indices = {} # modality_name -> FAISS index self.profile_ids = [] # ordered list of profile IDs (position = FAISS index position) self._id_file = os.path.join(db_dir, "faiss_profile_ids.json") # Build indices for mod_name, mod_info in MODALITIES.items(): self.indices[mod_name] = self._create_index(mod_info["dim"]) # Load existing self._load() def _create_index(self, dim): """Create a FAISS IndexFlatIP (inner product = cosine on normalized vecs).""" if FAISS_AVAILABLE: index = faiss.IndexFlatIP(dim) return index else: return _NumpyFallbackIndex(dim) def _load(self): """Load saved indices and profile ID mapping.""" if os.path.exists(self._id_file): try: with open(self._id_file, "r") as f: self.profile_ids = json.load(f) except Exception: self.profile_ids = [] if FAISS_AVAILABLE: for mod_name in MODALITIES: idx_path = os.path.join(self.db_dir, f"faiss_{mod_name}.index") if os.path.exists(idx_path): try: self.indices[mod_name] = faiss.read_index(idx_path) except Exception: self.indices[mod_name] = self._create_index(MODALITIES[mod_name]["dim"]) def save(self): """Persist indices and ID mapping to disk.""" with open(self._id_file, "w") as f: json.dump(self.profile_ids, f) if FAISS_AVAILABLE: for mod_name in MODALITIES: idx_path = os.path.join(self.db_dir, f"faiss_{mod_name}.index") try: faiss.write_index(self.indices[mod_name], idx_path) except Exception as e: print(f"[VectorDB] Failed to save {mod_name} index: {e}") def add_profile(self, profile_id, vectors): """ Add a profile's vectors to all indices. Args: profile_id: str vectors: dict with keys matching MODALITIES, values are np.array or None """ self.profile_ids.append(profile_id) for mod_name, mod_info in MODALITIES.items(): vec = vectors.get(mod_name) if vec is not None: vec = np.array(vec, dtype=np.float32).reshape(1, -1) # L2 normalize norm = np.linalg.norm(vec) if norm > 0: vec = vec / norm self.indices[mod_name].add(vec) else: # Add zero vector as placeholder to keep indices aligned zero = np.zeros((1, mod_info["dim"]), dtype=np.float32) self.indices[mod_name].add(zero) self.save() def remove_profile(self, profile_id): """Remove a profile. Requires rebuilding indices (FAISS doesn't support deletion).""" if profile_id not in self.profile_ids: return idx = self.profile_ids.index(profile_id) self.profile_ids.pop(idx) # Rebuild indices without the removed profile # This is expensive but deletion is rare for mod_name, mod_info in MODALITIES.items(): old_index = self.indices[mod_name] dim = mod_info["dim"] if old_index.ntotal <= 1: self.indices[mod_name] = self._create_index(dim) continue if FAISS_AVAILABLE: # Extract all vectors, remove the one at idx, rebuild all_vecs = faiss.rev_swig_ptr( old_index.get_xb(), old_index.ntotal * dim ).reshape(old_index.ntotal, dim).copy() all_vecs = np.delete(all_vecs, idx, axis=0) new_index = self._create_index(dim) if len(all_vecs) > 0: new_index.add(all_vecs) self.indices[mod_name] = new_index else: # Numpy fallback all_vecs = old_index.vectors.copy() all_vecs = np.delete(all_vecs, idx, axis=0) new_index = _NumpyFallbackIndex(dim) if len(all_vecs) > 0: new_index.add(all_vecs) self.indices[mod_name] = new_index self.save() def search(self, query_vectors, k=5): """ Search for closest profiles using multi-modal weighted ensemble. Args: query_vectors: dict with modality names as keys, np.array as values e.g., {"gait": vec_256, "appearance": vec_512, ...} k: number of results to return Returns: list of (profile_id, weighted_score, per_modality_scores) sorted by score desc. """ if len(self.profile_ids) == 0: return [] k = min(k, len(self.profile_ids)) profile_scores = {} # profile_id -> {mod_name: score} for mod_name, mod_info in MODALITIES.items(): query = query_vectors.get(mod_name) if query is None: continue if np.all(query == 0): continue query = np.array(query, dtype=np.float32).reshape(1, -1) norm = np.linalg.norm(query) if norm > 0: query = query / norm index = self.indices[mod_name] if index.ntotal == 0: continue scores, indices = index.search(query, min(k, index.ntotal)) for i in range(scores.shape[1]): idx = int(indices[0][i]) if 0 <= idx < len(self.profile_ids): pid = self.profile_ids[idx] if pid not in profile_scores: profile_scores[pid] = {} profile_scores[pid][mod_name] = max(0.0, float(scores[0][i])) # Compute weighted ensemble score per profile results = [] for pid, mod_scores in profile_scores.items(): # Standard 4-modal ensemble gait_sim = mod_scores.get("gait", 0.0) biomech_sim = mod_scores.get("biomech", 0.0) appearance_sim = mod_scores.get("appearance", 0.0) height_sim = mod_scores.get("height", 0.5) weighted = (0.35 * gait_sim + 0.25 * biomech_sim + 0.25 * appearance_sim + 0.15 * height_sim) # Optional face modality face_sim = mod_scores.get("face", None) if face_sim is not None and face_sim > 0: weighted = 0.70 * weighted + 0.30 * face_sim results.append((pid, weighted, mod_scores)) results.sort(key=lambda x: x[1], reverse=True) return results[:k] def get_profile_count(self): return len(self.profile_ids)