""" FAISS index for fast similarity search. Handles index building, searching, and persistence. """ import faiss import torch import numpy as np from pathlib import Path from typing import Tuple, Optional class FAISSIndex: """ FAISS index for cosine similarity search. Uses IndexFlatIP (inner product) which works as cosine similarity when embeddings are L2-normalized. Supports: - Global embeddings (standard CLIP) - Multiscale embeddings (fused global + patch features) - Multiple index types for different use cases """ def __init__(self, embed_dim: int = 768, index_type: str = "flat"): """ Initialize FAISS index. Args: embed_dim: Embedding dimension index_type: Type of index ("flat", "ivf", "pq") """ self.embed_dim = embed_dim self.index_type = index_type self.is_built = False self._n_embeddings = 0 # Create index based on type if index_type == "flat": self.index = faiss.IndexFlatIP(embed_dim) elif index_type == "ivf": # IVF index for faster search on large datasets quantizer = faiss.IndexFlatIP(embed_dim) self.index = faiss.IndexIVFFlat(quantizer, embed_dim, 100) elif index_type == "pq": # Product quantization for memory efficiency self.index = faiss.IndexPQ(embed_dim, 8, 8) else: raise ValueError(f"Unknown index type: {index_type}") @property def size(self) -> int: """Number of embeddings in index.""" return self._n_embeddings def build(self, embeddings: torch.Tensor, train_index: bool = False) -> None: """ Build index from embeddings. Args: embeddings: Tensor of shape (N, embed_dim), L2-normalized train_index: Whether to train IVF/PQ index (requires enough data) """ if isinstance(embeddings, torch.Tensor): embeddings = embeddings.numpy().astype(np.float32) if embeddings.ndim == 1: embeddings = embeddings.reshape(1, -1) assert embeddings.shape[1] == self.embed_dim, \ f"Expected dim {self.embed_dim}, got {embeddings.shape[1]}" # Recreate index with correct type if self.index_type == "flat": self.index = faiss.IndexFlatIP(self.embed_dim) elif self.index_type == "ivf": quantizer = faiss.IndexFlatIP(self.embed_dim) self.index = faiss.IndexIVFFlat(quantizer, self.embed_dim, 100) if train_index and embeddings.shape[0] >= 100: self.index.train(embeddings) elif self.index_type == "pq": self.index = faiss.IndexPQ(self.embed_dim, 8, 8) if train_index and embeddings.shape[0] >= 100: self.index.train(embeddings) self.index.add(embeddings) self._n_embeddings = self.index.ntotal self.is_built = True def search( self, query: torch.Tensor, k: int = 5 ) -> Tuple[np.ndarray, np.ndarray]: """ Search for top-K similar embeddings. Args: query: Query embedding(s), shape (embed_dim,) or (N, embed_dim) k: Number of results to return Returns: (scores, indices) where: - scores: shape (N, k) with similarity scores - indices: shape (N, k) with indices into gallery """ if not self.is_built: raise RuntimeError("Index not built. Call build() first.") # Convert to numpy if isinstance(query, torch.Tensor): query = query.numpy().astype(np.float32) # Ensure 2D if query.ndim == 1: query = query.reshape(1, -1) # Clamp k to index size k = min(k, self._n_embeddings) # Search scores, indices = self.index.search(query, k) return scores, indices def save(self, path: str) -> None: """ Save index to disk. Args: path: Path to save index (without extension) """ path = Path(path) path.parent.mkdir(parents=True, exist_ok=True) # Save FAISS index faiss.write_index(self.index, str(path)) def load(self, path: str) -> None: """ Load index from disk. Args: path: Path to saved index """ self.index = faiss.read_index(str(path)) self._n_embeddings = self.index.ntotal self.is_built = True def get_embeddings(self) -> np.ndarray: """Get all embeddings from index.""" if not self.is_built: return np.array([]) embeddings = np.array([ self.index.reconstruct(i) for i in range(self._n_embeddings) ]) return embeddings def search_multiscale( self, query: torch.Tensor, k: int = 5, global_weight: float = 0.7 ) -> Tuple[np.ndarray, np.ndarray]: """ Search with weighted global + patch features. Args: query: Query embedding (fused global + patch) k: Number of results global_weight: Weight for global features (0-1) Returns: (scores, indices) """ if not self.is_built: raise RuntimeError("Index not built. Call build() first.") if isinstance(query, torch.Tensor): query = query.numpy().astype(np.float32) if query.ndim == 1: query = query.reshape(1, -1) k = min(k, self._n_embeddings) scores, indices = self.index.search(query, k) return scores, indices # Self-check if __name__ == "__main__": print("Testing FAISSIndex...") # Create dummy embeddings n_gallery = 100 embed_dim = 768 embeddings = torch.randn(n_gallery, embed_dim) embeddings = torch.nn.functional.normalize(embeddings, dim=1) # Build index index = FAISSIndex(embed_dim) index.build(embeddings) print(f"Index built with {index.size} embeddings") # Search query = torch.randn(embed_dim) query = torch.nn.functional.normalize(query, dim=0) scores, indices = index.search(query, k=5) print(f"Query results:") print(f" Scores shape: {scores.shape}") print(f" Indices shape: {indices.shape}") print(f" Top-5 scores: {scores[0]}") print(f" Top-5 indices: {indices[0]}") # Save/load roundtrip import tempfile with tempfile.TemporaryDirectory() as tmpdir: save_path = Path(tmpdir) / "test_index.faiss" index.save(save_path) loaded_index = FAISSIndex(embed_dim) loaded_index.load(save_path) print(f"\nLoaded index size: {loaded_index.size}") # Verify search results match scores2, indices2 = loaded_index.search(query, k=5) assert np.allclose(scores, scores2), "Scores mismatch!" assert np.array_equal(indices, indices2), "Indices mismatch!" print("\nFAISSIndex test passed!")