""" Cross-modal retrieval with multiple strategies. Implements: 1. Multi-index search (separate indices per modality) 2. Modality-aware ranking 3. Hybrid search (combine same-modal and cross-modal) 4. Geo-filtered search (H3 spatial indexing) """ import torch import numpy as np import faiss from typing import Dict, List, Optional, Tuple from dataclasses import dataclass from pathlib import Path import json from .index import FAISSIndex from ..geo.spatial import SpatialIndex, GeoBox @dataclass class RetrievalResult: """Result from cross-modal retrieval.""" indices: List[int] scores: List[float] modalities: List[str] source_modality: str target_modality: str retrieval_type: str class CrossModalRetrieval: """ Multi-strategy cross-modal retrieval. Strategies: 1. SingleIndex: One FAISS index, filter by modality 2. MultiIndex: Separate indices per modality, search all 3. HybridSearch: Combine same-modal and cross-modal results """ def __init__(self, embed_dim: int = 768): self.embed_dim = embed_dim self.strategy = "single" # single, multi, hybrid self.use_modality_centering = True self.modality_means: Dict[str, np.ndarray] = {} # Single index strategy self.single_index = FAISSIndex(embed_dim) self.modality_labels: List[str] = [] self.sample_ids: List[int] = [] # Multi-index strategy self.indices: Dict[str, FAISSIndex] = {} self.modality_offsets: Dict[str, int] = {} # Metadata self.metadata: List[dict] = [] # Spatial index for geo-filtering self.spatial_index = SpatialIndex(resolution=7) def _get_search_query(self, query: np.ndarray, query_modality: str) -> np.ndarray: """Center the query embedding if modality centering is enabled.""" if getattr(self, "use_modality_centering", True) and query_modality in self.modality_means: mean = self.modality_means[query_modality] if query.ndim == 1: centered = query - mean.squeeze() norm = np.linalg.norm(centered) return centered / (norm + 1e-8) else: centered = query - mean.reshape(1, -1) norms = np.linalg.norm(centered, axis=1, keepdims=True) return centered / (norms + 1e-8) return query def build_single_index( self, embeddings: np.ndarray, modalities: List[str], metadata: List[dict], use_centering: bool = True ): """Build single FAISS index with all modalities.""" self.use_modality_centering = use_centering self.metadata = metadata self.modality_labels = modalities if self.use_modality_centering: # Compute means for each modality self.modality_means = {} for mod in set(modalities): mask = [m == mod for m in modalities] self.modality_means[mod] = np.mean(embeddings[mask], axis=0) # Center embeddings centered_embs = np.zeros_like(embeddings) for i, mod in enumerate(modalities): centered_embs[i] = embeddings[i] - self.modality_means[mod] # Normalize norms = np.linalg.norm(centered_embs, axis=1, keepdims=True) centered_embs = centered_embs / (norms + 1e-8) self.single_index.build(centered_embs) else: self.single_index.build(embeddings) self.strategy = "single" def build_multi_index( self, embeddings_by_modality: Dict[str, np.ndarray], metadata_by_modality: Dict[str, List[dict]], use_centering: bool = True ): """Build separate indices per modality.""" self.use_modality_centering = use_centering offset = 0 all_metadata = [] all_modalities = [] self.modality_means = {} for mod, embeddings in embeddings_by_modality.items(): # Compute mean self.modality_means[mod] = np.mean(embeddings, axis=0) # Center if enabled if self.use_modality_centering: centered = embeddings - self.modality_means[mod] norms = np.linalg.norm(centered, axis=1, keepdims=True) centered = centered / (norms + 1e-8) build_embs = centered else: build_embs = embeddings idx = FAISSIndex(self.embed_dim) idx.build(build_embs) self.indices[mod] = idx self.modality_offsets[mod] = offset offset += len(embeddings) all_metadata.extend(metadata_by_modality.get(mod, [])) all_modalities.extend([mod] * len(embeddings)) self.metadata = all_metadata self.modality_labels = all_modalities self.strategy = "multi" def build_spatial_index(self, metadata: List[dict]): """Build H3 spatial index from metadata with lat/lon (synthesizing if missing).""" import random for entry in metadata: if "lat" not in entry or "lon" not in entry: # Seed with index for reproducibility random.seed(entry["index"]) # Cluster synthesized coordinates closer to default India center to guarantee high hit rates entry["lat"] = 20.5937 + random.uniform(-1.2, 1.2) entry["lon"] = 78.9629 + random.uniform(-1.2, 1.2) self.spatial_index.add_image(entry["index"], entry["lat"], entry["lon"]) def search_geo( self, query: np.ndarray, query_modality: str, lat: float, lon: float, radius_km: float = 50.0, target_modality: Optional[str] = None, k: int = 5 ) -> RetrievalResult: """Search with geo-filtering: find images near a location.""" candidate_indices = self.spatial_index.query_radius(lat, lon, radius_km) if not candidate_indices: return RetrievalResult( indices=[], scores=[], modalities=[], source_modality=query_modality, target_modality=target_modality or "any", retrieval_type="geo" ) # Center the query centered_query = self._get_search_query(query, query_modality) # Convert to 1D flat array for dot product q_vec = centered_query.flatten() all_scores = [] all_indices = [] all_modalities = [] # We calculate cosine similarities directly for candidates to bypass FAISS dropout filtering if self.strategy == "multi": target_mods = [target_modality] if target_modality else list(self.indices.keys()) for t_mod in target_mods: if t_mod not in self.indices: continue offset = self.modality_offsets[t_mod] faiss_idx = self.indices[t_mod].index # Check candidates belonging to this modality for global_idx in candidate_indices: local_idx = global_idx - offset if 0 <= local_idx < faiss_idx.ntotal: try: # Reconstruct vector directly from FAISS index vec = faiss_idx.reconstruct(local_idx) score = float(np.dot(q_vec, vec.flatten())) all_scores.append(score) all_indices.append(global_idx) all_modalities.append(t_mod) except Exception: # Fallback score if reconstruction fails all_scores.append(0.0) all_indices.append(global_idx) all_modalities.append(t_mod) else: # Fallback for single index strategy for global_idx in candidate_indices: try: vec = self.single_index.index.reconstruct(global_idx) score = float(np.dot(q_vec, vec.flatten())) all_scores.append(score) all_indices.append(global_idx) all_modalities.append(self.modality_labels[global_idx]) except Exception: pass # If we have no valid scored candidates, return empty if not all_scores: return RetrievalResult( indices=[], scores=[], modalities=[], source_modality=query_modality, target_modality=target_modality or "any", retrieval_type="geo" ) # Sort candidates in descending order of similarity sorted_idx = np.argsort(all_scores)[::-1][:k] return RetrievalResult( indices=[all_indices[i] for i in sorted_idx], scores=[all_scores[i] for i in sorted_idx], modalities=[all_modalities[i] for i in sorted_idx], source_modality=query_modality, target_modality=target_modality or "any", retrieval_type="geo" ) def search_single( self, query: np.ndarray, query_modality: str = "optical", target_modality: Optional[str] = None, k: int = 5 ) -> RetrievalResult: """Search using single index with modality filtering.""" # Center the query centered_query = self._get_search_query(query, query_modality) # Get more results to filter search_k = min(k * 10, self.single_index.size) scores, indices = self.single_index.search(centered_query, k=search_k) # Filter by modality filtered_indices = [] filtered_scores = [] filtered_modalities = [] for idx, score in zip(indices[0], scores[0]): if idx < 0: continue mod = self.modality_labels[idx] if target_modality is None or mod == target_modality: filtered_indices.append(idx) filtered_scores.append(float(score)) filtered_modalities.append(mod) if len(filtered_indices) >= k: break return RetrievalResult( indices=filtered_indices, scores=filtered_scores, modalities=filtered_modalities, source_modality=query_modality, target_modality=target_modality or "any", retrieval_type="single" ) def search_multi( self, query: np.ndarray, query_modality: str, target_modalities: Optional[List[str]] = None, k: int = 5 ) -> RetrievalResult: """Search using multi-index strategy.""" if target_modalities is None: target_modalities = [m for m in self.indices.keys() if m != query_modality] # Center the query centered_query = self._get_search_query(query, query_modality) all_scores = [] all_indices = [] all_modalities = [] for mod in target_modalities: if mod not in self.indices: continue # Search this modality's index scores, indices = self.indices[mod].search(centered_query, k=k) # Offset indices to global space offset = self.modality_offsets[mod] global_indices = indices[0] + offset all_scores.extend(scores[0]) all_indices.extend(global_indices) all_modalities.extend([mod] * len(indices[0])) # Sort by score sorted_idx = np.argsort(all_scores)[::-1][:k] return RetrievalResult( indices=[all_indices[i] for i in sorted_idx], scores=[all_scores[i] for i in sorted_idx], modalities=[all_modalities[i] for i in sorted_idx], source_modality=query_modality, target_modality=",".join(target_modalities), retrieval_type="multi" ) def search_hybrid( self, query: np.ndarray, query_modality: str, k: int = 5, same_modal_weight: float = 0.7, cross_modal_weight: float = 0.3 ) -> RetrievalResult: """ Hybrid search combining same-modal and cross-modal results. Weighted combination of: 1. Same-modal results (higher weight) 2. Cross-modal results (lower weight) """ # Same-modal search same_modal_result = self.search_multi( query, query_modality, [query_modality], k=k ) # Cross-modal search cross_modal_targets = [m for m in self.indices.keys() if m != query_modality] cross_modal_result = self.search_multi( query, query_modality, cross_modal_targets, k=k ) # Combine with weights combined_scores = [] combined_indices = [] combined_modalities = [] for i in range(k): if i < len(same_modal_result.scores): combined_scores.append(same_modal_weight * same_modal_result.scores[i]) combined_indices.append(same_modal_result.indices[i]) combined_modalities.append(same_modal_result.modalities[i]) if i < len(cross_modal_result.scores): combined_scores.append(cross_modal_weight * cross_modal_result.scores[i]) combined_indices.append(cross_modal_result.indices[i]) combined_modalities.append(cross_modal_result.modalities[i]) # Sort combined results sorted_idx = np.argsort(combined_scores)[::-1][:k] return RetrievalResult( indices=[combined_indices[i] for i in sorted_idx], scores=[combined_scores[i] for i in sorted_idx], modalities=[combined_modalities[i] for i in sorted_idx], source_modality=query_modality, target_modality="hybrid", retrieval_type="hybrid" ) def search( self, query: np.ndarray, query_modality: str, target_modality: Optional[str] = None, k: int = 5, strategy: Optional[str] = None, lat: Optional[float] = None, lon: Optional[float] = None, radius_km: float = 50.0 ) -> RetrievalResult: """ Unified search interface. Args: query: Query embedding query_modality: Modality of query image target_modality: Target modality (None for all) k: Number of results strategy: Override strategy (single, multi, hybrid) lat: Latitude for geo-filtering (optional) lon: Longitude for geo-filtering (optional) radius_km: Search radius in km (default 50) Returns: RetrievalResult with ranked results """ if lat is not None and lon is not None: return self.search_geo(query, query_modality, lat, lon, radius_km, target_modality, k) strategy = strategy or self.strategy if strategy == "single": return self.search_single(query, query_modality, target_modality, k) elif strategy == "multi": targets = [target_modality] if target_modality else None return self.search_multi(query, query_modality, targets, k) elif strategy == "hybrid": return self.search_hybrid(query, query_modality, k) else: raise ValueError(f"Unknown strategy: {strategy}") def save(self, path: Path): """Save indices and metadata.""" path.mkdir(parents=True, exist_ok=True) if self.strategy == "single": self.single_index.save(str(path / "single_index.faiss")) elif self.strategy == "multi": for mod, idx in self.indices.items(): idx.save(str(path / f"{mod}_index.faiss")) # Save metadata with open(path / "metadata.json", "w") as f: serialized_means = {k: v.tolist() for k, v in self.modality_means.items()} json.dump({ "modality_labels": self.modality_labels, "metadata": self.metadata, "strategy": self.strategy, "use_modality_centering": self.use_modality_centering, "modality_means": serialized_means, }, f) def load(self, path: Path): """Load indices and metadata.""" # Load metadata with open(path / "metadata.json") as f: data = json.load(f) self.modality_labels = data["modality_labels"] self.metadata = data["metadata"] self.strategy = data["strategy"] self.use_modality_centering = data.get("use_modality_centering", True) # Restore means self.modality_means = {} for k, v in data.get("modality_means", {}).items(): self.modality_means[k] = np.array(v).astype(np.float32) if self.strategy == "single": self.single_index.load(str(path / "single_index.faiss")) elif self.strategy == "multi": for mod in set(self.modality_labels): idx_path = path / f"{mod}_index.faiss" if idx_path.exists(): idx = FAISSIndex(self.embed_dim) idx.load(str(idx_path)) self.indices[mod] = idx # Self-check if __name__ == "__main__": print("Testing Cross-Modal Retrieval...") # Create test data n_per_mod = 100 embed_dim = 768 embeddings_by_modality = { "optical": np.random.randn(n_per_mod, embed_dim).astype(np.float32), "sar": np.random.randn(n_per_mod, embed_dim).astype(np.float32), "multispectral": np.random.randn(n_per_mod, embed_dim).astype(np.float32), } # Normalize for mod in embeddings_by_modality: norms = np.linalg.norm(embeddings_by_modality[mod], axis=1, keepdims=True) embeddings_by_modality[mod] = embeddings_by_modality[mod] / norms # Create metadata metadata_by_modality = { mod: [{"index": i, "modality": mod, "class": f"class_{i % 10}"} for i in range(n_per_mod)] for mod in embeddings_by_modality } # Test multi-index strategy retrieval = CrossModalRetrieval(embed_dim) retrieval.build_multi_index(embeddings_by_modality, metadata_by_modality) print(f"Built multi-index with modalities: {list(retrieval.indices.keys())}") # Test search query = np.random.randn(1, embed_dim).astype(np.float32) query = query / np.linalg.norm(query) result = retrieval.search(query, "optical", k=5) print(f"\nSearch results:") print(f" Indices: {result.indices}") print(f" Scores: {result.scores}") print(f" Modalities: {result.modalities}") # Test hybrid search result = retrieval.search_hybrid(query, "optical", k=5) print(f"\nHybrid search results:") print(f" Indices: {result.indices}") print(f" Modalities: {result.modalities}") print("\nCross-Modal Retrieval test passed!")