""" Multi-modal retrieval for satellite imagery. Handles same-modal and cross-modal retrieval with modality filtering. """ import torch import numpy as np from typing import Dict, List, Optional, Tuple from dataclasses import dataclass from .index import FAISSIndex @dataclass class ModalityResult: """Result with modality information.""" indices: List[int] scores: List[float] modalities: List[str] query_modality: str class MultiModalRetrieval: """ Multi-modal retrieval with modality-aware search. Supports same-modal and cross-modal queries with filtering. """ # Modality to index mapping MODALITY_MAP = { "optical": 0, "sar": 1, "multispectral": 2, } def __init__(self, embed_dim: int = 768): """ Initialize multi-modal retrieval. Args: embed_dim: Embedding dimension """ self.embed_dim = embed_dim self.index = FAISSIndex(embed_dim) # Track modality for each embedding self.modality_labels: List[str] = [] self.sample_ids: List[int] = [] @property def size(self) -> int: """Total number of embeddings.""" return self.index.size def build_index( self, embeddings_by_modality: Dict[str, torch.Tensor], sample_ids_by_modality: Optional[Dict[str, List[int]]] = None ) -> None: """ Build index with modality labels. Args: embeddings_by_modality: Dict mapping modality to embeddings tensor sample_ids_by_modality: Optional sample IDs per modality """ all_embeddings = [] all_modalities = [] all_sample_ids = [] for modality, embeddings in embeddings_by_modality.items(): # Convert to numpy if needed if isinstance(embeddings, torch.Tensor): embeddings = embeddings.numpy().astype(np.float32) all_embeddings.append(embeddings) all_modalities.extend([modality] * len(embeddings)) # Sample IDs if sample_ids_by_modality and modality in sample_ids_by_modality: all_sample_ids.extend(sample_ids_by_modality[modality]) else: all_sample_ids.extend(range(len(embeddings))) # Concatenate all embeddings combined_embeddings = np.concatenate(all_embeddings, axis=0) # Build index self.index.build(combined_embeddings) self.modality_labels = all_modalities self.sample_ids = all_sample_ids def _filter_by_modality( self, indices: np.ndarray, scores: np.ndarray, target_modality: Optional[str] = None ) -> Tuple[List[int], List[float], List[str]]: """ Filter results by modality. Args: indices: Raw indices from FAISS scores: Raw scores from FAISS target_modality: If specified, only return results from this modality Returns: (filtered_indices, filtered_scores, modalities) """ filtered_indices = [] filtered_scores = [] filtered_modalities = [] for idx, score in zip(indices[0], scores[0]): if idx < 0: # FAISS returns -1 for empty slots continue modality = self.modality_labels[idx] if target_modality is None or modality == target_modality: filtered_indices.append(idx) filtered_scores.append(float(score)) filtered_modalities.append(modality) return filtered_indices, filtered_scores, filtered_modalities def same_modal_query( self, query_embedding: torch.Tensor, modality: str, k: int = 5 ) -> ModalityResult: """ Query for same modality. Args: query_embedding: Query embedding modality: Modality to search k: Number of results Returns: ModalityResult with filtered results """ # Search with no filter first scores, indices = self.index.search(query_embedding, k=k * 10) # Get more to filter # Filter by modality filtered_indices, filtered_scores, modalities = self._filter_by_modality( indices, scores, target_modality=modality ) # Take top-k return ModalityResult( indices=filtered_indices[:k], scores=filtered_scores[:k], modalities=modalities[:k], query_modality=modality ) def cross_modal_query( self, query_embedding: torch.Tensor, source_modality: str, target_modality: str, k: int = 5 ) -> ModalityResult: """ Query across modalities. Args: query_embedding: Query embedding source_modality: Modality of query image target_modality: Modality to search in k: Number of results Returns: ModalityResult with filtered results """ # Search with no filter first scores, indices = self.index.search(query_embedding, k=k * 10) # Filter by target modality (excluding source) filtered_indices, filtered_scores, modalities = self._filter_by_modality( indices, scores, target_modality=target_modality ) # Take top-k return ModalityResult( indices=filtered_indices[:k], scores=filtered_scores[:k], modalities=modalities[:k], query_modality=source_modality ) def mixed_query( self, query_embedding: torch.Tensor, source_modality: str, k: int = 5 ) -> ModalityResult: """ Query across all modalities. Args: query_embedding: Query embedding source_modality: Modality of query image k: Number of results Returns: ModalityResult with results from all modalities """ # Search scores, indices = self.index.search(query_embedding, k=k) # Get modalities modalities = [ self.modality_labels[idx] for idx in indices[0] if idx >= 0 ] return ModalityResult( indices=indices[0].tolist(), scores=scores[0].tolist(), modalities=modalities, query_modality=source_modality ) def get_modality_distribution(self) -> Dict[str, int]: """ Get distribution of modalities in index. Returns: Dict mapping modality to count """ dist = {} for mod in self.modality_labels: dist[mod] = dist.get(mod, 0) + 1 return dist # Self-check if __name__ == "__main__": print("Testing MultiModalRetrieval...") # Create dummy embeddings n_per_modality = 50 embed_dim = 768 embeddings_by_modality = { "optical": torch.randn(n_per_modality, embed_dim), "sar": torch.randn(n_per_modality, embed_dim), "multispectral": torch.randn(n_per_modality, embed_dim), } # Normalize for mod in embeddings_by_modality: embeddings_by_modality[mod] = torch.nn.functional.normalize( embeddings_by_modality[mod], dim=1 ) # Build index retrieval = MultiModalRetrieval(embed_dim) retrieval.build_index(embeddings_by_modality) print(f"Index size: {retrieval.size}") print(f"Modality distribution: {retrieval.get_modality_distribution()}") # Same-modal query query = torch.randn(embed_dim) query = torch.nn.functional.normalize(query, dim=0) result = retrieval.same_modal_query(query, modality="optical", k=5) print(f"\nSame-modal (optical→optical):") print(f" Results: {len(result.indices)}") print(f" Modalities: {result.modalities}") # Cross-modal query result = retrieval.cross_modal_query( query, source_modality="optical", target_modality="sar", k=5 ) print(f"\nCross-modal (optical→sar):") print(f" Results: {len(result.indices)}") print(f" Modalities: {result.modalities}") print("\nMultiModalRetrieval test passed!")