""" Gradio UI for satellite image retrieval. Vaporwave/Outrun interface: neon grids, pink-cyan-purple palette, retro-futurism. """ import gradio as gr import time import traceback import numpy as np from PIL import Image from pathlib import Path from typing import Optional from ..retrieval.cross_modal_retrieval import CrossModalRetrieval from ..features.extractor import FeatureExtractor _retrieval: Optional[CrossModalRetrieval] = None _feature_extractor: Optional[FeatureExtractor] = None _gallery_dir: Optional[Path] = None _gallery_metadata: Optional[list] = None def initialize( retrieval: CrossModalRetrieval, feature_extractor: Optional[FeatureExtractor], gallery_dir: Optional[Path] = None, gallery_metadata: Optional[list] = None, ) -> None: global _retrieval, _feature_extractor, _gallery_dir, _gallery_metadata _retrieval = retrieval _feature_extractor = feature_extractor _gallery_dir = Path(gallery_dir) if gallery_dir else None _gallery_metadata = gallery_metadata def _gallery_image_path(idx: int, modality: str) -> Optional[str]: if _gallery_metadata is not None and idx < len(_gallery_metadata): entry = _gallery_metadata[idx] path = Path(entry["gallery_path"]).resolve() if path.exists(): return str(path) if _gallery_dir is not None: path = (_gallery_dir / f"{modality}_{idx}.png").resolve() if path.exists(): return str(path) return None def _load_image_tensor(path, modality): """Load an image and return (PIL preview, torch tensor with proper channels).""" import torch ext = Path(path).suffix.lower() # Try multi-channel TIFF first if ext in (".tif", ".tiff"): try: import tifffile arr = tifffile.imread(str(path)) # Handle different channel arrangements if arr.ndim == 2: # Grayscale → make 3-channel for preview, keep 1ch for features preview = Image.fromarray(arr).convert("RGB") tensor = torch.from_numpy(arr).float().unsqueeze(0) # (1, H, W) tensor = tensor.unsqueeze(0) # (1, 1, H, W) return preview, tensor elif arr.ndim == 3: if arr.shape[-1] in (2, 3, 4, 13): # Channels-last: (H, W, C) tensor = torch.from_numpy(arr).float() tensor = tensor.permute(2, 0, 1).unsqueeze(0) # (1, C, H, W) # RGB preview if arr.shape[-1] >= 3: preview = Image.fromarray(arr[:, :, :3].astype(np.uint8)) else: preview = Image.fromarray(arr[:, :, 0].astype(np.uint8)).convert("RGB") return preview, tensor elif arr.shape[0] in (2, 3, 4, 13): # Channels-first: (C, H, W) tensor = torch.from_numpy(arr).float().unsqueeze(0) # For preview: use first 3 channels or repeat if arr.shape[0] >= 3: preview_arr = np.transpose(arr[:3], (1, 2, 0)) else: preview_arr = np.stack([arr[0]] * 3, axis=-1) if arr.dtype == np.uint16: preview_arr = (preview_arr / 65535.0 * 255).astype(np.uint8) preview = Image.fromarray(preview_arr) return preview, tensor except ImportError: pass # Fallback: PIL img = Image.open(path).convert("RGB") return img, None def retrieve(image, modality: str, k: int, retrieval_type: str, use_sar_adapter: bool = False, use_multiscale: bool = False, lat: float = None, lon: float = None, radius_km: float = 50.0): if image is None: return [], "", "Please upload an image first." if _retrieval is None: return [], "", "System not initialized. Please restart the app." start = time.perf_counter() try: import torch if isinstance(image, str): pil_img, img_tensor = _load_image_tensor(image, modality) else: pil_img = image img_tensor = None if _feature_extractor is not None: if img_tensor is not None and img_tensor.shape[1] not in (3,): # Multi-channel TIFF → use tensor extractor query_embedding = _feature_extractor.extract_features_from_tensor( img_tensor, modality=modality, normalize=True ) elif use_sar_adapter and modality == "sar": from ..features.sar_adapter import SARAdapter adapter = SARAdapter() adapter.eval() img_t = torch.from_numpy(np.array(pil_img)).permute(2, 0, 1).float() / 255.0 if img_t.shape[0] == 3: img_t = img_t[:2] img_t = img_t.unsqueeze(0) with torch.no_grad(): adapted = adapter(img_t) adapted_pil = Image.fromarray( (adapted.squeeze(0).permute(1, 2, 0).numpy() * 255).astype(np.uint8)) query_embedding = _feature_extractor.extract_features( adapted_pil, modality=modality, normalize=True) else: query_embedding = _feature_extractor.extract_features( pil_img, modality=modality, normalize=True) else: embed_dim = _retrieval.embed_dim query_embedding = torch.randn(embed_dim) query_embedding = torch.nn.functional.normalize(query_embedding, dim=0) query_np = query_embedding.unsqueeze(0).numpy().astype(np.float32) if lat is not None and lon is not None: result = _retrieval.search(query_np, modality, k=k, lat=lat, lon=lon, radius_km=radius_km) elif retrieval_type == "same-modal": result = _retrieval.search(query_np, modality, target_modality=modality, k=k) else: result = _retrieval.search(query_np, modality, k=k, strategy="multi") elapsed_ms = (time.perf_counter() - start) * 1000 gallery_images = [] for i, (idx, score) in enumerate(zip(result.indices, result.scores)): mod = result.modalities[i] if result.modalities else modality img_path = _gallery_image_path(idx, mod) if img_path: gallery_images.append(Image.open(img_path)) if not gallery_images: for idx, _ in zip(result.indices, result.scores): np.random.seed(idx) arr = np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8) gallery_images.append(Image.fromarray(arr)) timing_text = f"{elapsed_ms:.0f}ms" n_results = len(result.indices) mod_str = ", ".join(set(result.modalities)) if result.modalities else modality status_text = f"{n_results} results | {mod_str} | {elapsed_ms:.0f}ms" return gallery_images, timing_text, status_text except Exception as exc: tb = traceback.format_exc() return [], "", f"Error: {exc}\n\n{tb}" # --------------------------------------------------------------------------- # Vaporwave Design System # --------------------------------------------------------------------------- VAPORWAVE_CSS = """ """ def _open_image(file): if file is None: return None if isinstance(file, dict): return Image.open(file.get('path') or file.get('url')) if hasattr(file, 'path'): return Image.open(file.path) if hasattr(file, 'name'): return Image.open(file.name) if isinstance(file, str): return Image.open(file) return Image.open(file) def create_app() -> gr.Blocks: def on_upload(file): if file is None: return None path = None if isinstance(file, dict): path = file.get('path') or file.get('url') elif hasattr(file, 'path'): path = file.path elif hasattr(file, 'name'): path = file.name elif isinstance(file, str): path = file if path: pil_img, _ = _load_image_tensor(path, "optical") return pil_img return _open_image(file) def on_retrieve(file, modality, k, retrieval_type): if file is None: return [], "", "Upload an image first." return retrieve(file, modality, int(float(k)), retrieval_type) with gr.Blocks(title="SATCOM // Cross-Modal Retrieval") as app: gr.HTML(VAPORWAVE_CSS) gr.HTML("""
SATCOM // RETRIEVAL
Cross-Modal Satellite Image Search // Optical · SAR · Multispectral
""") with gr.Row(): with gr.Column(scale=1, elem_classes=["vapor-card-left"]): gr.HTML('// INPUT') file_input = gr.File( label="Upload Satellite Image", file_types=[".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp"], ) preview = gr.Image( label="Preview", interactive=False, height=160, ) gr.HTML('// SETTINGS') modality = gr.Dropdown( ["optical", "sar", "multispectral"], value="optical", label="Query Modality", ) retrieval_type = gr.Radio( ["same-modal", "cross-modal"], value="same-modal", label="Retrieval Type", ) k_slider = gr.Slider( 1, 10, value=5, step=1, label="Results (K)", ) btn = gr.Button( "▶ EXECUTE SEARCH", variant="primary", elem_classes=["vapor-btn"], ) gr.HTML("""
OPTICAL RGB · 3ch
SAR Radar · 2ch
MULTI All · 13ch
""") with gr.Column(scale=2, elem_classes=["vapor-card-right"]): gr.HTML('// RESULTS') status = gr.Textbox( label="Status", interactive=False, lines=1, elem_classes=["vapor-status"], ) gallery = gr.Gallery( label="Retrieved Images", columns=5, rows=2, height=360, elem_classes=["vapor-gallery"], ) timing = gr.Textbox( label="Query Time", interactive=False, lines=1, ) gr.HTML("""
OPTICAL SAR MULTISPECTRAL   //   SatCLIP + FAISS + Multi-Index   //   Ayush · Karan · Anurag
""") file_input.change(fn=on_upload, inputs=[file_input], outputs=[preview]) btn.click( fn=on_retrieve, inputs=[file_input, modality, k_slider, retrieval_type], outputs=[gallery, timing, status], ) return app if __name__ == "__main__": app = create_app() print("Gradio app created. Run with: app.launch()")