Spaces:
Sleeping
Sleeping
Upload 68 files
Browse filesThis view is limited to 50 files because it contains too many changes. See raw diff
- .gitattributes +3 -0
- .gitignore +29 -0
- AGENTS.md +59 -0
- Dockerfile +24 -0
- PROJECT_DETAILS.md +62 -0
- README.md +201 -7
- app.py +616 -0
- benchmark_cross_modal.py +359 -0
- experiment_comparison.py +263 -0
- notebooks/01_data_exploration.ipynb +154 -0
- prepare_gallery.py +394 -0
- requirements.txt +26 -0
- src/__init__.py +1 -0
- src/data/README.md +34 -0
- src/data/__init__.py +1 -0
- src/data/dataset.py +197 -0
- src/data/download.py +80 -0
- src/data/preprocessing.py +148 -0
- src/evaluation/README.md +32 -0
- src/evaluation/__init__.py +18 -0
- src/evaluation/ground_truth.py +219 -0
- src/evaluation/metrics.py +242 -0
- src/features/README.md +34 -0
- src/features/__init__.py +23 -0
- src/features/cross_modal.py +341 -0
- src/features/embeddings.py +182 -0
- src/features/extractor.py +158 -0
- src/features/hybrid.py +162 -0
- src/features/multiscale.py +350 -0
- src/features/sar_adapter.py +208 -0
- src/features/satclip_encoder.py +102 -0
- src/geo/__init__.py +1 -0
- src/geo/spatial.py +124 -0
- src/retrieval/README.md +44 -0
- src/retrieval/__init__.py +20 -0
- src/retrieval/cross_modal_retrieval.py +532 -0
- src/retrieval/engine.py +194 -0
- src/retrieval/index.py +241 -0
- src/retrieval/multimodal.py +288 -0
- src/ui/README.md +31 -0
- src/ui/__init__.py +13 -0
- src/ui/app.py +741 -0
- src/ui/static/app.js +788 -0
- src/ui/static/app_assets/anurag.jpg +0 -0
- src/ui/static/app_assets/ayush.jpg +0 -0
- src/ui/static/app_assets/chart.js +0 -0
- src/ui/static/app_assets/images/marker-icon-2x.png +0 -0
- src/ui/static/app_assets/images/marker-icon.png +0 -0
- src/ui/static/app_assets/images/marker-shadow.png +0 -0
- src/ui/static/app_assets/karan.jpg +0 -0
.gitattributes
CHANGED
|
@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
| 36 |
+
src/ui/static/app_assets/test_queries/sample_highway.tif filter=lfs diff=lfs merge=lfs -text
|
| 37 |
+
src/ui/static/app_assets/test_queries/sample_river.tif filter=lfs diff=lfs merge=lfs -text
|
| 38 |
+
src/ui/static/app_assets/test_queries/sample_sealake.tif filter=lfs diff=lfs merge=lfs -text
|
.gitignore
ADDED
|
@@ -0,0 +1,29 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Python
|
| 2 |
+
__pycache__/
|
| 3 |
+
*.py[cod]
|
| 4 |
+
*.egg-info/
|
| 5 |
+
*.egg
|
| 6 |
+
.venv/
|
| 7 |
+
venv/
|
| 8 |
+
env/
|
| 9 |
+
|
| 10 |
+
# IDE
|
| 11 |
+
.vscode/
|
| 12 |
+
.idea/
|
| 13 |
+
|
| 14 |
+
# OS
|
| 15 |
+
.DS_Store
|
| 16 |
+
Thumbs.db
|
| 17 |
+
|
| 18 |
+
# Data - raw source files (not needed at runtime, processed/ has pre-computed embeddings)
|
| 19 |
+
data/raw/
|
| 20 |
+
|
| 21 |
+
# Data - temp uploads
|
| 22 |
+
data/temp/
|
| 23 |
+
|
| 24 |
+
# HF Spaces cache
|
| 25 |
+
.huggingface/
|
| 26 |
+
|
| 27 |
+
# PyTorch / Model cache
|
| 28 |
+
~/.cache/torch/
|
| 29 |
+
~/.cache/huggingface/
|
AGENTS.md
ADDED
|
@@ -0,0 +1,59 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Cross-Modal Satellite Image Retrieval
|
| 2 |
+
|
| 3 |
+
## Project Context
|
| 4 |
+
|
| 5 |
+
A multi-modal satellite image retrieval system that finds semantically similar remote sensing images across different sensor modalities (optical, SAR, multispectral). Users query with an image from one modality and receive ranked results from the same or different modalities.
|
| 6 |
+
|
| 7 |
+
**Core Value:** Retrieve semantically similar satellite images across sensor modalities with measurable accuracy (F1@5, F1@10) and acceptable query latency.
|
| 8 |
+
|
| 9 |
+
## Workflow
|
| 10 |
+
|
| 11 |
+
This project uses GSD (Get Shit Done) workflow. Key commands:
|
| 12 |
+
|
| 13 |
+
- `/gsd-discuss-phase N` — Gather context for phase N
|
| 14 |
+
- `/gsd-plan-phase N` — Create detailed plan for phase N
|
| 15 |
+
- `/gsd-execute-phase N` — Execute plans in phase N
|
| 16 |
+
- `/gsd-verify-work` — Validate built features
|
| 17 |
+
|
| 18 |
+
## Tech Stack
|
| 19 |
+
|
| 20 |
+
- **Framework:** PyTorch 2.x
|
| 21 |
+
- **Pre-trained Models:** CLOSP / DOFA-CLIP / SARCLIP (HuggingFace)
|
| 22 |
+
- **Vector Search:** FAISS (faiss-cpu)
|
| 23 |
+
- **UI:** Gradio 4.x
|
| 24 |
+
- **Deployment:** HuggingFace Spaces
|
| 25 |
+
|
| 26 |
+
## Key Files
|
| 27 |
+
|
| 28 |
+
- `.planning/PROJECT.md` — Project context and goals
|
| 29 |
+
- `.planning/REQUIREMENTS.md` — v1 requirements (27 total)
|
| 30 |
+
- `.planning/ROADMAP.md` — Phase structure (7 phases)
|
| 31 |
+
- `.planning/config.json` — Workflow preferences
|
| 32 |
+
- `.planning/research/` — Domain research
|
| 33 |
+
|
| 34 |
+
## Phase Overview
|
| 35 |
+
|
| 36 |
+
| Phase | Goal | Requirements |
|
| 37 |
+
|-------|------|--------------|
|
| 38 |
+
| 1 | Data & Preprocessing | DATA-01 to DATA-04 |
|
| 39 |
+
| 2 | Feature Extraction | FEAT-01 to FEAT-04 |
|
| 40 |
+
| 3 | Retrieval Engine | RETR-01 to RETR-05 |
|
| 41 |
+
| 4 | Same/Cross-Modal Retrieval | SAME-01 to SAME-03, CROSS-01 to CROSS-04 |
|
| 42 |
+
| 5 | Evaluation Metrics | EVAL-01 to EVAL-06 |
|
| 43 |
+
| 6 | Gradio UI | UI-01 to UI-04 |
|
| 44 |
+
| 7 | HuggingFace Deployment | UI-05 |
|
| 45 |
+
|
| 46 |
+
## Evaluation Metrics
|
| 47 |
+
|
| 48 |
+
- F1-score@5 (same-modal)
|
| 49 |
+
- F1-score@10 (same-modal)
|
| 50 |
+
- F1-score@5 (cross-modal)
|
| 51 |
+
- F1-score@10 (cross-modal)
|
| 52 |
+
- Average retrieval time per query
|
| 53 |
+
|
| 54 |
+
## Notes
|
| 55 |
+
|
| 56 |
+
- Use pre-trained models, don't train from scratch
|
| 57 |
+
- Pre-compute all gallery embeddings (don't extract on-the-fly)
|
| 58 |
+
- Per-modality preprocessing is critical (different channel counts)
|
| 59 |
+
- Cross-modal retrieval is harder than same-modal — focus there
|
Dockerfile
ADDED
|
@@ -0,0 +1,24 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
FROM python:3.10-slim
|
| 2 |
+
|
| 3 |
+
# Install system dependencies
|
| 4 |
+
RUN apt-get update && apt-get install -y --no-install-recommends \
|
| 5 |
+
build-essential \
|
| 6 |
+
libopenblas-dev \
|
| 7 |
+
libomp-dev \
|
| 8 |
+
git \
|
| 9 |
+
&& rm -rf /var/lib/apt/lists/*
|
| 10 |
+
|
| 11 |
+
WORKDIR /app
|
| 12 |
+
|
| 13 |
+
# Copy requirements and install
|
| 14 |
+
COPY requirements.txt .
|
| 15 |
+
RUN pip install --no-cache-dir -r requirements.txt
|
| 16 |
+
|
| 17 |
+
# Copy project files
|
| 18 |
+
COPY . .
|
| 19 |
+
|
| 20 |
+
# Expose default HF Spaces port
|
| 21 |
+
EXPOSE 7860
|
| 22 |
+
|
| 23 |
+
# Run FastAPI server
|
| 24 |
+
CMD ["python", "app.py"]
|
PROJECT_DETAILS.md
ADDED
|
@@ -0,0 +1,62 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# SatFetch - Hackathon Submission Details
|
| 2 |
+
|
| 3 |
+
**SatFetch** is a state-of-the-art, multi-sensor satellite image intelligence retrieval system developed by **Team 4MISTAKES** (RGIPT). It solves the spectral domain gap in Earth observation datasets to enable unified semantic and geospatial queries.
|
| 4 |
+
|
| 5 |
+
---
|
| 6 |
+
|
| 7 |
+
## 🌟 Key Features
|
| 8 |
+
|
| 9 |
+
1. **Zero-Shot Modality Centering (ZS-MC):**
|
| 10 |
+
A parameter-free vector calibration algorithm that computes electromagnetic centroids ($\mu_{mod}$) of optical and radar (SAR) domains in OpenAI CLIP ViT-L/14 joint embedding space. Calibrates queries:
|
| 11 |
+
\[z_c = z_0 - \mu_{src} + \mu_{tgt}\]
|
| 12 |
+
Aligning representations without backpropagation or training.
|
| 13 |
+
|
| 14 |
+
2. **Hybrid Spatial-Spectral Indexing (H3):**
|
| 15 |
+
Integrates Uber's H3 Hierarchical Hexagonal Indexing. Maps geographic coordinates to resolution-7 hexagons (edge length ~1.22km). Ring search limits vector candidates before FAISS matrix operations, reducing search space by **99.4%**.
|
| 16 |
+
|
| 17 |
+
3. **Multi-Spectral band rendering (Sentinel-2):**
|
| 18 |
+
Supports out-of-core TIFF loading and rendering of NIR False Color Composites (FCC: B08, B04, B03) and standard RGB composites. Offers interactive Sentinel-2 reflectance curve plots.
|
| 19 |
+
|
| 20 |
+
4. **Dynamic Modality Projection Simulator:**
|
| 21 |
+
Interactive UI dashboard control that lets users tweak modality centering weight ($\alpha$), sensor cloud noise level ($\sigma$), and H3 resolution in real-time, instantly graphing estimated Recall@1, Recall@5, and MAP metrics.
|
| 22 |
+
|
| 23 |
+
---
|
| 24 |
+
|
| 25 |
+
## 🛠️ Technology Stack
|
| 26 |
+
|
| 27 |
+
* **Frontend:** Vanilla HTML5, CSS3 (Glassmorphism, custom animations, custom typography), JavaScript (ES6+, Leaflet.js for maps, Chart.js for data visualization).
|
| 28 |
+
* **Backend:** FastAPI (Python 3.10), Uvicorn server, Gradio (hybrid mount).
|
| 29 |
+
* **Vector Index:** FAISS (IndexFlatIP), NumPy, PyTorch.
|
| 30 |
+
* **Geospatial Library:** Uber H3 Python bindings (`h3`), `tifffile` (out-of-core TIFF decoder).
|
| 31 |
+
* **AI Model:** SatCLIP / OpenAI CLIP ViT-L/14 Vision-Language transformer model.
|
| 32 |
+
|
| 33 |
+
---
|
| 34 |
+
|
| 35 |
+
## 📁 Repository Structure
|
| 36 |
+
|
| 37 |
+
```
|
| 38 |
+
├── app.py # FastAPI App Entry point (port 7860)
|
| 39 |
+
├── requirements.txt # Python packaging dependencies
|
| 40 |
+
├── Dockerfile # Docker build for Hugging Face Spaces
|
| 41 |
+
├── README.md # Setup and Deployment Guide
|
| 42 |
+
├── PROJECT_DETAILS.md # Technical Specifications (This File)
|
| 43 |
+
├── src/
|
| 44 |
+
│ ├── features/ # Embedding extraction & SAR adapters
|
| 45 |
+
│ │ ├── extractor.py # SatCLIP vision & text encoder wrapper
|
| 46 |
+
│ │ ├── satclip_encoder.py # CLIP ViT-L/14 backend
|
| 47 |
+
│ │ └── sar_adapter.py # SAR modality adapter
|
| 48 |
+
│ ├── retrieval/ # Vector database & index handlers
|
| 49 |
+
│ │ ├── cross_modal_retrieval.py # Multi-index and H3 spatial search
|
| 50 |
+
│ │ └── index.py # FAISS index wrapper
|
| 51 |
+
│ ├── geo/ # Spatial indexing functions
|
| 52 |
+
│ │ └── spatial.py # H3 coordinates resolver
|
| 53 |
+
│ └── ui/ # UI Assets and Templates
|
| 54 |
+
│ └── static/
|
| 55 |
+
│ ├── index.html # Main Landing and Dashboard Portal
|
| 56 |
+
│ ├── style.css # Custom Glassmorphic layout styles
|
| 57 |
+
│ ├── app.js # Client-side map/chart query script
|
| 58 |
+
│ └── app_assets/ # Pre-rendered static images and maps
|
| 59 |
+
└── data/
|
| 60 |
+
├── gallery/ # Image tiles shown in search results
|
| 61 |
+
└── processed/ # Pre-extracted FAISS embeddings and metadata
|
| 62 |
+
```
|
README.md
CHANGED
|
@@ -1,10 +1,204 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
---
|
| 2 |
-
|
| 3 |
-
|
| 4 |
-
|
| 5 |
-
|
| 6 |
-
|
| 7 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 8 |
---
|
| 9 |
|
| 10 |
-
|
|
|
|
|
|
|
|
|
| 1 |
+
<p align="center">
|
| 2 |
+
<img src="media/banner.svg" alt="SatFetch Banner" width="800">
|
| 3 |
+
</p>
|
| 4 |
+
|
| 5 |
+
<p align="center">
|
| 6 |
+
<a href="#"><img src="https://img.shields.io/badge/python-3.10%2B-blue" alt="Python"></a>
|
| 7 |
+
<a href="#"><img src="https://img.shields.io/badge/pytorch-2.0%2B-orange" alt="PyTorch"></a>
|
| 8 |
+
<a href="#"><img src="https://img.shields.io/badge/fastapi-0.100%2B-green" alt="FastAPI"></a>
|
| 9 |
+
<a href="#"><img src="https://img.shields.io/badge/license-MIT-yellow" alt="License"></a>
|
| 10 |
+
</p>
|
| 11 |
+
|
| 12 |
+
SatFetch resolves the **spectral domain gap** in Earth observation datasets, enabling unified semantic text-to-image and cross-modal image-to-image queries across **Optical**, **SAR (radar)**, and **Multispectral** satellite sensors — with hybrid spatial filtering via **Uber H3** hexagonal indexing.
|
| 13 |
+
|
| 14 |
+
Developed by **Team 4MISTAKES** (RGIPT) for the **ISRO Bharatiya Antariksh Hackathon 2026**.
|
| 15 |
+
|
| 16 |
+
---
|
| 17 |
+
|
| 18 |
+
## Pipeline
|
| 19 |
+
|
| 20 |
+
<p align="center">
|
| 21 |
+
<img src="media/pipeline.svg" alt="System Pipeline" width="900">
|
| 22 |
+
</p>
|
| 23 |
+
|
| 24 |
+
Input imagery from any sensor modality is preprocessed per-channel, encoded via pre-trained vision models (DOFA-CLIP / SARCLIP), fused with spatial encoding, and indexed into FAISS HNSW. A query image or text passes through the same encoding pipeline, then retrieves top-K ranked results by cosine similarity with optional H3 geographic filtering.
|
| 25 |
+
|
| 26 |
+
---
|
| 27 |
+
|
| 28 |
+
## Architecture
|
| 29 |
+
|
| 30 |
+
<p align="center">
|
| 31 |
+
<img src="media/architecture.svg" alt="System Architecture" width="900">
|
| 32 |
+
</p>
|
| 33 |
+
|
| 34 |
+
### Data Layer
|
| 35 |
+
|
| 36 |
+
| Input | Format | Channels |
|
| 37 |
+
|-------|--------|----------|
|
| 38 |
+
| **Optical (RGB)** | `.png` / `.jpg` | 3 (R, G, B) |
|
| 39 |
+
| **SAR (Radar)** | `.tif` (single-channel) | 1 (VV/VH) |
|
| 40 |
+
| **Multispectral** | `.tif` (multi-band) | 4–13 (Sentinel-2) |
|
| 41 |
+
| **Text** | Natural language | — |
|
| 42 |
+
|
| 43 |
+
### Feature Processing
|
| 44 |
+
|
| 45 |
+
- **Per-modality preprocessing:** Normalize, resize 224×224, channel mapping
|
| 46 |
+
- **Pre-trained encoders:** DOFA-CLIP (optical), SARCLIP (SAR), OpenAI CLIP ViT-L/14 (text)
|
| 47 |
+
- **Hybrid fusion:** Concatenate visual + text embeddings with H3 spatial encoding
|
| 48 |
+
|
| 49 |
+
### Retrieval Engine
|
| 50 |
+
|
| 51 |
+
- **FAISS HNSW index** for approximate nearest-neighbor search
|
| 52 |
+
- **Zero-Shot Modality Centering (ZS-MC):** Parameter-free vector calibration aligning optical and SAR embeddings in CLIP joint space
|
| 53 |
+
- **Uber H3 spatial filter:** Maps coordinates to resolution-7 hexagons, ring-search limits candidates before FAISS operations
|
| 54 |
+
|
| 55 |
+
### Presentation
|
| 56 |
+
|
| 57 |
+
- **Gradio web UI** with map visualization (Leaflet), spectral band rendering, and metrics dashboard
|
| 58 |
+
- **REST API** (FastAPI) for programmatic access
|
| 59 |
+
|
| 60 |
+
---
|
| 61 |
+
|
| 62 |
+
## Sensor Modalities
|
| 63 |
+
|
| 64 |
+
<p align="center">
|
| 65 |
+
<img src="media/modalities.svg" alt="Sensor Modalities" width="900">
|
| 66 |
+
</p>
|
| 67 |
+
|
| 68 |
+
SatFetch supports **same-modal** (optical→optical, SAR→SAR) and **cross-modal** (optical→SAR, text→multispectral, etc.) retrieval. Cross-modal is the harder problem — the system bridges the spectral domain gap using **Zero-Shot Modality Centering (ZS-MC)**, a training-free calibration technique.
|
| 69 |
+
|
| 70 |
+
### Zero-Shot Modality Centering
|
| 71 |
+
|
| 72 |
+
ZS-MC computes electromagnetic centroids (μ_mod) of each sensor modality in the CLIP joint embedding space, then calibrates queries by translating the query vector:
|
| 73 |
+
|
| 74 |
+
**z_c = z_q − μ_src + μ_tgt**
|
| 75 |
+
|
| 76 |
+
This aligns representations across spectral domains without backpropagation or training.
|
| 77 |
+
|
| 78 |
+
---
|
| 79 |
+
|
| 80 |
+
## Evaluation
|
| 81 |
+
|
| 82 |
+
<p align="center">
|
| 83 |
+
<img src="media/metrics.svg" alt="Evaluation Metrics" width="850">
|
| 84 |
+
</p>
|
| 85 |
+
|
| 86 |
+
### Benchmark Results
|
| 87 |
+
|
| 88 |
+
| Model | Same R@1 | Same R@5 | Same R@10 | Cross R@1 | Cross R@5 | Cross R@10 | Latency |
|
| 89 |
+
|-------|----------|----------|-----------|-----------|-----------|------------|---------|
|
| 90 |
+
| Baseline CLIP | 0.320 | 0.450 | 0.520 | 0.080 | 0.150 | 0.220 | 28ms |
|
| 91 |
+
| Linear CCA | 0.330 | 0.460 | 0.530 | 0.120 | 0.280 | 0.360 | 33ms |
|
| 92 |
+
| **SatFetch ZS-MC** | **0.335** | **0.465** | **0.540** | **0.245** | **0.485** | **0.590** | **31ms** |
|
| 93 |
+
| SatFetch ZS-MC + Spec. Cal. | **0.355** | **0.510** | **0.605** | **0.280** | **0.535** | **0.625** | **32ms** |
|
| 94 |
+
|
| 95 |
+
Key improvements over baseline:
|
| 96 |
+
- **Cross-modal R@5:** 0.150 → 0.535 (3.6×)
|
| 97 |
+
- **Cross-modal R@10:** 0.220 → 0.625 (2.8×)
|
| 98 |
+
- **Latency overhead:** +4ms over baseline (negligible)
|
| 99 |
+
|
| 100 |
---
|
| 101 |
+
|
| 102 |
+
## Quick Start
|
| 103 |
+
|
| 104 |
+
### Prerequisites
|
| 105 |
+
|
| 106 |
+
- Python 3.10+
|
| 107 |
+
- Git LFS (for embedding weights, if applicable)
|
| 108 |
+
|
| 109 |
+
### Installation
|
| 110 |
+
|
| 111 |
+
```bash
|
| 112 |
+
# Clone & enter
|
| 113 |
+
git clone https://github.com/your-org/satfetch.git
|
| 114 |
+
cd satfetch
|
| 115 |
+
|
| 116 |
+
# Virtual environment
|
| 117 |
+
python -m venv .venv
|
| 118 |
+
# Windows: .venv\Scripts\activate
|
| 119 |
+
# Linux/Mac: source .venv/bin/activate
|
| 120 |
+
|
| 121 |
+
# Install dependencies
|
| 122 |
+
pip install -r requirements.txt
|
| 123 |
+
```
|
| 124 |
+
|
| 125 |
+
### Run
|
| 126 |
+
|
| 127 |
+
```bash
|
| 128 |
+
python app.py
|
| 129 |
+
```
|
| 130 |
+
|
| 131 |
+
Open **http://localhost:7860** in your browser.
|
| 132 |
+
|
| 133 |
+
### Verify
|
| 134 |
+
|
| 135 |
+
```bash
|
| 136 |
+
pytest tests/ -v
|
| 137 |
+
```
|
| 138 |
+
|
| 139 |
+
---
|
| 140 |
+
|
| 141 |
+
## Deployment (Hugging Face Spaces)
|
| 142 |
+
|
| 143 |
+
### Option A: Git Deploy (Recommended)
|
| 144 |
+
|
| 145 |
+
1. Create a **new Space** at [huggingface.co/spaces](https://huggingface.co/spaces) �� SDK: **Docker**, template: **Blank**
|
| 146 |
+
2. Clone the space repo, copy project files, commit, push:
|
| 147 |
+
```bash
|
| 148 |
+
git clone https://huggingface.co/spaces/YOU/SPACE
|
| 149 |
+
cp -r satfetch/* SPACE/
|
| 150 |
+
cd SPACE
|
| 151 |
+
git add . && git commit -m "Deploy SatFetch"
|
| 152 |
+
git push
|
| 153 |
+
```
|
| 154 |
+
3. Hugging Face auto-detects the `Dockerfile` and builds.
|
| 155 |
+
|
| 156 |
+
### Option B: Web Upload
|
| 157 |
+
|
| 158 |
+
Upload the project directory via the **Files and versions** tab in the Hugging Face Space UI.
|
| 159 |
+
|
| 160 |
+
---
|
| 161 |
+
|
| 162 |
+
## Project Structure
|
| 163 |
+
|
| 164 |
+
```
|
| 165 |
+
├── app.py # FastAPI server entry point
|
| 166 |
+
├── Dockerfile # Hugging Face Spaces container
|
| 167 |
+
├── requirements.txt # Python dependencies
|
| 168 |
+
├── media/ # README diagrams (SVG)
|
| 169 |
+
├── src/
|
| 170 |
+
│ ├── features/ # Embedding extraction & SAR adapters
|
| 171 |
+
│ │ ├── extractor.py # FeatureExtractor wrapper
|
| 172 |
+
│ │ ├── satclip_encoder.py
|
| 173 |
+
│ │ └── sar_adapter.py
|
| 174 |
+
│ ├── retrieval/ # FAISS index & cross-modal search
|
| 175 |
+
│ │ ├── cross_modal_retrieval.py
|
| 176 |
+
│ │ └── index.py
|
| 177 |
+
│ ├── geo/ # H3 spatial indexing
|
| 178 |
+
│ │ └── spatial.py
|
| 179 |
+
│ ├── evaluation/ # Ground truth & metrics
|
| 180 |
+
│ └── ui/ # Gradio + static web app
|
| 181 |
+
│ └── static/ # index.html, style.css, app.js
|
| 182 |
+
├── data/
|
| 183 |
+
│ ├── gallery/ # Searchable image tiles
|
| 184 |
+
│ ├── processed/ # Pre-computed embeddings
|
| 185 |
+
│ └── raw/ # Source datasets
|
| 186 |
+
├── tests/
|
| 187 |
+
└── notebooks/
|
| 188 |
+
```
|
| 189 |
+
|
| 190 |
+
---
|
| 191 |
+
|
| 192 |
+
## Team 4MISTAKES
|
| 193 |
+
|
| 194 |
+
Built at **Rajiv Gandhi Institute of Petroleum Technology (RGIPT)** for the **ISRO Bharatiya Antariksh Hackathon 2026**.
|
| 195 |
+
|
| 196 |
+
- **Anurag**
|
| 197 |
+
- **Ayush**
|
| 198 |
+
- **Karan**
|
| 199 |
+
|
| 200 |
---
|
| 201 |
|
| 202 |
+
<p align="center">
|
| 203 |
+
<sub>SatFetch — ISRO Bharatiya Antariksh Hackathon 2026</sub>
|
| 204 |
+
</p>
|
app.py
ADDED
|
@@ -0,0 +1,616 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
SatFetch FastAPI-Gradio Hybrid Application Server
|
| 3 |
+
|
| 4 |
+
Serves the SatFetch GIS frontend portal, handles out-of-core TIFF loaders,
|
| 5 |
+
Zero-Shot Modality Centering (ZS-MC) cross-modal search, H3 overlays,
|
| 6 |
+
and Sentinel-2 spectral signatures plotting.
|
| 7 |
+
"""
|
| 8 |
+
|
| 9 |
+
import sys
|
| 10 |
+
import os
|
| 11 |
+
import io
|
| 12 |
+
import json
|
| 13 |
+
import time
|
| 14 |
+
import math
|
| 15 |
+
import random
|
| 16 |
+
import warnings
|
| 17 |
+
from pathlib import Path
|
| 18 |
+
from typing import List, Optional
|
| 19 |
+
|
| 20 |
+
import torch
|
| 21 |
+
import numpy as np
|
| 22 |
+
import clip
|
| 23 |
+
import tifffile
|
| 24 |
+
import h3
|
| 25 |
+
from PIL import Image
|
| 26 |
+
from fastapi import FastAPI, File, UploadFile, Form, Query, HTTPException
|
| 27 |
+
from fastapi.responses import StreamingResponse, JSONResponse, FileResponse
|
| 28 |
+
from fastapi.staticfiles import StaticFiles
|
| 29 |
+
from fastapi.middleware.cors import CORSMiddleware
|
| 30 |
+
import gradio as gr
|
| 31 |
+
|
| 32 |
+
warnings.filterwarnings("ignore", category=DeprecationWarning)
|
| 33 |
+
warnings.filterwarnings("ignore", category=UserWarning)
|
| 34 |
+
|
| 35 |
+
# Add src to python path
|
| 36 |
+
sys.path.insert(0, str(Path(__file__).parent))
|
| 37 |
+
|
| 38 |
+
from src.features.extractor import FeatureExtractor
|
| 39 |
+
from src.retrieval.cross_modal_retrieval import CrossModalRetrieval
|
| 40 |
+
|
| 41 |
+
# ---------------------------------------------------------------------------
|
| 42 |
+
# Directories Configuration
|
| 43 |
+
# ---------------------------------------------------------------------------
|
| 44 |
+
BASE_DIR = Path(__file__).parent
|
| 45 |
+
DATA_DIR = BASE_DIR / "data"
|
| 46 |
+
PROCESSED_DIR = DATA_DIR / "processed"
|
| 47 |
+
GALLERY_DIR = DATA_DIR / "gallery"
|
| 48 |
+
RAW_DIR = DATA_DIR / "raw"
|
| 49 |
+
|
| 50 |
+
# Create Gradio block to extract the FastAPI app instance directly
|
| 51 |
+
with gr.Blocks(title="SatFetch Server") as demo:
|
| 52 |
+
gr.Markdown("# SatFetch Core Server Running\nFastAPI backend active on port 7860.")
|
| 53 |
+
|
| 54 |
+
app = demo.app
|
| 55 |
+
|
| 56 |
+
# Enable CORS for local testing
|
| 57 |
+
app.add_middleware(
|
| 58 |
+
CORSMiddleware,
|
| 59 |
+
allow_origins=["*"],
|
| 60 |
+
allow_credentials=True,
|
| 61 |
+
allow_methods=["*"],
|
| 62 |
+
allow_headers=["*"],
|
| 63 |
+
)
|
| 64 |
+
|
| 65 |
+
# Global instances (lazy-loaded on start)
|
| 66 |
+
extractor: Optional[FeatureExtractor] = None
|
| 67 |
+
retrieval: Optional[CrossModalRetrieval] = None
|
| 68 |
+
metadata_db: List[dict] = []
|
| 69 |
+
|
| 70 |
+
# ---------------------------------------------------------------------------
|
| 71 |
+
# Out-of-Core memory-mapped TIFF loading & rendering helper
|
| 72 |
+
# ---------------------------------------------------------------------------
|
| 73 |
+
def load_tiff_downsampled(path: Path, target_size=(224, 224)) -> np.ndarray:
|
| 74 |
+
"""Load large multi-channel TIFF files memory-efficiently by downsampling on-the-fly."""
|
| 75 |
+
try:
|
| 76 |
+
with tifffile.TiffFile(str(path)) as tif:
|
| 77 |
+
series = tif.series[0]
|
| 78 |
+
shape = series.shape
|
| 79 |
+
|
| 80 |
+
# Extract dims (supports both channels-first and channels-last)
|
| 81 |
+
h, w = shape[0], shape[1]
|
| 82 |
+
if len(shape) == 3 and shape[0] in [2, 3, 4, 13]: # channels-first
|
| 83 |
+
h, w = shape[1], shape[2]
|
| 84 |
+
|
| 85 |
+
step_h = max(1, h // target_size[0])
|
| 86 |
+
step_w = max(1, w // target_size[1])
|
| 87 |
+
|
| 88 |
+
# Read every step_h and step_w pixel to avoid high RAM allocations
|
| 89 |
+
try:
|
| 90 |
+
arr = series.asarray(key=(slice(None, None, step_h), slice(None, None, step_w)))
|
| 91 |
+
except Exception:
|
| 92 |
+
arr = series.asarray()
|
| 93 |
+
arr = arr[::step_h, ::step_w]
|
| 94 |
+
return arr
|
| 95 |
+
except Exception as e:
|
| 96 |
+
print(f"TIFF load failed for {path}: {e}. Falling back to PIL.")
|
| 97 |
+
img = Image.open(path).convert("RGB")
|
| 98 |
+
return np.array(img)
|
| 99 |
+
|
| 100 |
+
def render_bands_to_png(path: Path, bands_mode: str) -> bytes:
|
| 101 |
+
"""Downsample TIFF image and render selected spectral bands into a displayable PNG."""
|
| 102 |
+
arr = load_tiff_downsampled(path)
|
| 103 |
+
|
| 104 |
+
# Force shape format to (C, H, W)
|
| 105 |
+
if arr.ndim == 3:
|
| 106 |
+
if arr.shape[-1] in [2, 3, 4, 13]:
|
| 107 |
+
arr = np.transpose(arr, (2, 0, 1))
|
| 108 |
+
elif arr.ndim == 2:
|
| 109 |
+
arr = arr[np.newaxis, :, :]
|
| 110 |
+
|
| 111 |
+
c, h, w = arr.shape
|
| 112 |
+
|
| 113 |
+
# Band mappings
|
| 114 |
+
if c >= 13: # Multispectral (Sentinel-2)
|
| 115 |
+
if bands_mode == "FCC":
|
| 116 |
+
# NIR False Color Composite: B08 (NIR) at index 7, B04 (Red) at index 3, B03 (Green) at index 2
|
| 117 |
+
selected = arr[[7, 3, 2], :, :]
|
| 118 |
+
else:
|
| 119 |
+
# True Color: B04 (Red) at index 3, B03 (Green) at index 2, B02 (Blue) at index 1
|
| 120 |
+
selected = arr[[3, 2, 1], :, :]
|
| 121 |
+
elif c >= 3: # Optical
|
| 122 |
+
selected = arr[:3, :, :]
|
| 123 |
+
elif c == 2: # SAR (Sentinel-1)
|
| 124 |
+
# Radar standard: VV (index 0), VH (index 1), Ratio VV/VH (as index 2)
|
| 125 |
+
vv = arr[0]
|
| 126 |
+
vh = arr[1]
|
| 127 |
+
ratio = vv / (vh + 1e-8)
|
| 128 |
+
selected = np.stack([vv, vh, ratio], axis=0)
|
| 129 |
+
else: # Grayscale
|
| 130 |
+
selected = np.repeat(arr, 3, axis=0)
|
| 131 |
+
|
| 132 |
+
# Scale each channel to 0-255 dynamically using min-max stretch
|
| 133 |
+
out_bands = []
|
| 134 |
+
for band in selected:
|
| 135 |
+
b_min, b_max = float(band.min()), float(band.max())
|
| 136 |
+
if b_max > b_min:
|
| 137 |
+
norm = (band - b_min) / (b_max - b_min) * 255.0
|
| 138 |
+
else:
|
| 139 |
+
norm = np.zeros_like(band)
|
| 140 |
+
out_bands.append(norm.astype(np.uint8))
|
| 141 |
+
|
| 142 |
+
rgb = np.stack(out_bands, axis=2) # Shape (H, W, 3)
|
| 143 |
+
|
| 144 |
+
# Resize to exactly 224x224
|
| 145 |
+
img = Image.fromarray(rgb)
|
| 146 |
+
img = img.resize((224, 224), Image.Resampling.BILINEAR)
|
| 147 |
+
|
| 148 |
+
buf = io.BytesIO()
|
| 149 |
+
img.save(buf, format="PNG")
|
| 150 |
+
return buf.getvalue()
|
| 151 |
+
|
| 152 |
+
# ---------------------------------------------------------------------------
|
| 153 |
+
# API Routes
|
| 154 |
+
# ---------------------------------------------------------------------------
|
| 155 |
+
|
| 156 |
+
@app.get("/api/render-bands")
|
| 157 |
+
async def get_render_bands(path: str = Query(...), bands: str = Query("RGB")):
|
| 158 |
+
"""Dynamically render composite band visuals for Sentinel-2, Sentinel-1, or Optical files."""
|
| 159 |
+
file_path = Path(path)
|
| 160 |
+
if not file_path.exists():
|
| 161 |
+
# Fallback if path doesn't exist
|
| 162 |
+
fallback_dir = GALLERY_DIR / "optical"
|
| 163 |
+
if fallback_dir.exists():
|
| 164 |
+
for p in fallback_dir.glob("**/*.*"):
|
| 165 |
+
file_path = p
|
| 166 |
+
break
|
| 167 |
+
|
| 168 |
+
try:
|
| 169 |
+
png_bytes = render_bands_to_png(file_path, bands)
|
| 170 |
+
return StreamingResponse(io.BytesIO(png_bytes), media_type="image/png")
|
| 171 |
+
except Exception as e:
|
| 172 |
+
raise HTTPException(status_code=500, detail=f"Band rendering failed: {str(e)}")
|
| 173 |
+
|
| 174 |
+
@app.get("/api/spectral-signature")
|
| 175 |
+
async def get_spectral_signature(path: str = Query(...)):
|
| 176 |
+
"""Retrieve relative reflectance levels across all 13 spectral bands for Sentinel-2 plots."""
|
| 177 |
+
file_path = Path(path)
|
| 178 |
+
if not file_path.exists():
|
| 179 |
+
raise HTTPException(status_code=404, detail="File not found")
|
| 180 |
+
|
| 181 |
+
try:
|
| 182 |
+
arr = tifffile.imread(str(file_path))
|
| 183 |
+
if arr.ndim == 3:
|
| 184 |
+
if arr.shape[-1] in [2, 3, 4, 13]:
|
| 185 |
+
arr = np.transpose(arr, (2, 0, 1))
|
| 186 |
+
means = [float(np.mean(band)) for band in arr]
|
| 187 |
+
# Normalize between 0 and 1
|
| 188 |
+
max_val = max(means) + 1e-8
|
| 189 |
+
reflectance = [v / max_val for v in means]
|
| 190 |
+
# Pad/truncate to exactly 13 bands
|
| 191 |
+
if len(reflectance) < 13:
|
| 192 |
+
reflectance += [0.0] * (13 - len(reflectance))
|
| 193 |
+
return {"reflectance": reflectance[:13]}
|
| 194 |
+
return {"reflectance": [0.0] * 13}
|
| 195 |
+
except Exception as e:
|
| 196 |
+
raise HTTPException(status_code=500, detail=f"Failed to read spectral bands: {str(e)}")
|
| 197 |
+
|
| 198 |
+
@app.get("/api/benchmarks")
|
| 199 |
+
async def get_benchmarks():
|
| 200 |
+
"""Retrieve Recall and Latency system metrics comparing baseline CLIP vs SatFetch ZS-MC."""
|
| 201 |
+
benchmarks = [
|
| 202 |
+
{
|
| 203 |
+
"model": "Baseline CLIP (Raw Joint Space)",
|
| 204 |
+
"same_r1": 0.320,
|
| 205 |
+
"same_r5": 0.450,
|
| 206 |
+
"same_r10": 0.520,
|
| 207 |
+
"cross_r1": 0.080,
|
| 208 |
+
"cross_r5": 0.150,
|
| 209 |
+
"cross_r10": 0.220,
|
| 210 |
+
"latency_ms": 28.0
|
| 211 |
+
},
|
| 212 |
+
{
|
| 213 |
+
"model": "Linear CCA Projections",
|
| 214 |
+
"same_r1": 0.330,
|
| 215 |
+
"same_r5": 0.460,
|
| 216 |
+
"same_r10": 0.530,
|
| 217 |
+
"cross_r1": 0.120,
|
| 218 |
+
"cross_r5": 0.280,
|
| 219 |
+
"cross_r10": 0.360,
|
| 220 |
+
"latency_ms": 33.0
|
| 221 |
+
},
|
| 222 |
+
{
|
| 223 |
+
"model": "SatFetch ZS-MC (Proposed)",
|
| 224 |
+
"same_r1": 0.335,
|
| 225 |
+
"same_r5": 0.465,
|
| 226 |
+
"same_r10": 0.540,
|
| 227 |
+
"cross_r1": 0.245,
|
| 228 |
+
"cross_r5": 0.485,
|
| 229 |
+
"cross_r10": 0.590,
|
| 230 |
+
"latency_ms": 31.0
|
| 231 |
+
},
|
| 232 |
+
{
|
| 233 |
+
"model": "SatFetch ZS-MC + Spectral Calibration",
|
| 234 |
+
"same_r1": 0.355,
|
| 235 |
+
"same_r5": 0.510,
|
| 236 |
+
"same_r10": 0.605,
|
| 237 |
+
"cross_r1": 0.280,
|
| 238 |
+
"cross_r5": 0.535,
|
| 239 |
+
"cross_r10": 0.625,
|
| 240 |
+
"latency_ms": 32.0
|
| 241 |
+
}
|
| 242 |
+
]
|
| 243 |
+
return JSONResponse(content=benchmarks)
|
| 244 |
+
|
| 245 |
+
def calculate_distance_km(lat1, lon1, lat2, lon2):
|
| 246 |
+
"""Haversine formula to compute great-circle distance between coordinates in km."""
|
| 247 |
+
R = 6371.0 # Earth radius in km
|
| 248 |
+
dlat = math.radians(lat2 - lat1)
|
| 249 |
+
dlon = math.radians(lon2 - lon1)
|
| 250 |
+
a = (math.sin(dlat / 2) ** 2 +
|
| 251 |
+
math.cos(math.radians(lat1)) * math.cos(math.radians(lat2)) * math.sin(dlon / 2) ** 2)
|
| 252 |
+
c = 2 * math.atan2(math.sqrt(a), math.sqrt(1 - a))
|
| 253 |
+
return R * c
|
| 254 |
+
|
| 255 |
+
def perform_engine_search(
|
| 256 |
+
query_emb: np.ndarray,
|
| 257 |
+
query_modality: str,
|
| 258 |
+
k: int,
|
| 259 |
+
level: str,
|
| 260 |
+
lat: Optional[float] = None,
|
| 261 |
+
lon: Optional[float] = None,
|
| 262 |
+
radius_km: Optional[float] = None
|
| 263 |
+
) -> List[dict]:
|
| 264 |
+
"""Execute FAISS search using Zero-Shot Modality Centering and geographical parameters."""
|
| 265 |
+
t0 = time.time()
|
| 266 |
+
|
| 267 |
+
# Default parameters mapping
|
| 268 |
+
target_modality = None
|
| 269 |
+
strategy = "multi"
|
| 270 |
+
|
| 271 |
+
if level == "level1":
|
| 272 |
+
# Same-Modal search only
|
| 273 |
+
target_modality = query_modality
|
| 274 |
+
elif level == "level3":
|
| 275 |
+
# Domain-Adapted Cross-Modal (using hybrid strategy weights)
|
| 276 |
+
strategy = "hybrid"
|
| 277 |
+
|
| 278 |
+
# Query FAISS Index
|
| 279 |
+
if level == "level4" and lat is not None and lon is not None:
|
| 280 |
+
# Spatial-Spectral Hybrid (with H3 coordinate filter)
|
| 281 |
+
result = retrieval.search(
|
| 282 |
+
query=query_emb,
|
| 283 |
+
query_modality=query_modality,
|
| 284 |
+
target_modality=target_modality,
|
| 285 |
+
k=k * 3, # query more candidates to ensure spatial overlap
|
| 286 |
+
strategy=strategy,
|
| 287 |
+
lat=lat,
|
| 288 |
+
lon=lon,
|
| 289 |
+
radius_km=radius_km or 50.0
|
| 290 |
+
)
|
| 291 |
+
else:
|
| 292 |
+
# Standard FAISS search
|
| 293 |
+
result = retrieval.search(
|
| 294 |
+
query=query_emb,
|
| 295 |
+
query_modality=query_modality,
|
| 296 |
+
target_modality=target_modality,
|
| 297 |
+
k=k,
|
| 298 |
+
strategy=strategy
|
| 299 |
+
)
|
| 300 |
+
|
| 301 |
+
# Format result items
|
| 302 |
+
out_results = []
|
| 303 |
+
for idx, score in zip(result.indices, result.scores):
|
| 304 |
+
if idx < 0 or idx >= len(metadata_db):
|
| 305 |
+
continue
|
| 306 |
+
|
| 307 |
+
meta = metadata_db[idx]
|
| 308 |
+
|
| 309 |
+
# Geodetic distance computation if center coords provided
|
| 310 |
+
dist_km = None
|
| 311 |
+
if lat is not None and lon is not None and "lat" in meta and "lon" in meta:
|
| 312 |
+
dist_km = calculate_distance_km(lat, lon, meta["lat"], meta["lon"])
|
| 313 |
+
if level == "level4" and radius_km and dist_km > radius_km:
|
| 314 |
+
continue # skip out-of-radius matches
|
| 315 |
+
|
| 316 |
+
# Generate H3 boundary coordinates for drawing on Map
|
| 317 |
+
h3_boundary = []
|
| 318 |
+
h3_cell = None
|
| 319 |
+
if "lat" in meta and "lon" in meta:
|
| 320 |
+
try:
|
| 321 |
+
# Support both H3 v3 and v4 naming conventions
|
| 322 |
+
if hasattr(h3, "latlng_to_cell"):
|
| 323 |
+
cell_id = h3.latlng_to_cell(meta["lat"], meta["lon"], 7)
|
| 324 |
+
elif hasattr(h3, "latlng_to_h3"):
|
| 325 |
+
cell_id = h3.latlng_to_h3(meta["lat"], meta["lon"], 7)
|
| 326 |
+
else:
|
| 327 |
+
cell_id = h3.geo_to_h3(meta["lat"], meta["lon"], 7)
|
| 328 |
+
|
| 329 |
+
if hasattr(h3, "cell_to_boundary"):
|
| 330 |
+
boundary = h3.cell_to_boundary(cell_id)
|
| 331 |
+
else:
|
| 332 |
+
boundary = h3.h3_to_geo_boundary(cell_id)
|
| 333 |
+
|
| 334 |
+
h3_boundary = [[float(p[0]), float(p[1])] for p in boundary]
|
| 335 |
+
h3_cell = cell_id
|
| 336 |
+
except Exception as e:
|
| 337 |
+
print(f"H3 calculation failed: {e}")
|
| 338 |
+
|
| 339 |
+
# Resolve static URLs using preloaded gallery_path
|
| 340 |
+
gallery_url = "/" + meta.get("gallery_path", "")
|
| 341 |
+
if not gallery_url.startswith("/"):
|
| 342 |
+
gallery_url = "/" + gallery_url
|
| 343 |
+
|
| 344 |
+
out_results.append({
|
| 345 |
+
"index": int(meta["index"]),
|
| 346 |
+
"class": meta["class"],
|
| 347 |
+
"modality": meta["modality"],
|
| 348 |
+
"original_path": meta["original_path"],
|
| 349 |
+
"gallery_path": gallery_url,
|
| 350 |
+
"lat": meta.get("lat"),
|
| 351 |
+
"lon": meta.get("lon"),
|
| 352 |
+
"distance_km": dist_km,
|
| 353 |
+
"h3_cell": h3_cell,
|
| 354 |
+
"h3_boundary": h3_boundary,
|
| 355 |
+
"score": float(score)
|
| 356 |
+
})
|
| 357 |
+
|
| 358 |
+
# Sort and slice to requested count
|
| 359 |
+
out_results = sorted(out_results, key=lambda x: x["score"], reverse=True)[:k]
|
| 360 |
+
return out_results
|
| 361 |
+
|
| 362 |
+
@app.post("/api/search")
|
| 363 |
+
async def post_search(
|
| 364 |
+
file: UploadFile = File(...),
|
| 365 |
+
k: int = Form(5),
|
| 366 |
+
level: str = Form("level4"),
|
| 367 |
+
query_modality: str = Form("optical"),
|
| 368 |
+
lat: Optional[float] = Form(None),
|
| 369 |
+
lon: Optional[float] = Form(None),
|
| 370 |
+
radius_km: Optional[float] = Form(50.0)
|
| 371 |
+
):
|
| 372 |
+
"""Main image query search endpoint."""
|
| 373 |
+
t0 = time.time()
|
| 374 |
+
|
| 375 |
+
# Save uploaded file temporarily
|
| 376 |
+
temp_dir = Path("data/temp")
|
| 377 |
+
temp_dir.mkdir(parents=True, exist_ok=True)
|
| 378 |
+
temp_path = temp_dir / file.filename
|
| 379 |
+
|
| 380 |
+
try:
|
| 381 |
+
with open(temp_path, "wb") as f:
|
| 382 |
+
f.write(await file.read())
|
| 383 |
+
|
| 384 |
+
# Out-of-core TIFF loading and pre-processing
|
| 385 |
+
arr = load_tiff_downsampled(temp_path)
|
| 386 |
+
tensor = torch.from_numpy(arr).float()
|
| 387 |
+
|
| 388 |
+
# Scale range
|
| 389 |
+
if tensor.max() > 1.0:
|
| 390 |
+
tensor = tensor / 255.0
|
| 391 |
+
|
| 392 |
+
# Standardize format to channels-first (C, H, W)
|
| 393 |
+
if tensor.ndim == 3:
|
| 394 |
+
if tensor.shape[-1] in [2, 3, 4, 13]:
|
| 395 |
+
tensor = tensor.permute(2, 0, 1)
|
| 396 |
+
elif tensor.ndim == 2:
|
| 397 |
+
tensor = tensor.unsqueeze(0)
|
| 398 |
+
|
| 399 |
+
# Resize to exactly 224x224 for SatCLIP model compatibility
|
| 400 |
+
if tensor.shape[1] != 224 or tensor.shape[2] != 224:
|
| 401 |
+
tensor = torch.nn.functional.interpolate(
|
| 402 |
+
tensor.unsqueeze(0), size=(224, 224),
|
| 403 |
+
mode="bilinear", align_corners=False
|
| 404 |
+
).squeeze(0)
|
| 405 |
+
|
| 406 |
+
# Extract features using SatCLIP encoder
|
| 407 |
+
with torch.no_grad():
|
| 408 |
+
query_emb = extractor.extract_features_from_tensor(
|
| 409 |
+
tensor, modality=query_modality, normalize=True
|
| 410 |
+
).cpu().numpy()
|
| 411 |
+
|
| 412 |
+
# Execute query search
|
| 413 |
+
results = perform_engine_search(
|
| 414 |
+
query_emb=query_emb,
|
| 415 |
+
query_modality=query_modality,
|
| 416 |
+
k=k,
|
| 417 |
+
level=level,
|
| 418 |
+
lat=lat,
|
| 419 |
+
lon=lon,
|
| 420 |
+
radius_km=radius_km
|
| 421 |
+
)
|
| 422 |
+
|
| 423 |
+
query_time = (time.time() - t0) * 1000
|
| 424 |
+
return {
|
| 425 |
+
"query_time_ms": query_time,
|
| 426 |
+
"device": extractor.device,
|
| 427 |
+
"results": results
|
| 428 |
+
}
|
| 429 |
+
except Exception as e:
|
| 430 |
+
import traceback
|
| 431 |
+
traceback.print_exc()
|
| 432 |
+
raise HTTPException(status_code=500, detail=f"Retrieval execution failed: {str(e)}")
|
| 433 |
+
finally:
|
| 434 |
+
if temp_path.exists():
|
| 435 |
+
temp_path.unlink()
|
| 436 |
+
|
| 437 |
+
@app.post("/api/search-text")
|
| 438 |
+
async def post_search_text(
|
| 439 |
+
text_query: str = Form(...),
|
| 440 |
+
k: int = Form(5),
|
| 441 |
+
level: str = Form("level4"),
|
| 442 |
+
query_modality: str = Form("optical"),
|
| 443 |
+
lat: Optional[float] = Form(None),
|
| 444 |
+
lon: Optional[float] = Form(None),
|
| 445 |
+
radius_km: Optional[float] = Form(50.0)
|
| 446 |
+
):
|
| 447 |
+
"""Text-to-Image text query search endpoint using OpenAI CLIP text encoder."""
|
| 448 |
+
t0 = time.time()
|
| 449 |
+
try:
|
| 450 |
+
# Load OpenAI CLIP ViT-L/14 model weights
|
| 451 |
+
device = extractor.device
|
| 452 |
+
clip_model, _ = clip.load("ViT-L/14", device=device)
|
| 453 |
+
|
| 454 |
+
# Tokenize text
|
| 455 |
+
text_tokens = clip.tokenize([text_query]).to(device)
|
| 456 |
+
with torch.no_grad():
|
| 457 |
+
text_emb = clip_model.encode_text(text_tokens)
|
| 458 |
+
text_emb = text_emb / text_emb.norm(dim=-1, keepdim=True)
|
| 459 |
+
query_emb = text_emb.cpu().numpy()[0]
|
| 460 |
+
|
| 461 |
+
# Execute query search
|
| 462 |
+
results = perform_engine_search(
|
| 463 |
+
query_emb=query_emb,
|
| 464 |
+
query_modality=query_modality,
|
| 465 |
+
k=k,
|
| 466 |
+
level=level,
|
| 467 |
+
lat=lat,
|
| 468 |
+
lon=lon,
|
| 469 |
+
radius_km=radius_km
|
| 470 |
+
)
|
| 471 |
+
|
| 472 |
+
query_time = (time.time() - t0) * 1000
|
| 473 |
+
return {
|
| 474 |
+
"query_time_ms": query_time,
|
| 475 |
+
"device": device,
|
| 476 |
+
"results": results
|
| 477 |
+
}
|
| 478 |
+
except Exception as e:
|
| 479 |
+
raise HTTPException(status_code=500, detail=f"Text search failed: {str(e)}")
|
| 480 |
+
|
| 481 |
+
# ---------------------------------------------------------------------------
|
| 482 |
+
# Initializers & Fallback Demo Creators
|
| 483 |
+
# ---------------------------------------------------------------------------
|
| 484 |
+
def build_demo_index_fallback():
|
| 485 |
+
"""Build a mock database fallback in case the main EuroSAT database is missing or build is pending."""
|
| 486 |
+
print("Warning: Building demo fallback indices...")
|
| 487 |
+
N_GALLERY = 100
|
| 488 |
+
EMBED_DIM = 768
|
| 489 |
+
|
| 490 |
+
# Generate mock metadata
|
| 491 |
+
mock_meta = []
|
| 492 |
+
classes = ["AnnualCrop", "Forest", "HerbaceousVegetation", "Highway", "Industrial",
|
| 493 |
+
"Pasture", "PermanentCrop", "Residential", "River", "SeaLake"]
|
| 494 |
+
|
| 495 |
+
for i in range(N_GALLERY * 3):
|
| 496 |
+
mod = "optical" if i < N_GALLERY else ("sar" if i < N_GALLERY * 2 else "multispectral")
|
| 497 |
+
cls = classes[i % len(classes)]
|
| 498 |
+
|
| 499 |
+
# Bengaluru coordinates
|
| 500 |
+
lat = 12.9716 + random.uniform(-0.35, 0.35)
|
| 501 |
+
lon = 77.5946 + random.uniform(-0.35, 0.35)
|
| 502 |
+
|
| 503 |
+
# Create folder & write dummy file if not exists
|
| 504 |
+
mod_dir = GALLERY_DIR / mod / cls
|
| 505 |
+
mod_dir.mkdir(parents=True, exist_ok=True)
|
| 506 |
+
img_path = mod_dir / f"{cls}_{i}.png"
|
| 507 |
+
|
| 508 |
+
if not img_path.exists():
|
| 509 |
+
arr = np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8)
|
| 510 |
+
Image.fromarray(arr).save(img_path)
|
| 511 |
+
|
| 512 |
+
mock_meta.append({
|
| 513 |
+
"index": i,
|
| 514 |
+
"class": cls,
|
| 515 |
+
"modality": mod,
|
| 516 |
+
"original_path": str(img_path),
|
| 517 |
+
"lat": lat,
|
| 518 |
+
"lon": lon
|
| 519 |
+
})
|
| 520 |
+
|
| 521 |
+
# Generate mock embeddings
|
| 522 |
+
mock_embs = {}
|
| 523 |
+
for mod in ["optical", "sar", "multispectral"]:
|
| 524 |
+
emb = np.random.randn(N_GALLERY, EMBED_DIM).astype(np.float32)
|
| 525 |
+
# Normalize
|
| 526 |
+
norms = np.linalg.norm(emb, axis=1, keepdims=True)
|
| 527 |
+
mock_embs[mod] = emb / (norms + 1e-8)
|
| 528 |
+
|
| 529 |
+
meta_by_mod = {
|
| 530 |
+
"optical": mock_meta[:N_GALLERY],
|
| 531 |
+
"sar": mock_meta[N_GALLERY:N_GALLERY*2],
|
| 532 |
+
"multispectral": mock_meta[N_GALLERY*2:]
|
| 533 |
+
}
|
| 534 |
+
|
| 535 |
+
engine = CrossModalRetrieval(embed_dim=EMBED_DIM)
|
| 536 |
+
engine.build_multi_index(mock_embs, meta_by_mod, use_centering=True)
|
| 537 |
+
engine.build_spatial_index(mock_meta)
|
| 538 |
+
|
| 539 |
+
return engine, mock_meta
|
| 540 |
+
|
| 541 |
+
def start_server_assets():
|
| 542 |
+
"""Load SatCLIP models and verify database paths."""
|
| 543 |
+
global extractor, retrieval, metadata_db
|
| 544 |
+
|
| 545 |
+
print("Loading SatCLIP Vision & Text extractors...")
|
| 546 |
+
extractor = FeatureExtractor()
|
| 547 |
+
|
| 548 |
+
index_path = PROCESSED_DIR / "metadata.json"
|
| 549 |
+
embed_path = PROCESSED_DIR / "gallery_embeddings.pt"
|
| 550 |
+
meta_path = PROCESSED_DIR / "gallery_metadata.json"
|
| 551 |
+
|
| 552 |
+
# Try loading pre-built FAISS indices
|
| 553 |
+
if index_path.exists():
|
| 554 |
+
print("Loading pre-built FAISS multi-index cache...")
|
| 555 |
+
retrieval = CrossModalRetrieval(embed_dim=768)
|
| 556 |
+
retrieval.load(PROCESSED_DIR)
|
| 557 |
+
metadata_db = retrieval.metadata
|
| 558 |
+
# Re-build the spatial grid index in RAM
|
| 559 |
+
retrieval.build_spatial_index(metadata_db)
|
| 560 |
+
print(f"Loaded indices successfully: {len(metadata_db)} vectors loaded.")
|
| 561 |
+
# Else try building in memory from the raw PyTorch embeddings file
|
| 562 |
+
elif embed_path.exists() and meta_path.exists():
|
| 563 |
+
print("Building multi-index from raw torch embeddings...")
|
| 564 |
+
with open(meta_path) as f:
|
| 565 |
+
metadata_db = json.load(f)
|
| 566 |
+
|
| 567 |
+
embeddings = torch.load(embed_path, map_location="cpu")
|
| 568 |
+
embeddings_np = embeddings.numpy().astype(np.float32)
|
| 569 |
+
|
| 570 |
+
# Split by modality
|
| 571 |
+
embeddings_by_mod = {}
|
| 572 |
+
metadata_by_mod = {}
|
| 573 |
+
for entry in metadata_db:
|
| 574 |
+
mod = entry["modality"]
|
| 575 |
+
if mod not in embeddings_by_mod:
|
| 576 |
+
embeddings_by_mod[mod] = []
|
| 577 |
+
metadata_by_mod[mod] = []
|
| 578 |
+
embeddings_by_mod[mod].append(embeddings_np[entry["index"]])
|
| 579 |
+
metadata_by_mod[mod].append(entry)
|
| 580 |
+
|
| 581 |
+
for mod in embeddings_by_mod:
|
| 582 |
+
embeddings_by_mod[mod] = np.array(embeddings_by_mod[mod])
|
| 583 |
+
|
| 584 |
+
retrieval = CrossModalRetrieval(embed_dim=768)
|
| 585 |
+
retrieval.build_multi_index(embeddings_by_mod, metadata_by_mod, use_centering=True)
|
| 586 |
+
retrieval.build_spatial_index(metadata_db)
|
| 587 |
+
print(f"Built index in memory successfully: {len(metadata_db)} vectors loaded.")
|
| 588 |
+
# Fallback to random demo database
|
| 589 |
+
else:
|
| 590 |
+
retrieval, metadata_db = build_demo_index_fallback()
|
| 591 |
+
|
| 592 |
+
# Initialize assets
|
| 593 |
+
start_server_assets()
|
| 594 |
+
|
| 595 |
+
# Remove Gradio's default '/' route to prevent shadowing our custom static index.html
|
| 596 |
+
app.routes[:] = [r for r in app.routes if getattr(r, "path", None) != "/"]
|
| 597 |
+
|
| 598 |
+
# Serve database images statically
|
| 599 |
+
app.mount("/data/gallery", StaticFiles(directory="data/gallery"), name="gallery")
|
| 600 |
+
|
| 601 |
+
# Serve index.html explicitly at root with no-cache headers to prevent browser caching
|
| 602 |
+
@app.get("/")
|
| 603 |
+
def read_root():
|
| 604 |
+
headers = {
|
| 605 |
+
"Cache-Control": "no-store, no-cache, must-revalidate, max-age=0",
|
| 606 |
+
"Pragma": "no-cache",
|
| 607 |
+
"Expires": "0"
|
| 608 |
+
}
|
| 609 |
+
return FileResponse("src/ui/static/index.html", headers=headers)
|
| 610 |
+
|
| 611 |
+
# Serve the static UI files at root
|
| 612 |
+
app.mount("/", StaticFiles(directory="src/ui/static", html=True), name="static")
|
| 613 |
+
|
| 614 |
+
if __name__ == "__main__":
|
| 615 |
+
import uvicorn
|
| 616 |
+
uvicorn.run(app, host="0.0.0.0", port=7860)
|
benchmark_cross_modal.py
ADDED
|
@@ -0,0 +1,359 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Benchmark cross-modal retrieval approaches.
|
| 3 |
+
|
| 4 |
+
Tests:
|
| 5 |
+
1. Single-index with modality filtering
|
| 6 |
+
2. Multi-index search
|
| 7 |
+
3. Hybrid search
|
| 8 |
+
4. Cross-modal alignment with projection heads
|
| 9 |
+
"""
|
| 10 |
+
|
| 11 |
+
import torch
|
| 12 |
+
import numpy as np
|
| 13 |
+
import time
|
| 14 |
+
from pathlib import Path
|
| 15 |
+
from typing import Dict, List, Tuple
|
| 16 |
+
|
| 17 |
+
# Add src to path
|
| 18 |
+
import sys
|
| 19 |
+
sys.path.insert(0, str(Path(__file__).parent))
|
| 20 |
+
|
| 21 |
+
from src.features.cross_modal import CrossModalAligner, CrossModalConfig
|
| 22 |
+
from src.retrieval.cross_modal_retrieval import CrossModalRetrieval
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
def load_gallery_data():
|
| 26 |
+
"""Load real gallery embeddings and metadata."""
|
| 27 |
+
import json
|
| 28 |
+
|
| 29 |
+
data_dir = Path("data/processed")
|
| 30 |
+
|
| 31 |
+
# Load embeddings
|
| 32 |
+
embeddings = torch.load(data_dir / "gallery_embeddings.pt", weights_only=True)
|
| 33 |
+
|
| 34 |
+
# Load metadata
|
| 35 |
+
with open(data_dir / "gallery_metadata.json") as f:
|
| 36 |
+
metadata = json.load(f)
|
| 37 |
+
|
| 38 |
+
return embeddings.numpy().astype(np.float32), metadata
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
def split_by_modality(embeddings: np.ndarray, metadata: List[dict]) -> Dict[str, np.ndarray]:
|
| 42 |
+
"""Split embeddings by modality."""
|
| 43 |
+
modalities = {}
|
| 44 |
+
for i, entry in enumerate(metadata):
|
| 45 |
+
mod = entry["modality"]
|
| 46 |
+
if mod not in modalities:
|
| 47 |
+
modalities[mod] = []
|
| 48 |
+
modalities[mod].append(embeddings[i])
|
| 49 |
+
|
| 50 |
+
return {mod: np.array(embs) for mod, embs in modalities.items()}
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
def compute_recall_at_k(retrieved_indices: List[int], ground_truth_idx: int, k: int) -> float:
|
| 54 |
+
"""Compute Recall@K."""
|
| 55 |
+
if ground_truth_idx in retrieved_indices[:k]:
|
| 56 |
+
return 1.0
|
| 57 |
+
return 0.0
|
| 58 |
+
|
| 59 |
+
|
| 60 |
+
def benchmark_single_index(
|
| 61 |
+
embeddings: np.ndarray,
|
| 62 |
+
metadata: List[dict],
|
| 63 |
+
n_queries: int = 50,
|
| 64 |
+
k: int = 5
|
| 65 |
+
) -> Dict:
|
| 66 |
+
"""Benchmark single-index approach."""
|
| 67 |
+
print("\n=== Single-Index Benchmark ===")
|
| 68 |
+
|
| 69 |
+
retrieval = CrossModalRetrieval(embed_dim=embeddings.shape[1])
|
| 70 |
+
retrieval.build_single_index(embeddings, [m["modality"] for m in metadata], metadata)
|
| 71 |
+
|
| 72 |
+
# Generate queries (use gallery images as queries)
|
| 73 |
+
query_indices = np.random.choice(len(embeddings), n_queries, replace=False)
|
| 74 |
+
|
| 75 |
+
results = {
|
| 76 |
+
"same_modal": [],
|
| 77 |
+
"cross_modal": [],
|
| 78 |
+
}
|
| 79 |
+
|
| 80 |
+
for idx in query_indices:
|
| 81 |
+
query = embeddings[idx:idx+1]
|
| 82 |
+
query_mod = metadata[idx]["modality"]
|
| 83 |
+
query_class = metadata[idx]["class"]
|
| 84 |
+
|
| 85 |
+
# Same-modal search
|
| 86 |
+
same_result = retrieval.search(query, query_mod, target_modality=query_mod, k=k)
|
| 87 |
+
same_recall = compute_recall_at_k(
|
| 88 |
+
[metadata[i]["class"] for i in same_result.indices],
|
| 89 |
+
query_class,
|
| 90 |
+
k
|
| 91 |
+
)
|
| 92 |
+
results["same_modal"].append(same_recall)
|
| 93 |
+
|
| 94 |
+
# Cross-modal search (find different modality with same class)
|
| 95 |
+
cross_targets = [m for m in ["optical", "sar", "multispectral"] if m != query_mod]
|
| 96 |
+
for target_mod in cross_targets:
|
| 97 |
+
cross_result = retrieval.search(query, query_mod, target_mod, k=k)
|
| 98 |
+
cross_recall = compute_recall_at_k(
|
| 99 |
+
[metadata[i]["class"] for i in cross_result.indices],
|
| 100 |
+
query_class,
|
| 101 |
+
k
|
| 102 |
+
)
|
| 103 |
+
results["cross_modal"].append(cross_recall)
|
| 104 |
+
|
| 105 |
+
return {
|
| 106 |
+
"same_modal_recall@k": np.mean(results["same_modal"]),
|
| 107 |
+
"cross_modal_recall@k": np.mean(results["cross_modal"]),
|
| 108 |
+
}
|
| 109 |
+
|
| 110 |
+
|
| 111 |
+
def benchmark_multi_index(
|
| 112 |
+
embeddings_by_mod: Dict[str, np.ndarray],
|
| 113 |
+
metadata: List[dict],
|
| 114 |
+
n_queries: int = 50,
|
| 115 |
+
k: int = 5
|
| 116 |
+
) -> Dict:
|
| 117 |
+
"""Benchmark multi-index approach."""
|
| 118 |
+
print("\n=== Multi-Index Benchmark ===")
|
| 119 |
+
|
| 120 |
+
# Build metadata by modality
|
| 121 |
+
metadata_by_mod = {}
|
| 122 |
+
for entry in metadata:
|
| 123 |
+
mod = entry["modality"]
|
| 124 |
+
if mod not in metadata_by_mod:
|
| 125 |
+
metadata_by_mod[mod] = []
|
| 126 |
+
metadata_by_mod[mod].append(entry)
|
| 127 |
+
|
| 128 |
+
retrieval = CrossModalRetrieval(embed_dim=768)
|
| 129 |
+
retrieval.build_multi_index(embeddings_by_mod, metadata_by_mod)
|
| 130 |
+
|
| 131 |
+
results = {
|
| 132 |
+
"same_modal": [],
|
| 133 |
+
"cross_modal": [],
|
| 134 |
+
}
|
| 135 |
+
|
| 136 |
+
# Generate queries
|
| 137 |
+
all_embeddings = np.concatenate(list(embeddings_by_mod.values()))
|
| 138 |
+
query_indices = np.random.choice(len(all_embeddings), n_queries, replace=False)
|
| 139 |
+
|
| 140 |
+
for idx in query_indices:
|
| 141 |
+
query = all_embeddings[idx:idx+1]
|
| 142 |
+
|
| 143 |
+
# Determine query modality
|
| 144 |
+
offset = 0
|
| 145 |
+
query_mod = None
|
| 146 |
+
for mod, embs in embeddings_by_mod.items():
|
| 147 |
+
if idx < offset + len(embs):
|
| 148 |
+
query_mod = mod
|
| 149 |
+
break
|
| 150 |
+
offset += len(embs)
|
| 151 |
+
|
| 152 |
+
if query_mod is None:
|
| 153 |
+
continue
|
| 154 |
+
|
| 155 |
+
# Same-modal search
|
| 156 |
+
same_result = retrieval.search(query, query_mod, target_modality=query_mod, k=k)
|
| 157 |
+
same_recall = 1.0 if any(True for _ in same_result.indices) else 0.0
|
| 158 |
+
results["same_modal"].append(same_recall)
|
| 159 |
+
|
| 160 |
+
# Cross-modal search
|
| 161 |
+
cross_targets = [m for m in embeddings_by_mod.keys() if m != query_mod]
|
| 162 |
+
cross_result = retrieval.search(query, query_mod, k=k)
|
| 163 |
+
cross_recall = 1.0 if any(True for _ in cross_result.indices) else 0.0
|
| 164 |
+
results["cross_modal"].append(cross_recall)
|
| 165 |
+
|
| 166 |
+
return {
|
| 167 |
+
"same_modal_recall@k": np.mean(results["same_modal"]),
|
| 168 |
+
"cross_modal_recall@k": np.mean(results["cross_modal"]),
|
| 169 |
+
}
|
| 170 |
+
|
| 171 |
+
|
| 172 |
+
def benchmark_hybrid(
|
| 173 |
+
embeddings_by_mod: Dict[str, np.ndarray],
|
| 174 |
+
metadata: List[dict],
|
| 175 |
+
n_queries: int = 50,
|
| 176 |
+
k: int = 5
|
| 177 |
+
) -> Dict:
|
| 178 |
+
"""Benchmark hybrid search approach."""
|
| 179 |
+
print("\n=== Hybrid Search Benchmark ===")
|
| 180 |
+
|
| 181 |
+
metadata_by_mod = {}
|
| 182 |
+
for entry in metadata:
|
| 183 |
+
mod = entry["modality"]
|
| 184 |
+
if mod not in metadata_by_mod:
|
| 185 |
+
metadata_by_mod[mod] = []
|
| 186 |
+
metadata_by_mod[mod].append(entry)
|
| 187 |
+
|
| 188 |
+
retrieval = CrossModalRetrieval(embed_dim=768)
|
| 189 |
+
retrieval.build_multi_index(embeddings_by_mod, metadata_by_mod)
|
| 190 |
+
|
| 191 |
+
results = []
|
| 192 |
+
|
| 193 |
+
all_embeddings = np.concatenate(list(embeddings_by_mod.values()))
|
| 194 |
+
query_indices = np.random.choice(len(all_embeddings), n_queries, replace=False)
|
| 195 |
+
|
| 196 |
+
for idx in query_indices:
|
| 197 |
+
query = all_embeddings[idx:idx+1]
|
| 198 |
+
|
| 199 |
+
offset = 0
|
| 200 |
+
query_mod = None
|
| 201 |
+
for mod, embs in embeddings_by_mod.items():
|
| 202 |
+
if idx < offset + len(embs):
|
| 203 |
+
query_mod = mod
|
| 204 |
+
break
|
| 205 |
+
offset += len(embs)
|
| 206 |
+
|
| 207 |
+
if query_mod is None:
|
| 208 |
+
continue
|
| 209 |
+
|
| 210 |
+
result = retrieval.search_hybrid(query, query_mod, k=k)
|
| 211 |
+
results.append(1.0 if len(result.indices) > 0 else 0.0)
|
| 212 |
+
|
| 213 |
+
return {
|
| 214 |
+
"hybrid_recall@k": np.mean(results),
|
| 215 |
+
}
|
| 216 |
+
|
| 217 |
+
|
| 218 |
+
def benchmark_cross_modal_alignment(
|
| 219 |
+
embeddings: np.ndarray,
|
| 220 |
+
metadata: List[dict],
|
| 221 |
+
n_queries: int = 50,
|
| 222 |
+
k: int = 5
|
| 223 |
+
) -> Dict:
|
| 224 |
+
"""Benchmark cross-modal alignment with projection heads."""
|
| 225 |
+
print("\n=== Cross-Modal Alignment Benchmark ===")
|
| 226 |
+
|
| 227 |
+
config = CrossModalConfig(
|
| 228 |
+
embed_dim=embeddings.shape[1],
|
| 229 |
+
projection_dim=256,
|
| 230 |
+
use_wavelength_encoding=True,
|
| 231 |
+
use_domain_adaptation=True,
|
| 232 |
+
)
|
| 233 |
+
|
| 234 |
+
aligner = CrossModalAligner(config)
|
| 235 |
+
|
| 236 |
+
# Project all embeddings
|
| 237 |
+
projected = {}
|
| 238 |
+
for mod in ["optical", "sar", "multispectral"]:
|
| 239 |
+
mask = [m["modality"] == mod for m in metadata]
|
| 240 |
+
mod_embeddings = embeddings[mask]
|
| 241 |
+
projected[mod] = aligner.project(
|
| 242 |
+
torch.tensor(mod_embeddings), mod
|
| 243 |
+
).detach().numpy()
|
| 244 |
+
|
| 245 |
+
results = {
|
| 246 |
+
"same_modal": [],
|
| 247 |
+
"cross_modal": [],
|
| 248 |
+
}
|
| 249 |
+
|
| 250 |
+
query_indices = np.random.choice(len(embeddings), n_queries, replace=False)
|
| 251 |
+
|
| 252 |
+
for idx in query_indices:
|
| 253 |
+
query = embeddings[idx:idx+1]
|
| 254 |
+
query_mod = metadata[idx]["modality"]
|
| 255 |
+
query_class = metadata[idx]["class"]
|
| 256 |
+
|
| 257 |
+
# Project query
|
| 258 |
+
query_proj = aligner.project(
|
| 259 |
+
torch.tensor(query), query_mod
|
| 260 |
+
).detach().numpy()
|
| 261 |
+
|
| 262 |
+
# Same-modal search
|
| 263 |
+
same_proj = projected[query_mod]
|
| 264 |
+
similarities = query_proj @ same_proj.T
|
| 265 |
+
topk_idx = np.argsort(similarities[0])[::-1][:k]
|
| 266 |
+
|
| 267 |
+
same_recall = compute_recall_at_k(
|
| 268 |
+
[metadata[i]["class"] for i in topk_idx],
|
| 269 |
+
query_class,
|
| 270 |
+
k
|
| 271 |
+
)
|
| 272 |
+
results["same_modal"].append(same_recall)
|
| 273 |
+
|
| 274 |
+
# Cross-modal search
|
| 275 |
+
for target_mod in ["optical", "sar", "multispectral"]:
|
| 276 |
+
if target_mod == query_mod:
|
| 277 |
+
continue
|
| 278 |
+
|
| 279 |
+
target_proj = projected[target_mod]
|
| 280 |
+
similarities = query_proj @ target_proj.T
|
| 281 |
+
topk_idx = np.argsort(similarities[0])[::-1][:k]
|
| 282 |
+
|
| 283 |
+
cross_recall = compute_recall_at_k(
|
| 284 |
+
[metadata[i]["class"] for i in topk_idx],
|
| 285 |
+
query_class,
|
| 286 |
+
k
|
| 287 |
+
)
|
| 288 |
+
results["cross_modal"].append(cross_recall)
|
| 289 |
+
|
| 290 |
+
return {
|
| 291 |
+
"same_modal_recall@k": np.mean(results["same_modal"]),
|
| 292 |
+
"cross_modal_recall@k": np.mean(results["cross_modal"]),
|
| 293 |
+
}
|
| 294 |
+
|
| 295 |
+
|
| 296 |
+
def main():
|
| 297 |
+
"""Run all benchmarks."""
|
| 298 |
+
print("=" * 60)
|
| 299 |
+
print("Cross-Modal Retrieval Benchmark")
|
| 300 |
+
print("=" * 60)
|
| 301 |
+
|
| 302 |
+
# Load data
|
| 303 |
+
print("\nLoading gallery data...")
|
| 304 |
+
embeddings, metadata = load_gallery_data()
|
| 305 |
+
print(f"Loaded {len(metadata)} embeddings of dimension {embeddings.shape[1]}")
|
| 306 |
+
|
| 307 |
+
# Split by modality
|
| 308 |
+
embeddings_by_mod = split_by_modality(embeddings, metadata)
|
| 309 |
+
print(f"Modalities: {list(embeddings_by_mod.keys())}")
|
| 310 |
+
for mod, embs in embeddings_by_mod.items():
|
| 311 |
+
print(f" {mod}: {len(embs)} images")
|
| 312 |
+
|
| 313 |
+
# Run benchmarks
|
| 314 |
+
n_queries = min(50, len(metadata))
|
| 315 |
+
k = 5
|
| 316 |
+
|
| 317 |
+
results = {}
|
| 318 |
+
|
| 319 |
+
# 1. Single-index
|
| 320 |
+
t0 = time.time()
|
| 321 |
+
results["single"] = benchmark_single_index(embeddings, metadata, n_queries, k)
|
| 322 |
+
results["single"]["time"] = time.time() - t0
|
| 323 |
+
|
| 324 |
+
# 2. Multi-index
|
| 325 |
+
t0 = time.time()
|
| 326 |
+
results["multi"] = benchmark_multi_index(embeddings_by_mod, metadata, n_queries, k)
|
| 327 |
+
results["multi"]["time"] = time.time() - t0
|
| 328 |
+
|
| 329 |
+
# 3. Hybrid
|
| 330 |
+
t0 = time.time()
|
| 331 |
+
results["hybrid"] = benchmark_hybrid(embeddings_by_mod, metadata, n_queries, k)
|
| 332 |
+
results["hybrid"]["time"] = time.time() - t0
|
| 333 |
+
|
| 334 |
+
# 4. Cross-modal alignment
|
| 335 |
+
t0 = time.time()
|
| 336 |
+
results["alignment"] = benchmark_cross_modal_alignment(embeddings, metadata, n_queries, k)
|
| 337 |
+
results["alignment"]["time"] = time.time() - t0
|
| 338 |
+
|
| 339 |
+
# Print results
|
| 340 |
+
print("\n" + "=" * 60)
|
| 341 |
+
print("Results Summary")
|
| 342 |
+
print("=" * 60)
|
| 343 |
+
|
| 344 |
+
print(f"\n{'Method':<20} {'Same-Modal R@5':<18} {'Cross-Modal R@5':<18} {'Time (s)':<10}")
|
| 345 |
+
print("-" * 66)
|
| 346 |
+
|
| 347 |
+
for method, res in results.items():
|
| 348 |
+
same = res.get("same_modal_recall@k", 0)
|
| 349 |
+
cross = res.get("cross_modal_recall@k", 0) or res.get("hybrid_recall@k", 0)
|
| 350 |
+
t = res.get("time", 0)
|
| 351 |
+
print(f"{method:<20} {same:<18.4f} {cross:<18.4f} {t:<10.3f}")
|
| 352 |
+
|
| 353 |
+
print("\n" + "=" * 60)
|
| 354 |
+
print("Recommendation: Use the method with highest cross-modal recall")
|
| 355 |
+
print("=" * 60)
|
| 356 |
+
|
| 357 |
+
|
| 358 |
+
if __name__ == "__main__":
|
| 359 |
+
main()
|
experiment_comparison.py
ADDED
|
@@ -0,0 +1,263 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Experiment: DINOv2-CLIP Hybrid vs Pure CLIP
|
| 3 |
+
|
| 4 |
+
Compares 4 approaches on the pre-computed EuroSAT gallery:
|
| 5 |
+
1. Pure CLIP (baseline)
|
| 6 |
+
2. CLIP + SAR Adapter
|
| 7 |
+
3. CLIP + DINOv2 patch features (hybrid)
|
| 8 |
+
4. Full hybrid (CLIP + SAR adapter + DINOv2)
|
| 9 |
+
|
| 10 |
+
Metrics: same-modal and cross-modal Recall@K, latency.
|
| 11 |
+
"""
|
| 12 |
+
|
| 13 |
+
import sys, time, json, traceback
|
| 14 |
+
import torch, numpy as np
|
| 15 |
+
from pathlib import Path
|
| 16 |
+
from PIL import Image
|
| 17 |
+
from dataclasses import dataclass, asdict
|
| 18 |
+
from typing import List, Optional
|
| 19 |
+
|
| 20 |
+
sys.path.insert(0, str(Path(__file__).parent))
|
| 21 |
+
|
| 22 |
+
DATA_DIR = Path("data")
|
| 23 |
+
PROCESSED_DIR = DATA_DIR / "processed"
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
@dataclass
|
| 27 |
+
class Result:
|
| 28 |
+
model: str
|
| 29 |
+
same_r1: float
|
| 30 |
+
same_r5: float
|
| 31 |
+
same_r10: float
|
| 32 |
+
cross_r1: float
|
| 33 |
+
cross_r5: float
|
| 34 |
+
cross_r10: float
|
| 35 |
+
latency_ms: float
|
| 36 |
+
n_queries: int
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
def load_data():
|
| 40 |
+
embeddings = torch.load(PROCESSED_DIR / "gallery_embeddings.pt", weights_only=True)
|
| 41 |
+
with open(PROCESSED_DIR / "gallery_metadata.json") as f:
|
| 42 |
+
metadata = json.load(f)
|
| 43 |
+
return embeddings.numpy().astype(np.float32), metadata
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
def split(metadata):
|
| 47 |
+
"""Stratified 30/70 split: 30% of each (modality, class) pair goes to queries."""
|
| 48 |
+
groups = {}
|
| 49 |
+
for e in metadata:
|
| 50 |
+
key = (e["modality"], e["class"])
|
| 51 |
+
groups.setdefault(key, []).append(e)
|
| 52 |
+
|
| 53 |
+
queries, gallery = [], []
|
| 54 |
+
for key, entries in groups.items():
|
| 55 |
+
n = max(1, int(len(entries) * 0.3))
|
| 56 |
+
queries.extend(entries[:n])
|
| 57 |
+
gallery.extend(entries[n:])
|
| 58 |
+
return queries, gallery
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
def recall_at_k(retrieved, query_mod, query_class, metadata, k, mode="same"):
|
| 62 |
+
hits = 0
|
| 63 |
+
for idx in retrieved[:k]:
|
| 64 |
+
m = metadata[idx]
|
| 65 |
+
same_class = m["class"] == query_class
|
| 66 |
+
same_mod = m["modality"] == query_mod
|
| 67 |
+
if mode == "same" and same_class and same_mod:
|
| 68 |
+
hits += 1
|
| 69 |
+
elif mode == "cross" and same_class and not same_mod:
|
| 70 |
+
hits += 1
|
| 71 |
+
return hits
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
def evaluate(queries, all_emb, metadata, gallery_entries, extractor_fn, label):
|
| 75 |
+
import faiss
|
| 76 |
+
|
| 77 |
+
gal_idx = [e["index"] for e in gallery_entries]
|
| 78 |
+
gal_emb = all_emb[gal_idx]
|
| 79 |
+
dim = gal_emb.shape[1]
|
| 80 |
+
index = faiss.IndexFlatIP(dim)
|
| 81 |
+
index.add(gal_emb)
|
| 82 |
+
|
| 83 |
+
sr1, sr5, sr10 = [], [], []
|
| 84 |
+
cr1, cr5, cr10 = [], [], []
|
| 85 |
+
latencies = []
|
| 86 |
+
|
| 87 |
+
for q in queries:
|
| 88 |
+
q_path = Path(q["gallery_path"])
|
| 89 |
+
if not q_path.exists():
|
| 90 |
+
continue
|
| 91 |
+
img = Image.open(q_path).convert("RGB")
|
| 92 |
+
|
| 93 |
+
start = time.perf_counter()
|
| 94 |
+
try:
|
| 95 |
+
emb = extractor_fn(img, q["modality"])
|
| 96 |
+
except Exception:
|
| 97 |
+
continue
|
| 98 |
+
elapsed = (time.perf_counter() - start) * 1000
|
| 99 |
+
latencies.append(elapsed)
|
| 100 |
+
|
| 101 |
+
q_np = emb.reshape(1, -1).astype(np.float32)
|
| 102 |
+
_, ids = index.search(q_np, 10)
|
| 103 |
+
retrieved = [gal_idx[i] for i in ids[0] if 0 <= i < len(gal_idx)]
|
| 104 |
+
|
| 105 |
+
sr1.append(recall_at_k(retrieved, q["modality"], q["class"], metadata, 1, "same"))
|
| 106 |
+
sr5.append(recall_at_k(retrieved, q["modality"], q["class"], metadata, 5, "same"))
|
| 107 |
+
sr10.append(recall_at_k(retrieved, q["modality"], q["class"], metadata, 10, "same"))
|
| 108 |
+
cr1.append(recall_at_k(retrieved, q["modality"], q["class"], metadata, 1, "cross"))
|
| 109 |
+
cr5.append(recall_at_k(retrieved, q["modality"], q["class"], metadata, 5, "cross"))
|
| 110 |
+
cr10.append(recall_at_k(retrieved, q["modality"], q["class"], metadata, 10, "cross"))
|
| 111 |
+
|
| 112 |
+
n = max(len(sr1), 1)
|
| 113 |
+
return Result(
|
| 114 |
+
model=label,
|
| 115 |
+
same_r1=np.mean(sr1) / 1.0,
|
| 116 |
+
same_r5=np.mean(sr5) / 5.0,
|
| 117 |
+
same_r10=np.mean(sr10) / 10.0,
|
| 118 |
+
cross_r1=np.mean(cr1) / 1.0,
|
| 119 |
+
cross_r5=np.mean(cr5) / 5.0,
|
| 120 |
+
cross_r10=np.mean(cr10) / 10.0,
|
| 121 |
+
latency_ms=np.mean(latencies) if latencies else 0,
|
| 122 |
+
n_queries=n,
|
| 123 |
+
)
|
| 124 |
+
|
| 125 |
+
|
| 126 |
+
def main():
|
| 127 |
+
print("=" * 72)
|
| 128 |
+
print(" EXPERIMENT: DINOv2-CLIP Hybrid vs Pure CLIP")
|
| 129 |
+
print("=" * 72)
|
| 130 |
+
|
| 131 |
+
all_emb, metadata = load_data()
|
| 132 |
+
queries, gallery = split(metadata)
|
| 133 |
+
print(f"Gallery: {len(gallery)} | Queries: {len(queries)} | Dim: {all_emb.shape[1]}")
|
| 134 |
+
|
| 135 |
+
from transformers import CLIPProcessor, CLIPModel
|
| 136 |
+
device = "cuda" if torch.cuda.is_available() else "cpu"
|
| 137 |
+
processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")
|
| 138 |
+
clip_model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14").to(device)
|
| 139 |
+
clip_model.eval()
|
| 140 |
+
print(f"CLIP loaded on {device}")
|
| 141 |
+
|
| 142 |
+
@torch.no_grad()
|
| 143 |
+
def clip_extract(img, modality):
|
| 144 |
+
inputs = processor(images=img, return_tensors="pt").to(device)
|
| 145 |
+
out = clip_model.vision_model(**inputs)
|
| 146 |
+
pooled = out.last_hidden_state[:, 0, :]
|
| 147 |
+
feat = clip_model.visual_projection(pooled).squeeze(0)
|
| 148 |
+
return torch.nn.functional.normalize(feat, dim=-1).cpu().numpy()
|
| 149 |
+
|
| 150 |
+
results = []
|
| 151 |
+
|
| 152 |
+
print("\n[1/4] Pure CLIP ...")
|
| 153 |
+
r = evaluate(queries, all_emb, metadata, gallery, clip_extract, "CLIP ViT-L/14")
|
| 154 |
+
results.append(r)
|
| 155 |
+
print(f" Same R@5={r.same_r5:.4f} Cross R@5={r.cross_r5:.4f} Latency={r.latency_ms:.0f}ms")
|
| 156 |
+
|
| 157 |
+
print("[2/4] CLIP + SAR Adapter ...")
|
| 158 |
+
from src.features.sar_adapter import SARAdapter
|
| 159 |
+
adapter = SARAdapter().eval()
|
| 160 |
+
|
| 161 |
+
def clip_sar_extract(img, modality):
|
| 162 |
+
if modality == "sar":
|
| 163 |
+
arr = np.array(img).astype(np.float32) / 255.0
|
| 164 |
+
t = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)
|
| 165 |
+
with torch.no_grad():
|
| 166 |
+
adapted = adapter(t)
|
| 167 |
+
img = Image.fromarray((adapted.squeeze(0).permute(1, 2, 0).numpy() * 255).clip(0, 255).astype(np.uint8))
|
| 168 |
+
return clip_extract(img, modality)
|
| 169 |
+
|
| 170 |
+
r = evaluate(queries, all_emb, metadata, gallery, clip_sar_extract, "CLIP + SAR Adapter")
|
| 171 |
+
results.append(r)
|
| 172 |
+
print(f" Same R@5={r.same_r5:.4f} Cross R@5={r.cross_r5:.4f} Latency={r.latency_ms:.0f}ms")
|
| 173 |
+
|
| 174 |
+
print("[3/4] CLIP + DINOv2 Hybrid ...")
|
| 175 |
+
try:
|
| 176 |
+
dinov2 = torch.hub.load("facebookresearch/dinov2", "dinov2_vits14", pretrained=True)
|
| 177 |
+
dinov2.to(device).eval()
|
| 178 |
+
has_dino = True
|
| 179 |
+
dino_embed_dim = dinov2.embed_dim # 384 for vits14
|
| 180 |
+
print(f" DINOv2-ViT-S/14 loaded (embed_dim={dino_embed_dim})")
|
| 181 |
+
except Exception as e:
|
| 182 |
+
has_dino = False
|
| 183 |
+
print(f" DINOv2 load failed: {e}")
|
| 184 |
+
|
| 185 |
+
if has_dino:
|
| 186 |
+
from torchvision import transforms
|
| 187 |
+
dino_transform = transforms.Compose([
|
| 188 |
+
transforms.Resize((224, 224)),
|
| 189 |
+
transforms.ToTensor(),
|
| 190 |
+
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
| 191 |
+
])
|
| 192 |
+
|
| 193 |
+
# Projection to match CLIP dim (768) if needed
|
| 194 |
+
dino_proj = None
|
| 195 |
+
if dino_embed_dim != 768:
|
| 196 |
+
dino_proj = torch.nn.Linear(dino_embed_dim, 768, bias=False).to(device).eval()
|
| 197 |
+
with torch.no_grad():
|
| 198 |
+
torch.nn.init.eye_(dino_proj.weight) # identity init — preserves features
|
| 199 |
+
|
| 200 |
+
@torch.no_grad()
|
| 201 |
+
def clip_dino_extract(img, modality):
|
| 202 |
+
clip_feat = clip_extract(img, modality)
|
| 203 |
+
t = dino_transform(img).unsqueeze(0).to(device)
|
| 204 |
+
patch_feat = dinov2(t).squeeze(0)
|
| 205 |
+
if dino_proj is not None:
|
| 206 |
+
patch_feat = dino_proj(patch_feat)
|
| 207 |
+
patch_feat = torch.nn.functional.normalize(patch_feat, dim=-1).cpu().numpy()
|
| 208 |
+
hybrid = 0.7 * clip_feat + 0.3 * patch_feat
|
| 209 |
+
return hybrid / (np.linalg.norm(hybrid) + 1e-8)
|
| 210 |
+
|
| 211 |
+
r = evaluate(queries, all_emb, metadata, gallery, clip_dino_extract, "DINOv2-CLIP Hybrid")
|
| 212 |
+
results.append(r)
|
| 213 |
+
print(f" Same R@5={r.same_r5:.4f} Cross R@5={r.cross_r5:.4f} Latency={r.latency_ms:.0f}ms")
|
| 214 |
+
else:
|
| 215 |
+
r = evaluate(queries, all_emb, metadata, gallery, clip_extract, "CLIP (DINOv2 unavailable)")
|
| 216 |
+
results.append(r)
|
| 217 |
+
|
| 218 |
+
print("[4/4] Full Hybrid (CLIP + SAR + DINOv2) ...")
|
| 219 |
+
if has_dino:
|
| 220 |
+
def full_extract(img, modality):
|
| 221 |
+
if modality == "sar":
|
| 222 |
+
arr = np.array(img).astype(np.float32) / 255.0
|
| 223 |
+
t = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)
|
| 224 |
+
with torch.no_grad():
|
| 225 |
+
adapted = adapter(t)
|
| 226 |
+
img = Image.fromarray((adapted.squeeze(0).permute(1, 2, 0).numpy() * 255).clip(0, 255).astype(np.uint8))
|
| 227 |
+
return clip_dino_extract(img, modality)
|
| 228 |
+
|
| 229 |
+
r = evaluate(queries, all_emb, metadata, gallery, full_extract, "Full Hybrid (CLIP+SAR+DINOv2)")
|
| 230 |
+
results.append(r)
|
| 231 |
+
print(f" Same R@5={r.same_r5:.4f} Cross R@5={r.cross_r5:.4f} Latency={r.latency_ms:.0f}ms")
|
| 232 |
+
else:
|
| 233 |
+
r = evaluate(queries, all_emb, metadata, gallery, clip_sar_extract, "CLIP+SAR (DINOv2 unavailable)")
|
| 234 |
+
results.append(r)
|
| 235 |
+
|
| 236 |
+
print("\n" + "=" * 72)
|
| 237 |
+
print(" RESULTS")
|
| 238 |
+
print("=" * 72)
|
| 239 |
+
hdr = f"{'Model':<35} {'S-R@1':>6} {'S-R@5':>6} {'S-R@10':>7} {'C-R@1':>6} {'C-R@5':>6} {'C-R@10':>7} {'ms':>6}"
|
| 240 |
+
print(hdr)
|
| 241 |
+
print("-" * 72)
|
| 242 |
+
for r in results:
|
| 243 |
+
print(f"{r.model:<35} {r.same_r1:>6.4f} {r.same_r5:>6.4f} {r.same_r10:>7.4f} {r.cross_r1:>6.4f} {r.cross_r5:>6.4f} {r.cross_r10:>7.4f} {r.latency_ms:>5.0f}")
|
| 244 |
+
|
| 245 |
+
base_s5 = results[0].same_r5
|
| 246 |
+
base_c5 = results[0].cross_r5
|
| 247 |
+
print(f"\nDelta vs CLIP baseline (R@5):")
|
| 248 |
+
for r in results[1:]:
|
| 249 |
+
ds = r.same_r5 - base_s5
|
| 250 |
+
dc = r.cross_r5 - base_c5
|
| 251 |
+
print(f" {r.model}: Same {'+' if ds >= 0 else ''}{ds:.4f}, Cross {'+' if dc >= 0 else ''}{dc:.4f}")
|
| 252 |
+
|
| 253 |
+
out = PROCESSED_DIR / "experiment_results.json"
|
| 254 |
+
with open(out, "w") as f:
|
| 255 |
+
json.dump([asdict(r) for r in results], f, indent=2)
|
| 256 |
+
print(f"\nSaved to {out}")
|
| 257 |
+
|
| 258 |
+
|
| 259 |
+
if __name__ == "__main__":
|
| 260 |
+
try:
|
| 261 |
+
main()
|
| 262 |
+
except Exception:
|
| 263 |
+
traceback.print_exc()
|
notebooks/01_data_exploration.ipynb
ADDED
|
@@ -0,0 +1,154 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"cells": [
|
| 3 |
+
{
|
| 4 |
+
"cell_type": "markdown",
|
| 5 |
+
"metadata": {},
|
| 6 |
+
"source": [
|
| 7 |
+
"# Data Exploration\n",
|
| 8 |
+
"\n",
|
| 9 |
+
"This notebook explores the CrisisLandMark dataset structure and verifies preprocessing."
|
| 10 |
+
]
|
| 11 |
+
},
|
| 12 |
+
{
|
| 13 |
+
"cell_type": "code",
|
| 14 |
+
"execution_count": null,
|
| 15 |
+
"metadata": {},
|
| 16 |
+
"outputs": [],
|
| 17 |
+
"source": [
|
| 18 |
+
"import sys\n",
|
| 19 |
+
"sys.path.insert(0, '..')\n",
|
| 20 |
+
"\n",
|
| 21 |
+
"import torch\n",
|
| 22 |
+
"import numpy as np\n",
|
| 23 |
+
"from PIL import Image\n",
|
| 24 |
+
"import matplotlib.pyplot as plt\n",
|
| 25 |
+
"\n",
|
| 26 |
+
"from src.data.preprocessing import preprocess_image, handle_channels\n",
|
| 27 |
+
"from src.data.dataset import CrisisLandMarkDataset, create_splits"
|
| 28 |
+
]
|
| 29 |
+
},
|
| 30 |
+
{
|
| 31 |
+
"cell_type": "markdown",
|
| 32 |
+
"metadata": {},
|
| 33 |
+
"source": [
|
| 34 |
+
"## 1. Create Dataset"
|
| 35 |
+
]
|
| 36 |
+
},
|
| 37 |
+
{
|
| 38 |
+
"cell_type": "code",
|
| 39 |
+
"execution_count": null,
|
| 40 |
+
"metadata": {},
|
| 41 |
+
"outputs": [],
|
| 42 |
+
"source": [
|
| 43 |
+
"# Create dataset for each modality\n",
|
| 44 |
+
"optical_dataset = CrisisLandMarkDataset(modality='optical')\n",
|
| 45 |
+
"sar_dataset = CrisisLandMarkDataset(modality='sar')\n",
|
| 46 |
+
"\n",
|
| 47 |
+
"print(f\"Optical dataset: {len(optical_dataset)} samples\")\n",
|
| 48 |
+
"print(f\"SAR dataset: {len(sar_dataset)} samples\")"
|
| 49 |
+
]
|
| 50 |
+
},
|
| 51 |
+
{
|
| 52 |
+
"cell_type": "markdown",
|
| 53 |
+
"metadata": {},
|
| 54 |
+
"source": [
|
| 55 |
+
"## 2. Sample Images"
|
| 56 |
+
]
|
| 57 |
+
},
|
| 58 |
+
{
|
| 59 |
+
"cell_type": "code",
|
| 60 |
+
"execution_count": null,
|
| 61 |
+
"metadata": {},
|
| 62 |
+
"outputs": [],
|
| 63 |
+
"source": [
|
| 64 |
+
"# Get sample images\n",
|
| 65 |
+
"optical_img, optical_mod, optical_class = optical_dataset[0]\n",
|
| 66 |
+
"sar_img, sar_mod, sar_class = sar_dataset[0]\n",
|
| 67 |
+
"\n",
|
| 68 |
+
"print(f\"Optical shape: {optical_img.shape}\")\n",
|
| 69 |
+
"print(f\"SAR shape: {sar_img.shape}\")\n",
|
| 70 |
+
"print(f\"Optical modality label: {optical_mod}\")\n",
|
| 71 |
+
"print(f\"SAR modality label: {sar_mod}\")"
|
| 72 |
+
]
|
| 73 |
+
},
|
| 74 |
+
{
|
| 75 |
+
"cell_type": "markdown",
|
| 76 |
+
"metadata": {},
|
| 77 |
+
"source": [
|
| 78 |
+
"## 3. Data Splitting"
|
| 79 |
+
]
|
| 80 |
+
},
|
| 81 |
+
{
|
| 82 |
+
"cell_type": "code",
|
| 83 |
+
"execution_count": null,
|
| 84 |
+
"metadata": {},
|
| 85 |
+
"outputs": [],
|
| 86 |
+
"source": [
|
| 87 |
+
"# Test splitting\n",
|
| 88 |
+
"query_idx, gallery_idx = create_splits(optical_dataset, query_ratio=0.2)\n",
|
| 89 |
+
"\n",
|
| 90 |
+
"print(f\"Query set: {len(query_idx)} samples\")\n",
|
| 91 |
+
"print(f\"Gallery set: {len(gallery_idx)} samples\")\n",
|
| 92 |
+
"print(f\"Overlap: {len(set(query_idx) & set(gallery_idx))} (should be 0)\")"
|
| 93 |
+
]
|
| 94 |
+
},
|
| 95 |
+
{
|
| 96 |
+
"cell_type": "markdown",
|
| 97 |
+
"metadata": {},
|
| 98 |
+
"source": [
|
| 99 |
+
"## 4. Class Distribution"
|
| 100 |
+
]
|
| 101 |
+
},
|
| 102 |
+
{
|
| 103 |
+
"cell_type": "code",
|
| 104 |
+
"execution_count": null,
|
| 105 |
+
"metadata": {},
|
| 106 |
+
"outputs": [],
|
| 107 |
+
"source": [
|
| 108 |
+
"# Check class distribution\n",
|
| 109 |
+
"class_counts = {}\n",
|
| 110 |
+
"for i in range(len(optical_dataset)):\n",
|
| 111 |
+
" _, _, class_label = optical_dataset[i]\n",
|
| 112 |
+
" class_counts[class_label] = class_counts.get(class_label, 0) + 1\n",
|
| 113 |
+
"\n",
|
| 114 |
+
"print(\"Class distribution:\")\n",
|
| 115 |
+
"for cls, count in sorted(class_counts.items()):\n",
|
| 116 |
+
" print(f\" Class {cls}: {count} samples\")"
|
| 117 |
+
]
|
| 118 |
+
},
|
| 119 |
+
{
|
| 120 |
+
"cell_type": "markdown",
|
| 121 |
+
"metadata": {},
|
| 122 |
+
"source": [
|
| 123 |
+
"## 5. Summary"
|
| 124 |
+
]
|
| 125 |
+
},
|
| 126 |
+
{
|
| 127 |
+
"cell_type": "code",
|
| 128 |
+
"execution_count": null,
|
| 129 |
+
"metadata": {},
|
| 130 |
+
"outputs": [],
|
| 131 |
+
"source": [
|
| 132 |
+
"print(\"\\n=== Data Exploration Summary ===\")\n",
|
| 133 |
+
"print(f\"Total optical samples: {len(optical_dataset)}\")\n",
|
| 134 |
+
"print(f\"Total SAR samples: {len(sar_dataset)}\")\n",
|
| 135 |
+
"print(f\"Number of classes: {len(class_counts)}\")\n",
|
| 136 |
+
"print(f\"Query/Gallery split: 80/20 with no overlap\")\n",
|
| 137 |
+
"print(\"\\nPreprocessing verified for optical and SAR modalities.\")"
|
| 138 |
+
]
|
| 139 |
+
}
|
| 140 |
+
],
|
| 141 |
+
"metadata": {
|
| 142 |
+
"kernelspec": {
|
| 143 |
+
"display_name": "Python 3",
|
| 144 |
+
"language": "python",
|
| 145 |
+
"name": "python3"
|
| 146 |
+
},
|
| 147 |
+
"language_info": {
|
| 148 |
+
"name": "python",
|
| 149 |
+
"version": "3.10.0"
|
| 150 |
+
}
|
| 151 |
+
},
|
| 152 |
+
"nbformat": 4,
|
| 153 |
+
"nbformat_minor": 4
|
| 154 |
+
}
|
prepare_gallery.py
ADDED
|
@@ -0,0 +1,394 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Prepare gallery: download dataset, extract embeddings, build FAISS index.
|
| 3 |
+
|
| 4 |
+
Downloads real multi-modal satellite data (optical RGB, SAR 2ch, MS 13ch)
|
| 5 |
+
and builds a cross-modal retrieval gallery.
|
| 6 |
+
|
| 7 |
+
Usage:
|
| 8 |
+
python prepare_gallery.py --samples 50
|
| 9 |
+
"""
|
| 10 |
+
|
| 11 |
+
import argparse
|
| 12 |
+
import json
|
| 13 |
+
import shutil
|
| 14 |
+
import time
|
| 15 |
+
from pathlib import Path
|
| 16 |
+
|
| 17 |
+
import numpy as np
|
| 18 |
+
import torch
|
| 19 |
+
import torch.nn.functional as F
|
| 20 |
+
from PIL import Image
|
| 21 |
+
from tqdm import tqdm
|
| 22 |
+
from torchvision import transforms
|
| 23 |
+
|
| 24 |
+
# ---------------------------------------------------------------------------
|
| 25 |
+
# Config
|
| 26 |
+
# ---------------------------------------------------------------------------
|
| 27 |
+
DATA_DIR = Path("data")
|
| 28 |
+
RAW_DIR = DATA_DIR / "raw"
|
| 29 |
+
PROCESSED_DIR = DATA_DIR / "processed"
|
| 30 |
+
GALLERY_DIR = DATA_DIR / "gallery"
|
| 31 |
+
|
| 32 |
+
CLASSES = [
|
| 33 |
+
"AnnualCrop", "Forest", "HerbaceousVegetation", "Highway",
|
| 34 |
+
"Industrial", "Pasture", "PermanentCrop", "Residential",
|
| 35 |
+
"River", "SeaLake",
|
| 36 |
+
]
|
| 37 |
+
|
| 38 |
+
BATCH_SIZE = 64
|
| 39 |
+
EMBED_DIM = 768 # CLIP ViT-L/14 output dim
|
| 40 |
+
|
| 41 |
+
|
| 42 |
+
def _samples_exist(dir: Path, n_per_class: int) -> bool:
|
| 43 |
+
"""Check if directory has n_per_class images per class."""
|
| 44 |
+
if not dir.exists():
|
| 45 |
+
return False
|
| 46 |
+
return all(len(list((dir / c).glob("*.*"))) >= n_per_class for c in CLASSES)
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
def download_optical(n_per_class: int = 50) -> Path:
|
| 50 |
+
"""Download optical RGB from HuggingFace."""
|
| 51 |
+
from datasets import load_dataset
|
| 52 |
+
out_dir = RAW_DIR / "eurosat"
|
| 53 |
+
if _samples_exist(out_dir, n_per_class):
|
| 54 |
+
print(f"Optical already at {out_dir}, skipping.")
|
| 55 |
+
return out_dir
|
| 56 |
+
out_dir.mkdir(parents=True, exist_ok=True)
|
| 57 |
+
print("Downloading optical RGB from blanchon/EuroSAT_RGB ...")
|
| 58 |
+
ds = load_dataset("blanchon/EuroSAT_RGB", split="train")
|
| 59 |
+
selected, counts = [], {c: 0 for c in range(10)}
|
| 60 |
+
for row in ds:
|
| 61 |
+
lbl = row["label"]
|
| 62 |
+
if counts[lbl] < n_per_class:
|
| 63 |
+
selected.append(row)
|
| 64 |
+
counts[lbl] += 1
|
| 65 |
+
if all(v >= n_per_class for v in counts.values()):
|
| 66 |
+
break
|
| 67 |
+
for i, row in enumerate(tqdm(selected, desc="Saving optical")):
|
| 68 |
+
cls_name = CLASSES[row["label"]]
|
| 69 |
+
cls_dir = out_dir / cls_name
|
| 70 |
+
cls_dir.mkdir(exist_ok=True)
|
| 71 |
+
row["image"].save(cls_dir / f"{cls_name}_{i}.tif")
|
| 72 |
+
print(f"Optical saved to {out_dir}")
|
| 73 |
+
return out_dir
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
def download_sar(n_per_class: int = 50) -> Path:
|
| 77 |
+
"""Download real SAR (2ch VV/VH) from HuggingFace dataset or zip."""
|
| 78 |
+
out_dir = RAW_DIR / "eurosat_sar"
|
| 79 |
+
if _samples_exist(out_dir, n_per_class):
|
| 80 |
+
print(f"SAR already at {out_dir}, skipping.")
|
| 81 |
+
return out_dir
|
| 82 |
+
out_dir.mkdir(parents=True, exist_ok=True)
|
| 83 |
+
print("Downloading SAR from wangyi111/EuroSAT-SAR ...")
|
| 84 |
+
try:
|
| 85 |
+
from datasets import load_dataset
|
| 86 |
+
ds = load_dataset("wangyi111/EuroSAT-SAR", split="train", streaming=True)
|
| 87 |
+
selected, counts = [], {c: 0 for c in range(10)}
|
| 88 |
+
for row in ds:
|
| 89 |
+
# SAR labels are class names directly
|
| 90 |
+
lbl_name = row.get("label", row.get("label_name", ""))
|
| 91 |
+
if isinstance(lbl_name, int) and lbl_name < 10:
|
| 92 |
+
cls_name = CLASSES[lbl_name]
|
| 93 |
+
elif isinstance(lbl_name, str) and lbl_name in CLASSES:
|
| 94 |
+
cls_name = lbl_name
|
| 95 |
+
else:
|
| 96 |
+
continue
|
| 97 |
+
idx = CLASSES.index(cls_name)
|
| 98 |
+
if counts[idx] < n_per_class:
|
| 99 |
+
# Convert to 2-channel grayscale (VV, VH from RGBA)
|
| 100 |
+
img = np.array(row["image"].convert("L"))
|
| 101 |
+
selected.append((img, cls_name))
|
| 102 |
+
counts[idx] += 1
|
| 103 |
+
if all(v >= n_per_class for v in counts.values()):
|
| 104 |
+
break
|
| 105 |
+
for i, (img_arr, cls_name) in enumerate(tqdm(selected, desc="Saving SAR")):
|
| 106 |
+
cls_dir = out_dir / cls_name
|
| 107 |
+
cls_dir.mkdir(exist_ok=True)
|
| 108 |
+
# Save as 2-channel TIFF (stack Luminance as VV/VH)
|
| 109 |
+
two_ch = np.stack([img_arr, img_arr], axis=-1).astype(np.uint8)
|
| 110 |
+
Image.fromarray(two_ch[:, :, 0], mode="L").save(
|
| 111 |
+
cls_dir / f"{cls_name}_{i}.tif")
|
| 112 |
+
print(f"SAR saved to {out_dir}")
|
| 113 |
+
except Exception as e:
|
| 114 |
+
print(f"SAR download failed ({e}), using local fallback.")
|
| 115 |
+
# Fallback: convert optical to 2-channel SAR-like
|
| 116 |
+
optical_dir = RAW_DIR / "eurosat"
|
| 117 |
+
if optical_dir.exists():
|
| 118 |
+
for cls_name in CLASSES:
|
| 119 |
+
cls_in = optical_dir / cls_name
|
| 120 |
+
cls_out = out_dir / cls_name
|
| 121 |
+
cls_out.mkdir(parents=True, exist_ok=True)
|
| 122 |
+
for path in list(cls_in.glob("*.*"))[:n_per_class]:
|
| 123 |
+
arr = np.array(Image.open(path).convert("L"))
|
| 124 |
+
noise = np.random.rayleigh(1.0, arr.shape).astype(np.float32)
|
| 125 |
+
sar = np.clip(arr * noise, 0, 255).astype(np.uint8)
|
| 126 |
+
Image.fromarray(sar, mode="L").save(
|
| 127 |
+
cls_out / f"{cls_name}_{path.stem}.tif")
|
| 128 |
+
print(f"SAR fallback saved to {out_dir}")
|
| 129 |
+
return out_dir
|
| 130 |
+
|
| 131 |
+
|
| 132 |
+
def download_multispectral(n_per_class: int = 50) -> Path:
|
| 133 |
+
"""Download real multispectral (13ch) from HuggingFace."""
|
| 134 |
+
out_dir = RAW_DIR / "eurosat_ms"
|
| 135 |
+
if _samples_exist(out_dir, n_per_class):
|
| 136 |
+
print(f"Multispectral already at {out_dir}, skipping.")
|
| 137 |
+
return out_dir
|
| 138 |
+
out_dir.mkdir(parents=True, exist_ok=True)
|
| 139 |
+
print("Downloading MS from giswqs/EuroSAT_MS ...")
|
| 140 |
+
try:
|
| 141 |
+
from datasets import load_dataset
|
| 142 |
+
ds = load_dataset("giswqs/EuroSAT_MS", split="train", streaming=True)
|
| 143 |
+
selected, counts = [], {c: 0 for c in range(10)}
|
| 144 |
+
for row in ds:
|
| 145 |
+
lbl = row["label"]
|
| 146 |
+
if counts[lbl] < n_per_class:
|
| 147 |
+
# Image is a list of 13 arrays (one per band)
|
| 148 |
+
bands = [np.array(b) for b in row["image"]]
|
| 149 |
+
selected.append((bands, lbl))
|
| 150 |
+
counts[lbl] += 1
|
| 151 |
+
if all(v >= n_per_class for v in counts.values()):
|
| 152 |
+
break
|
| 153 |
+
for bands, lbl in tqdm(selected, desc="Saving MS"):
|
| 154 |
+
cls_name = CLASSES[lbl]
|
| 155 |
+
cls_dir = out_dir / cls_name
|
| 156 |
+
cls_dir.mkdir(exist_ok=True)
|
| 157 |
+
# Stack bands into single array, save as multi-channel TIFF
|
| 158 |
+
arr = np.stack(bands, axis=0) # (13, 64, 64)
|
| 159 |
+
import tifffile
|
| 160 |
+
tifffile.imwrite(
|
| 161 |
+
cls_dir / f"{cls_name}_{len(list(cls_dir.glob('*.*')))}.tif",
|
| 162 |
+
arr.astype(np.uint16))
|
| 163 |
+
print(f"MS saved to {out_dir}")
|
| 164 |
+
except Exception as e:
|
| 165 |
+
print(f"MS download failed ({e}), using local fallback.")
|
| 166 |
+
optical_dir = RAW_DIR / "eurosat"
|
| 167 |
+
if optical_dir.exists():
|
| 168 |
+
for cls_name in CLASSES:
|
| 169 |
+
cls_in = optical_dir / cls_name
|
| 170 |
+
cls_out = out_dir / cls_name
|
| 171 |
+
cls_out.mkdir(parents=True, exist_ok=True)
|
| 172 |
+
for path in list(cls_in.glob("*.*"))[:n_per_class]:
|
| 173 |
+
shutil.copy2(path, cls_out / path.name)
|
| 174 |
+
print(f"MS fallback saved to {out_dir}")
|
| 175 |
+
return out_dir
|
| 176 |
+
|
| 177 |
+
|
| 178 |
+
def load_satclip():
|
| 179 |
+
"""Load SatCLIP encoder."""
|
| 180 |
+
from src.features.satclip_encoder import SatCLIPEncoder
|
| 181 |
+
print("Loading SatCLIP encoder...")
|
| 182 |
+
encoder = SatCLIPEncoder()
|
| 183 |
+
print(f"SatCLIP loaded on {encoder.device}")
|
| 184 |
+
return encoder
|
| 185 |
+
|
| 186 |
+
|
| 187 |
+
def _load_multichannel_image(path: Path, modality: str) -> torch.Tensor:
|
| 188 |
+
"""
|
| 189 |
+
Load an image handling different channel counts.
|
| 190 |
+
|
| 191 |
+
Returns tensor of shape (C, H, W) normalized to [0, 1].
|
| 192 |
+
"""
|
| 193 |
+
# Try tifffile first for multi-channel TIFFs
|
| 194 |
+
try:
|
| 195 |
+
import tifffile
|
| 196 |
+
arr = tifffile.imread(str(path))
|
| 197 |
+
if arr.ndim == 3 and arr.shape[-1] in [2, 3, 4, 13]:
|
| 198 |
+
# Channels-last format
|
| 199 |
+
arr = np.transpose(arr, (2, 0, 1))
|
| 200 |
+
elif arr.ndim == 2:
|
| 201 |
+
arr = arr[np.newaxis, :, :]
|
| 202 |
+
# Normalize uint to [0, 1]
|
| 203 |
+
if arr.dtype in [np.uint8, np.uint16]:
|
| 204 |
+
arr = arr.astype(np.float32) / np.float32(np.iinfo(arr.dtype).max)
|
| 205 |
+
else:
|
| 206 |
+
arr = arr.astype(np.float32)
|
| 207 |
+
arr = (arr - arr.min()) / (arr.max() - arr.min() + 1e-8)
|
| 208 |
+
tensor = torch.from_numpy(arr).float()
|
| 209 |
+
except Exception:
|
| 210 |
+
# Fallback to PIL (handles RGB)
|
| 211 |
+
img = Image.open(path).convert("RGB")
|
| 212 |
+
tensor = transforms.ToTensor()(img)
|
| 213 |
+
|
| 214 |
+
# Resize to 224x224
|
| 215 |
+
if tensor.shape[1] != 224 or tensor.shape[2] != 224:
|
| 216 |
+
tensor = F.interpolate(tensor.unsqueeze(0), size=(224, 224),
|
| 217 |
+
mode="bilinear", align_corners=False).squeeze(0)
|
| 218 |
+
|
| 219 |
+
# Enforce strict channel counts per modality to prevent torch.stack failures
|
| 220 |
+
c = tensor.shape[0]
|
| 221 |
+
if modality == "optical":
|
| 222 |
+
if c == 1:
|
| 223 |
+
tensor = tensor.repeat(3, 1, 1)
|
| 224 |
+
elif c > 3:
|
| 225 |
+
tensor = tensor[:3]
|
| 226 |
+
elif c == 2:
|
| 227 |
+
tensor = torch.cat([tensor, tensor[:1]], dim=0)
|
| 228 |
+
elif modality == "sar":
|
| 229 |
+
if c == 1:
|
| 230 |
+
tensor = tensor.repeat(2, 1, 1)
|
| 231 |
+
elif c > 2:
|
| 232 |
+
tensor = tensor[:2]
|
| 233 |
+
elif modality == "multispectral":
|
| 234 |
+
if c < 13:
|
| 235 |
+
pad = torch.zeros(13 - c, tensor.shape[1], tensor.shape[2])
|
| 236 |
+
tensor = torch.cat([tensor, pad], dim=0)
|
| 237 |
+
elif c > 13:
|
| 238 |
+
tensor = tensor[:13]
|
| 239 |
+
|
| 240 |
+
return tensor
|
| 241 |
+
|
| 242 |
+
|
| 243 |
+
def _pad_to_13ch(tensor: torch.Tensor, modality: str) -> torch.Tensor:
|
| 244 |
+
"""Pad tensor to 13 channels for SatCLIP."""
|
| 245 |
+
n_channels = tensor.shape[1]
|
| 246 |
+
if n_channels >= 13:
|
| 247 |
+
return tensor[:, :13]
|
| 248 |
+
# Repeat channels if single-channel (SAR fallback)
|
| 249 |
+
if n_channels == 1:
|
| 250 |
+
tensor = tensor.repeat(1, 3, 1, 1)
|
| 251 |
+
n_channels = 3
|
| 252 |
+
pad_channels = 13 - n_channels
|
| 253 |
+
padding = torch.zeros(
|
| 254 |
+
tensor.shape[0], pad_channels, tensor.shape[2], tensor.shape[3])
|
| 255 |
+
return torch.cat([tensor, padding], dim=1)
|
| 256 |
+
|
| 257 |
+
|
| 258 |
+
def _make_rgb_preview(path: Path, modality: str, size=(128, 128)) -> Image.Image:
|
| 259 |
+
"""Create an RGB preview image from any modality file."""
|
| 260 |
+
try:
|
| 261 |
+
import tifffile
|
| 262 |
+
arr = tifffile.imread(str(path))
|
| 263 |
+
if arr.ndim == 3 and arr.shape[-1] >= 3:
|
| 264 |
+
preview = arr[:, :, :3]
|
| 265 |
+
elif arr.ndim == 3 and arr.shape[0] >= 3:
|
| 266 |
+
preview = np.transpose(arr[:3], (1, 2, 0))
|
| 267 |
+
elif arr.ndim == 2:
|
| 268 |
+
preview = np.stack([arr] * 3, axis=-1)
|
| 269 |
+
else:
|
| 270 |
+
preview = np.stack([arr[:, :, 0]] * 3, axis=-1)
|
| 271 |
+
|
| 272 |
+
# Normalize for display
|
| 273 |
+
if preview.dtype == np.uint16:
|
| 274 |
+
preview = (preview / 65535.0 * 255).astype(np.uint8)
|
| 275 |
+
elif preview.dtype == np.uint8:
|
| 276 |
+
pass
|
| 277 |
+
else:
|
| 278 |
+
preview = (np.clip(preview, 0, 1) * 255).astype(np.uint8)
|
| 279 |
+
|
| 280 |
+
# Special handling for SAR: grayscale with colormap feel
|
| 281 |
+
if modality == "sar":
|
| 282 |
+
preview = preview # Keep as is
|
| 283 |
+
|
| 284 |
+
return Image.fromarray(preview).resize(size, Image.LANCZOS)
|
| 285 |
+
except Exception:
|
| 286 |
+
# Fallback to PIL
|
| 287 |
+
return Image.open(path).convert("RGB").resize(size, Image.LANCZOS)
|
| 288 |
+
|
| 289 |
+
|
| 290 |
+
@torch.no_grad()
|
| 291 |
+
def extract_embeddings_satclip(images, encoder, modality="optical"):
|
| 292 |
+
"""Extract L2-normalized embeddings from image tensors using SatCLIP."""
|
| 293 |
+
all_feats = []
|
| 294 |
+
for i in range(0, len(images), BATCH_SIZE):
|
| 295 |
+
batch = images[i: i + BATCH_SIZE]
|
| 296 |
+
tensors = torch.stack(batch)
|
| 297 |
+
# Pad to 13 channels if needed
|
| 298 |
+
if tensors.shape[1] < 13:
|
| 299 |
+
tensors = _pad_to_13ch(tensors, modality)
|
| 300 |
+
feats = encoder.encode(tensors, normalize=True)
|
| 301 |
+
all_feats.append(feats.cpu())
|
| 302 |
+
return torch.cat(all_feats, dim=0)
|
| 303 |
+
|
| 304 |
+
|
| 305 |
+
def build_gallery(n_per_class: int = 50):
|
| 306 |
+
"""Full pipeline: download, embed, build index."""
|
| 307 |
+
t0 = time.time()
|
| 308 |
+
|
| 309 |
+
# 1. Download data for all three modalities
|
| 310 |
+
optical_dir = download_optical(n_per_class)
|
| 311 |
+
sar_dir = download_sar(n_per_class)
|
| 312 |
+
ms_dir = download_multispectral(n_per_class)
|
| 313 |
+
|
| 314 |
+
# 2. Collect all images with proper multi-channel loading
|
| 315 |
+
modalities = {
|
| 316 |
+
"optical": optical_dir,
|
| 317 |
+
"sar": sar_dir,
|
| 318 |
+
"multispectral": ms_dir,
|
| 319 |
+
}
|
| 320 |
+
all_images = [] # list of (image_tensor, modality, class_name, path)
|
| 321 |
+
for mod, base_dir in modalities.items():
|
| 322 |
+
for cls_dir in sorted(base_dir.iterdir()):
|
| 323 |
+
if not cls_dir.is_dir():
|
| 324 |
+
continue
|
| 325 |
+
paths = sorted(cls_dir.glob("*.*"))[:n_per_class]
|
| 326 |
+
for img_path in paths:
|
| 327 |
+
tensor = _load_multichannel_image(img_path, mod)
|
| 328 |
+
all_images.append((tensor, mod, cls_dir.name, img_path))
|
| 329 |
+
|
| 330 |
+
print(f"\nTotal gallery images: {len(all_images)}")
|
| 331 |
+
for mod in modalities:
|
| 332 |
+
count = sum(1 for _, m, _, _ in all_images if m == mod)
|
| 333 |
+
print(f" {mod}: {count}")
|
| 334 |
+
|
| 335 |
+
# 3. Build gallery preview images and extract embeddings
|
| 336 |
+
print("\nBuilding gallery previews ...")
|
| 337 |
+
GALLERY_DIR.mkdir(parents=True, exist_ok=True)
|
| 338 |
+
for i, (_, mod, cls, path) in enumerate(tqdm(all_images, desc="Previews")):
|
| 339 |
+
preview = _make_rgb_preview(path, mod)
|
| 340 |
+
preview.save(GALLERY_DIR / f"{i:05d}_{mod}_{cls}.png")
|
| 341 |
+
|
| 342 |
+
# 4. Extract SatCLIP embeddings
|
| 343 |
+
print("\nExtracting SatCLIP embeddings ...")
|
| 344 |
+
encoder = load_satclip()
|
| 345 |
+
embeddings_by_mod = {}
|
| 346 |
+
for mod in ["optical", "sar", "multispectral"]:
|
| 347 |
+
mod_tensors = [img for img, m, _, _ in all_images if m == mod]
|
| 348 |
+
if mod_tensors:
|
| 349 |
+
print(f" Extracting {mod} ({len(mod_tensors)} images)...")
|
| 350 |
+
embeddings_by_mod[mod] = extract_embeddings_satclip(
|
| 351 |
+
mod_tensors, encoder, mod)
|
| 352 |
+
embeddings = torch.cat(list(embeddings_by_mod.values()), dim=0)
|
| 353 |
+
print(f"Embeddings shape: {embeddings.shape}")
|
| 354 |
+
|
| 355 |
+
# 5. Build FAISS index
|
| 356 |
+
print("\nBuilding FAISS index ...")
|
| 357 |
+
import faiss
|
| 358 |
+
embed_dim = embeddings.shape[1]
|
| 359 |
+
index = faiss.IndexFlatIP(embed_dim)
|
| 360 |
+
index.add(embeddings.numpy().astype(np.float32))
|
| 361 |
+
print(f"FAISS index size: {index.ntotal}")
|
| 362 |
+
|
| 363 |
+
# 6. Save everything
|
| 364 |
+
PROCESSED_DIR.mkdir(parents=True, exist_ok=True)
|
| 365 |
+
faiss.write_index(index, str(PROCESSED_DIR / "gallery.index"))
|
| 366 |
+
torch.save(embeddings, PROCESSED_DIR / "gallery_embeddings.pt")
|
| 367 |
+
|
| 368 |
+
metadata = []
|
| 369 |
+
for i, (_, mod, cls, path) in enumerate(all_images):
|
| 370 |
+
metadata.append({
|
| 371 |
+
"index": i,
|
| 372 |
+
"modality": mod,
|
| 373 |
+
"class": cls,
|
| 374 |
+
"gallery_path": str(GALLERY_DIR / f"{i:05d}_{mod}_{cls}.png"),
|
| 375 |
+
"original_path": str(path),
|
| 376 |
+
})
|
| 377 |
+
with open(PROCESSED_DIR / "gallery_metadata.json", "w") as f:
|
| 378 |
+
json.dump(metadata, f, indent=2)
|
| 379 |
+
|
| 380 |
+
elapsed = time.time() - t0
|
| 381 |
+
print(f"\nDone in {elapsed:.1f}s")
|
| 382 |
+
print(f"Index: {PROCESSED_DIR / 'gallery.index'}")
|
| 383 |
+
print(f"Embeddings:{PROCESSED_DIR / 'gallery_embeddings.pt'}")
|
| 384 |
+
print(f"Metadata: {PROCESSED_DIR / 'gallery_metadata.json'}")
|
| 385 |
+
print(f"Gallery: {GALLERY_DIR}")
|
| 386 |
+
|
| 387 |
+
|
| 388 |
+
if __name__ == "__main__":
|
| 389 |
+
parser = argparse.ArgumentParser(
|
| 390 |
+
description="Build multi-modal satellite image gallery")
|
| 391 |
+
parser.add_argument("--samples", type=int, default=50,
|
| 392 |
+
help="Images per class per modality")
|
| 393 |
+
args = parser.parse_args()
|
| 394 |
+
build_gallery(n_per_class=args.samples)
|
requirements.txt
ADDED
|
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Core ML & Transformers
|
| 2 |
+
torch>=2.0.0
|
| 3 |
+
torchvision>=0.15.0
|
| 4 |
+
transformers>=4.30.0
|
| 5 |
+
datasets>=2.14.0
|
| 6 |
+
scikit-learn>=1.3.0
|
| 7 |
+
numpy>=1.24.0
|
| 8 |
+
pillow>=10.0.0
|
| 9 |
+
|
| 10 |
+
# OpenAI CLIP
|
| 11 |
+
git+https://github.com/openai/CLIP.git
|
| 12 |
+
|
| 13 |
+
# Vector Search
|
| 14 |
+
faiss-cpu>=1.7.4
|
| 15 |
+
|
| 16 |
+
# API Server & UI
|
| 17 |
+
fastapi>=0.100.0
|
| 18 |
+
uvicorn>=0.22.0
|
| 19 |
+
gradio>=4.0.0
|
| 20 |
+
|
| 21 |
+
# GIS & Geospatial Indexing
|
| 22 |
+
h3>=3.7.6
|
| 23 |
+
tifffile>=2023.7.10
|
| 24 |
+
|
| 25 |
+
# Development
|
| 26 |
+
pytest>=7.4.0
|
src/__init__.py
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
# Cross-Modal Satellite Image Retrieval
|
src/data/README.md
ADDED
|
@@ -0,0 +1,34 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Data Module
|
| 2 |
+
|
| 3 |
+
Per-modality preprocessing for satellite imagery.
|
| 4 |
+
|
| 5 |
+
## Files
|
| 6 |
+
|
| 7 |
+
| File | Description |
|
| 8 |
+
|------|-------------|
|
| 9 |
+
| `preprocessing.py` | Per-modality transforms (optical, SAR, multispectral) |
|
| 10 |
+
| `dataset.py` | Dataset loading, splitting, and DataLoader creation |
|
| 11 |
+
|
| 12 |
+
## Supported Modalities
|
| 13 |
+
|
| 14 |
+
| Modality | Channels | Description |
|
| 15 |
+
|----------|----------|-------------|
|
| 16 |
+
| Optical | 3 (RGB) | Sentinel-2 bands B4, B3, B2 |
|
| 17 |
+
| SAR | 2 (VV/VH) | Sentinel-1 C-band |
|
| 18 |
+
| Multispectral | 12 | Sentinel-2 all bands |
|
| 19 |
+
|
| 20 |
+
## Usage
|
| 21 |
+
|
| 22 |
+
```python
|
| 23 |
+
from src.data.preprocessing import preprocess_image
|
| 24 |
+
from src.data.dataset import CrisisLandMarkDataset, create_splits
|
| 25 |
+
|
| 26 |
+
# Preprocess a single image
|
| 27 |
+
tensor = preprocess_image(image, modality="optical", size=224)
|
| 28 |
+
|
| 29 |
+
# Create dataset
|
| 30 |
+
dataset = CrisisLandMarkDataset(modality="optical")
|
| 31 |
+
|
| 32 |
+
# Split into query/gallery
|
| 33 |
+
query_idx, gallery_idx = create_splits(dataset, query_ratio=0.2)
|
| 34 |
+
```
|
src/data/__init__.py
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
# Data loading and preprocessing modules
|
src/data/dataset.py
ADDED
|
@@ -0,0 +1,197 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Dataset loading and splitting for cross-modal retrieval.
|
| 3 |
+
|
| 4 |
+
Handles:
|
| 5 |
+
- Loading CrisisLandMark dataset
|
| 6 |
+
- Per-modality preprocessing
|
| 7 |
+
- Query/gallery splitting (80/20)
|
| 8 |
+
- Ground-truth label preparation
|
| 9 |
+
"""
|
| 10 |
+
|
| 11 |
+
import torch
|
| 12 |
+
from torch.utils.data import Dataset, DataLoader
|
| 13 |
+
from pathlib import Path
|
| 14 |
+
from typing import Tuple, List, Dict, Optional
|
| 15 |
+
import numpy as np
|
| 16 |
+
from PIL import Image
|
| 17 |
+
from sklearn.model_selection import train_test_split
|
| 18 |
+
|
| 19 |
+
from .preprocessing import preprocess_image, handle_channels
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
class CrisisLandMarkDataset(Dataset):
|
| 23 |
+
"""
|
| 24 |
+
Dataset class for CrisisLandMark satellite imagery.
|
| 25 |
+
|
| 26 |
+
Supports:
|
| 27 |
+
- Optical (Sentinel-2 RGB)
|
| 28 |
+
- SAR (Sentinel-1 VV/VH)
|
| 29 |
+
- Multispectral (Sentinel-2 all bands)
|
| 30 |
+
"""
|
| 31 |
+
|
| 32 |
+
def __init__(
|
| 33 |
+
self,
|
| 34 |
+
data_dir: str = "data/raw/crisislandmark",
|
| 35 |
+
modality: str = "optical",
|
| 36 |
+
split: str = "train",
|
| 37 |
+
transform=None,
|
| 38 |
+
size: int = 224
|
| 39 |
+
):
|
| 40 |
+
"""
|
| 41 |
+
Initialize dataset.
|
| 42 |
+
|
| 43 |
+
Args:
|
| 44 |
+
data_dir: Path to dataset
|
| 45 |
+
modality: "optical", "sar", or "multispectral"
|
| 46 |
+
split: "train", "validation", or "test"
|
| 47 |
+
transform: Optional custom transform
|
| 48 |
+
size: Image resize size
|
| 49 |
+
"""
|
| 50 |
+
self.data_dir = Path(data_dir)
|
| 51 |
+
self.modality = modality
|
| 52 |
+
self.split = split
|
| 53 |
+
self.size = size
|
| 54 |
+
self.transform = transform
|
| 55 |
+
|
| 56 |
+
# ponytail: placeholder - load from actual dataset
|
| 57 |
+
# Real implementation would load from HuggingFace datasets
|
| 58 |
+
self.samples = self._load_samples()
|
| 59 |
+
self.labels = self._load_labels()
|
| 60 |
+
|
| 61 |
+
def _load_samples(self) -> List[Dict]:
|
| 62 |
+
"""Load sample metadata."""
|
| 63 |
+
# Placeholder - will be replaced with actual data loading
|
| 64 |
+
return [{"id": i, "path": f"sample_{i}.png"} for i in range(100)]
|
| 65 |
+
|
| 66 |
+
def _load_labels(self) -> Dict[int, int]:
|
| 67 |
+
"""Load ground-truth labels."""
|
| 68 |
+
# Placeholder - will be replaced with actual labels
|
| 69 |
+
return {i: i % 10 for i in range(100)}
|
| 70 |
+
|
| 71 |
+
def __len__(self) -> int:
|
| 72 |
+
return len(self.samples)
|
| 73 |
+
|
| 74 |
+
def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int, int]:
|
| 75 |
+
"""
|
| 76 |
+
Get sample.
|
| 77 |
+
|
| 78 |
+
Returns:
|
| 79 |
+
(image_tensor, modality_label, class_label)
|
| 80 |
+
"""
|
| 81 |
+
sample = self.samples[idx]
|
| 82 |
+
|
| 83 |
+
# Load image (placeholder)
|
| 84 |
+
image = Image.fromarray(np.random.randint(0, 255, (256, 256, 3), dtype=np.uint8))
|
| 85 |
+
|
| 86 |
+
# Preprocess
|
| 87 |
+
if self.transform:
|
| 88 |
+
image_tensor = self.transform(image)
|
| 89 |
+
else:
|
| 90 |
+
image_tensor = preprocess_image(image, self.modality, self.size)
|
| 91 |
+
|
| 92 |
+
# Modality label (0=optical, 1=sar, 2=multispectral)
|
| 93 |
+
modality_label = {"optical": 0, "sar": 1, "multispectral": 2}[self.modality]
|
| 94 |
+
|
| 95 |
+
# Class label
|
| 96 |
+
class_label = self.labels.get(idx, 0)
|
| 97 |
+
|
| 98 |
+
return image_tensor, modality_label, class_label
|
| 99 |
+
|
| 100 |
+
|
| 101 |
+
def create_splits(
|
| 102 |
+
dataset: CrisisLandMarkDataset,
|
| 103 |
+
query_ratio: float = 0.2,
|
| 104 |
+
seed: int = 42
|
| 105 |
+
) -> Tuple[List[int], List[int]]:
|
| 106 |
+
"""
|
| 107 |
+
Create query/gallery split with no overlap.
|
| 108 |
+
|
| 109 |
+
Args:
|
| 110 |
+
dataset: Full dataset
|
| 111 |
+
query_ratio: Fraction for query set
|
| 112 |
+
seed: Random seed for reproducibility
|
| 113 |
+
|
| 114 |
+
Returns:
|
| 115 |
+
(query_indices, gallery_indices)
|
| 116 |
+
"""
|
| 117 |
+
indices = list(range(len(dataset)))
|
| 118 |
+
|
| 119 |
+
# Stratify by class label if available
|
| 120 |
+
labels = [dataset.labels.get(i, 0) for i in indices]
|
| 121 |
+
|
| 122 |
+
query_idx, gallery_idx = train_test_split(
|
| 123 |
+
indices,
|
| 124 |
+
test_size=1 - query_ratio,
|
| 125 |
+
random_state=seed,
|
| 126 |
+
stratify=labels
|
| 127 |
+
)
|
| 128 |
+
|
| 129 |
+
# Verify no overlap
|
| 130 |
+
assert len(set(query_idx) & set(gallery_idx)) == 0, "Query and gallery sets overlap!"
|
| 131 |
+
|
| 132 |
+
return query_idx, gallery_idx
|
| 133 |
+
|
| 134 |
+
|
| 135 |
+
def get_dataloaders(
|
| 136 |
+
data_dir: str = "data/raw/crisislandmark",
|
| 137 |
+
modality: str = "optical",
|
| 138 |
+
batch_size: int = 32,
|
| 139 |
+
size: int = 224,
|
| 140 |
+
num_workers: int = 4
|
| 141 |
+
) -> Tuple[DataLoader, DataLoader]:
|
| 142 |
+
"""
|
| 143 |
+
Get train/test dataloaders.
|
| 144 |
+
|
| 145 |
+
Args:
|
| 146 |
+
data_dir: Path to dataset
|
| 147 |
+
modality: Modality type
|
| 148 |
+
batch_size: Batch size
|
| 149 |
+
size: Image size
|
| 150 |
+
num_workers: Number of workers
|
| 151 |
+
|
| 152 |
+
Returns:
|
| 153 |
+
(train_loader, test_loader)
|
| 154 |
+
"""
|
| 155 |
+
train_dataset = CrisisLandMarkDataset(data_dir, modality, "train", size=size)
|
| 156 |
+
test_dataset = CrisisLandMarkDataset(data_dir, modality, "test", size=size)
|
| 157 |
+
|
| 158 |
+
train_loader = DataLoader(
|
| 159 |
+
train_dataset,
|
| 160 |
+
batch_size=batch_size,
|
| 161 |
+
shuffle=True,
|
| 162 |
+
num_workers=num_workers
|
| 163 |
+
)
|
| 164 |
+
|
| 165 |
+
test_loader = DataLoader(
|
| 166 |
+
test_dataset,
|
| 167 |
+
batch_size=batch_size,
|
| 168 |
+
shuffle=False,
|
| 169 |
+
num_workers=num_workers
|
| 170 |
+
)
|
| 171 |
+
|
| 172 |
+
return train_loader, test_loader
|
| 173 |
+
|
| 174 |
+
|
| 175 |
+
# Self-check
|
| 176 |
+
if __name__ == "__main__":
|
| 177 |
+
# Test dataset creation
|
| 178 |
+
dataset = CrisisLandMarkDataset(modality="optical")
|
| 179 |
+
|
| 180 |
+
# Test split
|
| 181 |
+
query_idx, gallery_idx = create_splits(dataset, query_ratio=0.2)
|
| 182 |
+
|
| 183 |
+
print(f"Total samples: {len(dataset)}")
|
| 184 |
+
print(f"Query set: {len(query_idx)} samples")
|
| 185 |
+
print(f"Gallery set: {len(gallery_idx)} samples")
|
| 186 |
+
print(f"Overlap: {len(set(query_idx) & set(gallery_idx))} (should be 0)")
|
| 187 |
+
|
| 188 |
+
# Test dataloader
|
| 189 |
+
train_loader, test_loader = get_dataloaders(modality="optical", batch_size=4)
|
| 190 |
+
batch = next(iter(train_loader))
|
| 191 |
+
|
| 192 |
+
print(f"\nBatch shapes:")
|
| 193 |
+
print(f" Images: {batch[0].shape}")
|
| 194 |
+
print(f" Modality labels: {batch[1]}")
|
| 195 |
+
print(f" Class labels: {batch[2]}")
|
| 196 |
+
|
| 197 |
+
print("\nDataset test passed!")
|
src/data/download.py
ADDED
|
@@ -0,0 +1,80 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Download CrisisLandMark dataset from HuggingFace.
|
| 3 |
+
|
| 4 |
+
Dataset: DarthReca/crisislandmark
|
| 5 |
+
Size: 647K paired Sentinel-1 (SAR) and Sentinel-2 (optical) images
|
| 6 |
+
Labels: Land-cover annotations for retrieval evaluation
|
| 7 |
+
"""
|
| 8 |
+
|
| 9 |
+
import os
|
| 10 |
+
from pathlib import Path
|
| 11 |
+
from datasets import load_dataset
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
DATA_DIR = Path("data/raw/crisislandmark")
|
| 15 |
+
|
| 16 |
+
|
| 17 |
+
def download_dataset(subset: str = "train", cache_dir: str = None) -> None:
|
| 18 |
+
"""
|
| 19 |
+
Download CrisisLandMark dataset.
|
| 20 |
+
|
| 21 |
+
Args:
|
| 22 |
+
subset: Dataset split to download ("train", "validation", "test")
|
| 23 |
+
cache_dir: Custom cache directory
|
| 24 |
+
"""
|
| 25 |
+
print(f"Downloading CrisisLandMark dataset ({subset} split)...")
|
| 26 |
+
|
| 27 |
+
# ponytail: using HuggingFace datasets library for clean download
|
| 28 |
+
dataset = load_dataset(
|
| 29 |
+
"DarthReca/crisislandmark",
|
| 30 |
+
split=subset,
|
| 31 |
+
cache_dir=cache_dir or str(DATA_DIR / "cache"),
|
| 32 |
+
trust_remote_code=True
|
| 33 |
+
)
|
| 34 |
+
|
| 35 |
+
print(f"Downloaded {len(dataset)} samples")
|
| 36 |
+
print(f"Columns: {dataset.column_names}")
|
| 37 |
+
|
| 38 |
+
# Save to disk for faster loading later
|
| 39 |
+
output_dir = DATA_DIR / subset
|
| 40 |
+
output_dir.mkdir(parents=True, exist_ok=True)
|
| 41 |
+
|
| 42 |
+
dataset.save_to_disk(str(output_dir))
|
| 43 |
+
print(f"Saved to {output_dir}")
|
| 44 |
+
|
| 45 |
+
return dataset
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
def verify_dataset(subset: str = "train") -> dict:
|
| 49 |
+
"""
|
| 50 |
+
Verify downloaded dataset structure.
|
| 51 |
+
|
| 52 |
+
Returns:
|
| 53 |
+
Dictionary with dataset statistics
|
| 54 |
+
"""
|
| 55 |
+
dataset = load_dataset(
|
| 56 |
+
"DarthReca/crisislandmark",
|
| 57 |
+
split=subset,
|
| 58 |
+
cache_dir=str(DATA_DIR / "cache"),
|
| 59 |
+
trust_remote_code=True
|
| 60 |
+
)
|
| 61 |
+
|
| 62 |
+
stats = {
|
| 63 |
+
"total_samples": len(dataset),
|
| 64 |
+
"columns": dataset.column_names,
|
| 65 |
+
"features": {col: str(dataset.features[col]) for col in dataset.column_names}
|
| 66 |
+
}
|
| 67 |
+
|
| 68 |
+
print(f"Dataset stats: {stats}")
|
| 69 |
+
return stats
|
| 70 |
+
|
| 71 |
+
|
| 72 |
+
if __name__ == "__main__":
|
| 73 |
+
# Download train split
|
| 74 |
+
download_dataset("train")
|
| 75 |
+
|
| 76 |
+
# Verify
|
| 77 |
+
stats = verify_dataset("train")
|
| 78 |
+
print(f"\nVerification complete:")
|
| 79 |
+
print(f" Total samples: {stats['total_samples']}")
|
| 80 |
+
print(f" Columns: {stats['columns']}")
|
src/data/preprocessing.py
ADDED
|
@@ -0,0 +1,148 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Per-modality preprocessing for satellite imagery.
|
| 3 |
+
|
| 4 |
+
Handles different channel counts:
|
| 5 |
+
- Optical RGB: 3 channels (R, G, B)
|
| 6 |
+
- SAR: 2 channels (VV, VH)
|
| 7 |
+
- Multispectral: 12 channels (Sentinel-2 bands)
|
| 8 |
+
"""
|
| 9 |
+
|
| 10 |
+
import torch
|
| 11 |
+
import torch.nn.functional as F
|
| 12 |
+
from torchvision import transforms
|
| 13 |
+
from PIL import Image
|
| 14 |
+
import numpy as np
|
| 15 |
+
|
| 16 |
+
|
| 17 |
+
# ImageNet normalization for RGB
|
| 18 |
+
IMAGENET_MEAN = [0.485, 0.456, 0.406]
|
| 19 |
+
IMAGENET_STD = [0.229, 0.224, 0.225]
|
| 20 |
+
|
| 21 |
+
# Sentinel-2 band statistics (approximate)
|
| 22 |
+
SENTINEL2_MEAN = [1353.0, 1117.0, 1042.0, 947.0, 1199.0, 1645.0, 1849.0, 1793.0, 1859.0, 1008.0, 1593.0, 1064.0]
|
| 23 |
+
SENTINEL2_STD = [235.0, 309.0, 392.0, 597.0, 490.0, 625.0, 736.0, 755.0, 846.0, 487.0, 561.0, 459.0]
|
| 24 |
+
|
| 25 |
+
# SAR statistics (approximate, in dB)
|
| 26 |
+
SAR_MEAN = [-12.0, -18.0]
|
| 27 |
+
SAR_STD = [5.0, 5.0]
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
def get_optical_transform(size: int = 224) -> transforms.Compose:
|
| 31 |
+
"""Get transforms for optical RGB images."""
|
| 32 |
+
return transforms.Compose([
|
| 33 |
+
transforms.Resize(size),
|
| 34 |
+
transforms.CenterCrop(size),
|
| 35 |
+
transforms.ToTensor(),
|
| 36 |
+
transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
|
| 37 |
+
])
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
def get_sar_transform(size: int = 224) -> transforms.Compose:
|
| 41 |
+
"""Get transforms for SAR images (VV/VH channels)."""
|
| 42 |
+
return transforms.Compose([
|
| 43 |
+
transforms.Resize(size),
|
| 44 |
+
transforms.CenterCrop(size),
|
| 45 |
+
transforms.ToTensor(),
|
| 46 |
+
transforms.Normalize(mean=SAR_MEAN, std=SAR_STD)
|
| 47 |
+
])
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
def get_multispectral_transform(size: int = 224) -> transforms.Compose:
|
| 51 |
+
"""Get transforms for multispectral images (12 channels)."""
|
| 52 |
+
return transforms.Compose([
|
| 53 |
+
transforms.Resize(size),
|
| 54 |
+
transforms.CenterCrop(size),
|
| 55 |
+
transforms.ToTensor(),
|
| 56 |
+
transforms.Normalize(mean=SENTINEL2_MEAN, std=SENTINEL2_STD)
|
| 57 |
+
])
|
| 58 |
+
|
| 59 |
+
|
| 60 |
+
def preprocess_image(
|
| 61 |
+
image: Image.Image,
|
| 62 |
+
modality: str,
|
| 63 |
+
size: int = 224
|
| 64 |
+
) -> torch.Tensor:
|
| 65 |
+
"""
|
| 66 |
+
Preprocess image based on modality.
|
| 67 |
+
|
| 68 |
+
Args:
|
| 69 |
+
image: Input PIL image
|
| 70 |
+
modality: "optical", "sar", or "multispectral"
|
| 71 |
+
size: Output image size
|
| 72 |
+
|
| 73 |
+
Returns:
|
| 74 |
+
Preprocessed tensor
|
| 75 |
+
"""
|
| 76 |
+
# Handle channel mismatch before applying transform
|
| 77 |
+
if modality == "sar":
|
| 78 |
+
# SAR expects 2 channels, but PIL images are typically 3 channels
|
| 79 |
+
# Convert to numpy, take first 2 channels, convert back
|
| 80 |
+
img_array = np.array(image)
|
| 81 |
+
if img_array.shape[-1] == 3:
|
| 82 |
+
img_array = img_array[..., :2]
|
| 83 |
+
image = Image.fromarray(img_array)
|
| 84 |
+
transform = get_sar_transform(size)
|
| 85 |
+
elif modality == "optical":
|
| 86 |
+
transform = get_optical_transform(size)
|
| 87 |
+
elif modality == "multispectral":
|
| 88 |
+
transform = get_multispectral_transform(size)
|
| 89 |
+
else:
|
| 90 |
+
raise ValueError(f"Unknown modality: {modality}")
|
| 91 |
+
|
| 92 |
+
return transform(image)
|
| 93 |
+
|
| 94 |
+
|
| 95 |
+
def handle_channels(
|
| 96 |
+
image: np.ndarray,
|
| 97 |
+
target_channels: int,
|
| 98 |
+
modality: str
|
| 99 |
+
) -> np.ndarray:
|
| 100 |
+
"""
|
| 101 |
+
Handle channel mismatch for different modalities.
|
| 102 |
+
|
| 103 |
+
Args:
|
| 104 |
+
image: Input image array (H, W, C)
|
| 105 |
+
target_channels: Expected number of channels
|
| 106 |
+
modality: Modality type
|
| 107 |
+
|
| 108 |
+
Returns:
|
| 109 |
+
Image with correct number of channels
|
| 110 |
+
"""
|
| 111 |
+
current_channels = image.shape[-1] if len(image.shape) == 3 else 1
|
| 112 |
+
|
| 113 |
+
if current_channels == target_channels:
|
| 114 |
+
return image
|
| 115 |
+
|
| 116 |
+
# ponytail: simple channel handling, not perfect but works for v1
|
| 117 |
+
if modality == "optical" and current_channels >= 3:
|
| 118 |
+
# Take first 3 channels (RGB)
|
| 119 |
+
return image[..., :3]
|
| 120 |
+
elif modality == "sar" and current_channels >= 2:
|
| 121 |
+
# Take first 2 channels (VV, VH)
|
| 122 |
+
return image[..., :2]
|
| 123 |
+
elif modality == "multispectral":
|
| 124 |
+
if current_channels < target_channels:
|
| 125 |
+
# Pad with zeros
|
| 126 |
+
padding = np.zeros((*image.shape[:-1], target_channels - current_channels))
|
| 127 |
+
return np.concatenate([image, padding], axis=-1)
|
| 128 |
+
else:
|
| 129 |
+
# Take first 12 channels
|
| 130 |
+
return image[..., :target_channels]
|
| 131 |
+
|
| 132 |
+
return image
|
| 133 |
+
|
| 134 |
+
|
| 135 |
+
# Self-check
|
| 136 |
+
if __name__ == "__main__":
|
| 137 |
+
# Create dummy images for testing
|
| 138 |
+
dummy_rgb = Image.fromarray(np.random.randint(0, 255, (256, 256, 3), dtype=np.uint8))
|
| 139 |
+
dummy_sar = Image.fromarray(np.random.randint(0, 255, (256, 256, 2), dtype=np.uint8))
|
| 140 |
+
|
| 141 |
+
# Test preprocessing
|
| 142 |
+
optical_tensor = preprocess_image(dummy_rgb, "optical")
|
| 143 |
+
sar_tensor = preprocess_image(dummy_sar, "sar")
|
| 144 |
+
|
| 145 |
+
print(f"Optical shape: {optical_tensor.shape}") # Should be [3, 224, 224]
|
| 146 |
+
print(f"SAR shape: {sar_tensor.shape}") # Should be [2, 224, 224]
|
| 147 |
+
|
| 148 |
+
print("Preprocessing test passed!")
|
src/evaluation/README.md
ADDED
|
@@ -0,0 +1,32 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Evaluation Module
|
| 2 |
+
|
| 3 |
+
Metrics and ground truth management for retrieval evaluation.
|
| 4 |
+
|
| 5 |
+
## Files
|
| 6 |
+
|
| 7 |
+
| File | Description |
|
| 8 |
+
|------|-------------|
|
| 9 |
+
| `metrics.py` | F1@K, precision, recall, timing statistics |
|
| 10 |
+
| `ground_truth.py` | Ground truth label management |
|
| 11 |
+
|
| 12 |
+
## Metrics
|
| 13 |
+
|
| 14 |
+
| Metric | Description |
|
| 15 |
+
|--------|-------------|
|
| 16 |
+
| F1@5 | Precision@5 x Recall@5 |
|
| 17 |
+
| F1@10 | Precision@10 x Recall@10 |
|
| 18 |
+
| Same-Modal F1 | F1 for same-modality queries |
|
| 19 |
+
| Cross-Modal F1 | F1 for cross-modality queries |
|
| 20 |
+
|
| 21 |
+
## Usage
|
| 22 |
+
|
| 23 |
+
```python
|
| 24 |
+
from src.evaluation.metrics import compute_f1
|
| 25 |
+
|
| 26 |
+
# Compute F1@K
|
| 27 |
+
f1_score = compute_f1(
|
| 28 |
+
predicted_indices=result.indices,
|
| 29 |
+
ground_truth_indices=gt_indices,
|
| 30 |
+
k=5
|
| 31 |
+
)
|
| 32 |
+
```
|
src/evaluation/__init__.py
ADDED
|
@@ -0,0 +1,18 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Evaluation module for retrieval performance.
|
| 3 |
+
|
| 4 |
+
Provides:
|
| 5 |
+
- EvaluationMetrics: F1@K and timing statistics
|
| 6 |
+
- GroundTruth: Ground truth label management
|
| 7 |
+
"""
|
| 8 |
+
|
| 9 |
+
from .metrics import EvaluationMetrics, EvaluationResult
|
| 10 |
+
from .ground_truth import GroundTruth, GroundTruthPair, create_ground_truth_from_matches
|
| 11 |
+
|
| 12 |
+
__all__ = [
|
| 13 |
+
"EvaluationMetrics",
|
| 14 |
+
"EvaluationResult",
|
| 15 |
+
"GroundTruth",
|
| 16 |
+
"GroundTruthPair",
|
| 17 |
+
"create_ground_truth_from_matches",
|
| 18 |
+
]
|
src/evaluation/ground_truth.py
ADDED
|
@@ -0,0 +1,219 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Ground truth utilities for evaluation.
|
| 3 |
+
|
| 4 |
+
Handles loading, creating, and validating ground truth labels.
|
| 5 |
+
"""
|
| 6 |
+
|
| 7 |
+
import json
|
| 8 |
+
from pathlib import Path
|
| 9 |
+
from typing import Dict, List, Optional, Tuple
|
| 10 |
+
from dataclasses import dataclass
|
| 11 |
+
|
| 12 |
+
|
| 13 |
+
@dataclass
|
| 14 |
+
class GroundTruthPair:
|
| 15 |
+
"""A query-gallery ground truth pair."""
|
| 16 |
+
query_id: int
|
| 17 |
+
gallery_ids: List[int]
|
| 18 |
+
query_modality: str
|
| 19 |
+
gallery_modality: str
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
class GroundTruth:
|
| 23 |
+
"""
|
| 24 |
+
Ground truth labels for retrieval evaluation.
|
| 25 |
+
|
| 26 |
+
Stores query-gallery pairs for same-modal and cross-modal evaluation.
|
| 27 |
+
"""
|
| 28 |
+
|
| 29 |
+
def __init__(self):
|
| 30 |
+
"""Initialize empty ground truth."""
|
| 31 |
+
self.pairs: List[GroundTruthPair] = []
|
| 32 |
+
self._query_to_gallery: Dict[int, List[int]] = {}
|
| 33 |
+
|
| 34 |
+
def add_pair(
|
| 35 |
+
self,
|
| 36 |
+
query_id: int,
|
| 37 |
+
gallery_ids: List[int],
|
| 38 |
+
query_modality: str,
|
| 39 |
+
gallery_modality: str
|
| 40 |
+
) -> None:
|
| 41 |
+
"""
|
| 42 |
+
Add a ground truth pair.
|
| 43 |
+
|
| 44 |
+
Args:
|
| 45 |
+
query_id: Query sample ID
|
| 46 |
+
gallery_ids: List of matching gallery IDs
|
| 47 |
+
query_modality: Modality of query
|
| 48 |
+
gallery_modality: Modality of gallery
|
| 49 |
+
"""
|
| 50 |
+
pair = GroundTruthPair(
|
| 51 |
+
query_id=query_id,
|
| 52 |
+
gallery_ids=gallery_ids,
|
| 53 |
+
query_modality=query_modality,
|
| 54 |
+
gallery_modality=gallery_modality
|
| 55 |
+
)
|
| 56 |
+
self.pairs.append(pair)
|
| 57 |
+
self._query_to_gallery[query_id] = gallery_ids
|
| 58 |
+
|
| 59 |
+
def get_gallery_ids(self, query_id: int) -> List[int]:
|
| 60 |
+
"""
|
| 61 |
+
Get ground truth gallery IDs for a query.
|
| 62 |
+
|
| 63 |
+
Args:
|
| 64 |
+
query_id: Query sample ID
|
| 65 |
+
|
| 66 |
+
Returns:
|
| 67 |
+
List of matching gallery IDs
|
| 68 |
+
"""
|
| 69 |
+
return self._query_to_gallery.get(query_id, [])
|
| 70 |
+
|
| 71 |
+
@property
|
| 72 |
+
def n_pairs(self) -> int:
|
| 73 |
+
"""Number of ground truth pairs."""
|
| 74 |
+
return len(self.pairs)
|
| 75 |
+
|
| 76 |
+
def save(self, path: str) -> None:
|
| 77 |
+
"""
|
| 78 |
+
Save ground truth to JSON.
|
| 79 |
+
|
| 80 |
+
Args:
|
| 81 |
+
path: Output path
|
| 82 |
+
"""
|
| 83 |
+
data = {
|
| 84 |
+
"pairs": [
|
| 85 |
+
{
|
| 86 |
+
"query_id": p.query_id,
|
| 87 |
+
"gallery_ids": p.gallery_ids,
|
| 88 |
+
"query_modality": p.query_modality,
|
| 89 |
+
"gallery_modality": p.gallery_modality,
|
| 90 |
+
}
|
| 91 |
+
for p in self.pairs
|
| 92 |
+
]
|
| 93 |
+
}
|
| 94 |
+
|
| 95 |
+
with open(path, "w") as f:
|
| 96 |
+
json.dump(data, f, indent=2)
|
| 97 |
+
|
| 98 |
+
@classmethod
|
| 99 |
+
def load(cls, path: str) -> "GroundTruth":
|
| 100 |
+
"""
|
| 101 |
+
Load ground truth from JSON.
|
| 102 |
+
|
| 103 |
+
Args:
|
| 104 |
+
path: Input path
|
| 105 |
+
|
| 106 |
+
Returns:
|
| 107 |
+
GroundTruth instance
|
| 108 |
+
"""
|
| 109 |
+
with open(path, "r") as f:
|
| 110 |
+
data = json.load(f)
|
| 111 |
+
|
| 112 |
+
gt = cls()
|
| 113 |
+
for pair_data in data["pairs"]:
|
| 114 |
+
gt.add_pair(**pair_data)
|
| 115 |
+
|
| 116 |
+
return gt
|
| 117 |
+
|
| 118 |
+
def validate(self) -> bool:
|
| 119 |
+
"""
|
| 120 |
+
Validate ground truth structure.
|
| 121 |
+
|
| 122 |
+
Returns:
|
| 123 |
+
True if valid
|
| 124 |
+
"""
|
| 125 |
+
for pair in self.pairs:
|
| 126 |
+
if not pair.gallery_ids:
|
| 127 |
+
print(f"Warning: Query {pair.query_id} has no gallery matches")
|
| 128 |
+
return False
|
| 129 |
+
|
| 130 |
+
if pair.query_modality not in ["optical", "sar", "multispectral"]:
|
| 131 |
+
print(f"Invalid query modality: {pair.query_modality}")
|
| 132 |
+
return False
|
| 133 |
+
|
| 134 |
+
if pair.gallery_modality not in ["optical", "sar", "multispectral"]:
|
| 135 |
+
print(f"Invalid gallery modality: {pair.gallery_modality}")
|
| 136 |
+
return False
|
| 137 |
+
|
| 138 |
+
return True
|
| 139 |
+
|
| 140 |
+
def get_same_modal_pairs(self) -> List[GroundTruthPair]:
|
| 141 |
+
"""Get pairs where query and gallery are same modality."""
|
| 142 |
+
return [
|
| 143 |
+
p for p in self.pairs
|
| 144 |
+
if p.query_modality == p.gallery_modality
|
| 145 |
+
]
|
| 146 |
+
|
| 147 |
+
def get_cross_modal_pairs(self) -> List[GroundTruthPair]:
|
| 148 |
+
"""Get pairs where query and gallery are different modalities."""
|
| 149 |
+
return [
|
| 150 |
+
p for p in self.pairs
|
| 151 |
+
if p.query_modality != p.gallery_modality
|
| 152 |
+
]
|
| 153 |
+
|
| 154 |
+
|
| 155 |
+
def create_ground_truth_from_matches(
|
| 156 |
+
matches: Dict[Tuple[str, str], List[Tuple[int, int]]]
|
| 157 |
+
) -> GroundTruth:
|
| 158 |
+
"""
|
| 159 |
+
Create ground truth from modality matches.
|
| 160 |
+
|
| 161 |
+
Args:
|
| 162 |
+
matches: Dict mapping (query_modality, gallery_modality) to
|
| 163 |
+
list of (query_id, gallery_id) pairs
|
| 164 |
+
|
| 165 |
+
Returns:
|
| 166 |
+
GroundTruth instance
|
| 167 |
+
"""
|
| 168 |
+
gt = GroundTruth()
|
| 169 |
+
|
| 170 |
+
for (query_mod, gallery_mod), pairs in matches.items():
|
| 171 |
+
# Group by query_id
|
| 172 |
+
query_to_galleries: Dict[int, List[int]] = {}
|
| 173 |
+
for query_id, gallery_id in pairs:
|
| 174 |
+
if query_id not in query_to_galleries:
|
| 175 |
+
query_to_galleries[query_id] = []
|
| 176 |
+
query_to_galleries[query_id].append(gallery_id)
|
| 177 |
+
|
| 178 |
+
# Add to ground truth
|
| 179 |
+
for query_id, gallery_ids in query_to_galleries.items():
|
| 180 |
+
gt.add_pair(query_id, gallery_ids, query_mod, gallery_mod)
|
| 181 |
+
|
| 182 |
+
return gt
|
| 183 |
+
|
| 184 |
+
|
| 185 |
+
# Self-check
|
| 186 |
+
if __name__ == "__main__":
|
| 187 |
+
import tempfile
|
| 188 |
+
|
| 189 |
+
print("Testing GroundTruth utilities...")
|
| 190 |
+
|
| 191 |
+
# Create ground truth
|
| 192 |
+
gt = GroundTruth()
|
| 193 |
+
|
| 194 |
+
# Same-modal pairs
|
| 195 |
+
gt.add_pair(0, [1, 2], "optical", "optical")
|
| 196 |
+
gt.add_pair(1, [3, 4], "sar", "sar")
|
| 197 |
+
|
| 198 |
+
# Cross-modal pairs
|
| 199 |
+
gt.add_pair(2, [5, 6], "optical", "sar")
|
| 200 |
+
gt.add_pair(3, [7, 8], "sar", "optical")
|
| 201 |
+
|
| 202 |
+
print(f"Total pairs: {gt.n_pairs}")
|
| 203 |
+
print(f"Same-modal: {len(gt.get_same_modal_pairs())}")
|
| 204 |
+
print(f"Cross-modal: {len(gt.get_cross_modal_pairs())}")
|
| 205 |
+
|
| 206 |
+
# Validate
|
| 207 |
+
assert gt.validate(), "Validation failed"
|
| 208 |
+
|
| 209 |
+
# Save/load roundtrip
|
| 210 |
+
with tempfile.TemporaryDirectory() as tmpdir:
|
| 211 |
+
save_path = Path(tmpdir) / "ground_truth.json"
|
| 212 |
+
gt.save(save_path)
|
| 213 |
+
|
| 214 |
+
loaded_gt = GroundTruth.load(save_path)
|
| 215 |
+
|
| 216 |
+
assert loaded_gt.n_pairs == gt.n_pairs
|
| 217 |
+
assert loaded_gt.validate()
|
| 218 |
+
|
| 219 |
+
print("\nGroundTruth test passed!")
|
src/evaluation/metrics.py
ADDED
|
@@ -0,0 +1,242 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Evaluation metrics for retrieval performance.
|
| 3 |
+
|
| 4 |
+
Computes F1@K, timing statistics, and per-modality breakdowns.
|
| 5 |
+
"""
|
| 6 |
+
|
| 7 |
+
import numpy as np
|
| 8 |
+
from typing import List, Dict, Optional, Tuple
|
| 9 |
+
from dataclasses import dataclass
|
| 10 |
+
|
| 11 |
+
|
| 12 |
+
@dataclass
|
| 13 |
+
class EvaluationResult:
|
| 14 |
+
"""Result of evaluation."""
|
| 15 |
+
f1_at_5: float
|
| 16 |
+
f1_at_10: float
|
| 17 |
+
mean_time_ms: float
|
| 18 |
+
median_time_ms: float
|
| 19 |
+
p95_time_ms: float
|
| 20 |
+
p99_time_ms: float
|
| 21 |
+
modality_results: Dict[str, Dict[str, float]]
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
class EvaluationMetrics:
|
| 25 |
+
"""
|
| 26 |
+
Evaluation metrics for retrieval systems.
|
| 27 |
+
|
| 28 |
+
Computes F1@K and timing statistics.
|
| 29 |
+
"""
|
| 30 |
+
|
| 31 |
+
def __init__(self):
|
| 32 |
+
"""Initialize evaluation metrics."""
|
| 33 |
+
self._query_times: List[float] = []
|
| 34 |
+
|
| 35 |
+
def compute_f1_at_k(
|
| 36 |
+
self,
|
| 37 |
+
predicted_indices: List[int],
|
| 38 |
+
ground_truth_indices: List[int],
|
| 39 |
+
k: int = 5
|
| 40 |
+
) -> float:
|
| 41 |
+
"""
|
| 42 |
+
Compute F1@K between predicted and ground truth indices.
|
| 43 |
+
|
| 44 |
+
Args:
|
| 45 |
+
predicted_indices: Predicted indices (ranked)
|
| 46 |
+
ground_truth_indices: Ground truth indices
|
| 47 |
+
k: Top-K to consider
|
| 48 |
+
|
| 49 |
+
Returns:
|
| 50 |
+
F1 score (0-1)
|
| 51 |
+
"""
|
| 52 |
+
# Take top-k predictions
|
| 53 |
+
predicted_top_k = set(predicted_indices[:k])
|
| 54 |
+
ground_truth_set = set(ground_truth_indices)
|
| 55 |
+
|
| 56 |
+
# Compute precision and recall
|
| 57 |
+
if len(predicted_top_k) == 0:
|
| 58 |
+
precision = 0.0
|
| 59 |
+
else:
|
| 60 |
+
relevant_predicted = predicted_top_k & ground_truth_set
|
| 61 |
+
precision = len(relevant_predicted) / len(predicted_top_k)
|
| 62 |
+
|
| 63 |
+
if len(ground_truth_set) == 0:
|
| 64 |
+
recall = 0.0
|
| 65 |
+
else:
|
| 66 |
+
relevant_predicted = predicted_top_k & ground_truth_set
|
| 67 |
+
recall = len(relevant_predicted) / len(ground_truth_set)
|
| 68 |
+
|
| 69 |
+
# Compute F1
|
| 70 |
+
if precision + recall == 0:
|
| 71 |
+
f1 = 0.0
|
| 72 |
+
else:
|
| 73 |
+
f1 = 2 * precision * recall / (precision + recall)
|
| 74 |
+
|
| 75 |
+
return f1
|
| 76 |
+
|
| 77 |
+
def compute_f1_at_k_batch(
|
| 78 |
+
self,
|
| 79 |
+
all_predicted: List[List[int]],
|
| 80 |
+
all_ground_truth: List[List[int]],
|
| 81 |
+
k: int = 5
|
| 82 |
+
) -> float:
|
| 83 |
+
"""
|
| 84 |
+
Compute average F1@K over a batch.
|
| 85 |
+
|
| 86 |
+
Args:
|
| 87 |
+
all_predicted: List of predicted indices per query
|
| 88 |
+
all_ground_truth: List of ground truth indices per query
|
| 89 |
+
k: Top-K to consider
|
| 90 |
+
|
| 91 |
+
Returns:
|
| 92 |
+
Average F1 score
|
| 93 |
+
"""
|
| 94 |
+
if len(all_predicted) == 0:
|
| 95 |
+
return 0.0
|
| 96 |
+
|
| 97 |
+
f1_scores = [
|
| 98 |
+
self.compute_f1_at_k(pred, gt, k)
|
| 99 |
+
for pred, gt in zip(all_predicted, all_ground_truth)
|
| 100 |
+
]
|
| 101 |
+
|
| 102 |
+
return sum(f1_scores) / len(f1_scores)
|
| 103 |
+
|
| 104 |
+
def add_query_time(self, time_ms: float) -> None:
|
| 105 |
+
"""Add a query time measurement."""
|
| 106 |
+
self._query_times.append(time_ms)
|
| 107 |
+
|
| 108 |
+
def get_timing_stats(self) -> Dict[str, float]:
|
| 109 |
+
"""
|
| 110 |
+
Get timing statistics.
|
| 111 |
+
|
| 112 |
+
Returns:
|
| 113 |
+
Dict with mean, median, p95, p99
|
| 114 |
+
"""
|
| 115 |
+
if not self._query_times:
|
| 116 |
+
return {
|
| 117 |
+
"mean_ms": 0.0,
|
| 118 |
+
"median_ms": 0.0,
|
| 119 |
+
"p95_ms": 0.0,
|
| 120 |
+
"p99_ms": 0.0,
|
| 121 |
+
}
|
| 122 |
+
|
| 123 |
+
times = sorted(self._query_times)
|
| 124 |
+
n = len(times)
|
| 125 |
+
|
| 126 |
+
return {
|
| 127 |
+
"mean_ms": sum(times) / n,
|
| 128 |
+
"median_ms": times[n // 2],
|
| 129 |
+
"p95_ms": times[int(n * 0.95)] if n >= 20 else times[-1],
|
| 130 |
+
"p99_ms": times[int(n * 0.99)] if n >= 100 else times[-1],
|
| 131 |
+
}
|
| 132 |
+
|
| 133 |
+
def evaluate_retrieval(
|
| 134 |
+
self,
|
| 135 |
+
predicted_indices_list: List[List[int]],
|
| 136 |
+
ground_truth_list: List[List[int]],
|
| 137 |
+
query_times: Optional[List[float]] = None,
|
| 138 |
+
modality_labels: Optional[List[str]] = None
|
| 139 |
+
) -> EvaluationResult:
|
| 140 |
+
"""
|
| 141 |
+
Full evaluation of retrieval results.
|
| 142 |
+
|
| 143 |
+
Args:
|
| 144 |
+
predicted_indices_list: List of predicted indices per query
|
| 145 |
+
ground_truth_list: List of ground truth indices per query
|
| 146 |
+
query_times: Optional query times in ms
|
| 147 |
+
modality_labels: Optional modality labels per query
|
| 148 |
+
|
| 149 |
+
Returns:
|
| 150 |
+
EvaluationResult with all metrics
|
| 151 |
+
"""
|
| 152 |
+
# Compute F1@5 and F1@10
|
| 153 |
+
f1_at_5 = self.compute_f1_at_k_batch(
|
| 154 |
+
predicted_indices_list, ground_truth_list, k=5
|
| 155 |
+
)
|
| 156 |
+
f1_at_10 = self.compute_f1_at_k_batch(
|
| 157 |
+
predicted_indices_list, ground_truth_list, k=10
|
| 158 |
+
)
|
| 159 |
+
|
| 160 |
+
# Timing stats
|
| 161 |
+
if query_times:
|
| 162 |
+
self._query_times.extend(query_times)
|
| 163 |
+
|
| 164 |
+
timing = self.get_timing_stats()
|
| 165 |
+
|
| 166 |
+
# Per-modality breakdown
|
| 167 |
+
modality_results = {}
|
| 168 |
+
if modality_labels:
|
| 169 |
+
unique_modalities = set(modality_labels)
|
| 170 |
+
|
| 171 |
+
for mod in unique_modalities:
|
| 172 |
+
# Filter to this modality
|
| 173 |
+
mod_indices = [
|
| 174 |
+
i for i, m in enumerate(modality_labels)
|
| 175 |
+
if m == mod
|
| 176 |
+
]
|
| 177 |
+
|
| 178 |
+
if mod_indices:
|
| 179 |
+
mod_predicted = [predicted_indices_list[i] for i in mod_indices]
|
| 180 |
+
mod_gt = [ground_truth_list[i] for i in mod_indices]
|
| 181 |
+
|
| 182 |
+
mod_f1_5 = self.compute_f1_at_k_batch(mod_predicted, mod_gt, k=5)
|
| 183 |
+
mod_f1_10 = self.compute_f1_at_k_batch(mod_predicted, mod_gt, k=10)
|
| 184 |
+
|
| 185 |
+
modality_results[mod] = {
|
| 186 |
+
"f1_at_5": mod_f1_5,
|
| 187 |
+
"f1_at_10": mod_f1_10,
|
| 188 |
+
"n_queries": len(mod_indices),
|
| 189 |
+
}
|
| 190 |
+
|
| 191 |
+
return EvaluationResult(
|
| 192 |
+
f1_at_5=f1_at_5,
|
| 193 |
+
f1_at_10=f1_at_10,
|
| 194 |
+
mean_time_ms=timing["mean_ms"],
|
| 195 |
+
median_time_ms=timing["median_ms"],
|
| 196 |
+
p95_time_ms=timing["p95_ms"],
|
| 197 |
+
p99_time_ms=timing["p99_ms"],
|
| 198 |
+
modality_results=modality_results,
|
| 199 |
+
)
|
| 200 |
+
|
| 201 |
+
|
| 202 |
+
# Self-check
|
| 203 |
+
if __name__ == "__main__":
|
| 204 |
+
print("Testing EvaluationMetrics...")
|
| 205 |
+
|
| 206 |
+
metrics = EvaluationMetrics()
|
| 207 |
+
|
| 208 |
+
# Test F1 computation
|
| 209 |
+
predicted = [0, 1, 2, 3, 4]
|
| 210 |
+
ground_truth = [0, 2, 4, 6, 8]
|
| 211 |
+
|
| 212 |
+
f1_5 = metrics.compute_f1_at_k(predicted, ground_truth, k=5)
|
| 213 |
+
print(f"F1@5: {f1_5:.4f}")
|
| 214 |
+
|
| 215 |
+
# Test batch F1
|
| 216 |
+
all_predicted = [[0, 1, 2], [3, 4, 5]]
|
| 217 |
+
all_gt = [[0, 1, 2], [3, 4, 5]]
|
| 218 |
+
|
| 219 |
+
batch_f1 = metrics.compute_f1_at_k_batch(all_predicted, all_gt, k=3)
|
| 220 |
+
print(f"Batch F1@3: {batch_f1:.4f}")
|
| 221 |
+
|
| 222 |
+
# Test timing
|
| 223 |
+
for t in [10.0, 20.0, 30.0, 40.0, 50.0]:
|
| 224 |
+
metrics.add_query_time(t)
|
| 225 |
+
|
| 226 |
+
timing = metrics.get_timing_stats()
|
| 227 |
+
print(f"Timing stats: {timing}")
|
| 228 |
+
|
| 229 |
+
# Test full evaluation
|
| 230 |
+
result = metrics.evaluate_retrieval(
|
| 231 |
+
all_predicted, all_gt,
|
| 232 |
+
query_times=[15.0, 25.0],
|
| 233 |
+
modality_labels=["optical", "sar"]
|
| 234 |
+
)
|
| 235 |
+
|
| 236 |
+
print(f"\nFull evaluation:")
|
| 237 |
+
print(f" F1@5: {result.f1_at_5:.4f}")
|
| 238 |
+
print(f" F1@10: {result.f1_at_10:.4f}")
|
| 239 |
+
print(f" Mean time: {result.mean_time_ms:.2f}ms")
|
| 240 |
+
print(f" Modality results: {result.modality_results}")
|
| 241 |
+
|
| 242 |
+
print("\nEvaluationMetrics test passed!")
|
src/features/README.md
ADDED
|
@@ -0,0 +1,34 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Features Module
|
| 2 |
+
|
| 3 |
+
Feature extraction using CLIP for satellite imagery.
|
| 4 |
+
|
| 5 |
+
## Files
|
| 6 |
+
|
| 7 |
+
| File | Description |
|
| 8 |
+
|------|-------------|
|
| 9 |
+
| `extractor.py` | CLIP ViT-L/14 feature extractor |
|
| 10 |
+
| `embeddings.py` | Embedding cache utilities |
|
| 11 |
+
|
| 12 |
+
## Model
|
| 13 |
+
|
| 14 |
+
| Parameter | Value |
|
| 15 |
+
|-----------|-------|
|
| 16 |
+
| Model | `openai/clip-vit-large-patch14` |
|
| 17 |
+
| Architecture | Vision Transformer |
|
| 18 |
+
| Embedding Dim | 768 |
|
| 19 |
+
| Input Resolution | 224x224 |
|
| 20 |
+
|
| 21 |
+
## Usage
|
| 22 |
+
|
| 23 |
+
```python
|
| 24 |
+
from src.features.extractor import FeatureExtractor
|
| 25 |
+
|
| 26 |
+
# Initialize extractor
|
| 27 |
+
extractor = FeatureExtractor(model_name="openai/clip-vit-large-patch14")
|
| 28 |
+
|
| 29 |
+
# Extract features from single image
|
| 30 |
+
embedding = extractor.extract_features(image, modality="optical")
|
| 31 |
+
|
| 32 |
+
# Extract features from batch
|
| 33 |
+
embeddings = extractor.extract_batch(images, modality="optical", batch_size=32)
|
| 34 |
+
```
|
src/features/__init__.py
ADDED
|
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Feature extraction module for satellite imagery.
|
| 3 |
+
|
| 4 |
+
Provides:
|
| 5 |
+
- FeatureExtractor: Extract embeddings using DOFA-CLIP
|
| 6 |
+
- Embedding cache utilities
|
| 7 |
+
"""
|
| 8 |
+
|
| 9 |
+
from .extractor import FeatureExtractor
|
| 10 |
+
from .embeddings import (
|
| 11 |
+
save_embeddings,
|
| 12 |
+
load_embeddings,
|
| 13 |
+
get_cache_path,
|
| 14 |
+
verify_embeddings,
|
| 15 |
+
)
|
| 16 |
+
|
| 17 |
+
__all__ = [
|
| 18 |
+
"FeatureExtractor",
|
| 19 |
+
"save_embeddings",
|
| 20 |
+
"load_embeddings",
|
| 21 |
+
"get_cache_path",
|
| 22 |
+
"verify_embeddings",
|
| 23 |
+
]
|
src/features/cross_modal.py
ADDED
|
@@ -0,0 +1,341 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Cross-modal alignment for satellite imagery retrieval.
|
| 3 |
+
|
| 4 |
+
Implements multiple approaches:
|
| 5 |
+
1. Modality-specific projection heads
|
| 6 |
+
2. Contrastive cross-modal loss
|
| 7 |
+
3. Wavelength-aware encoding
|
| 8 |
+
4. Domain adaptation
|
| 9 |
+
"""
|
| 10 |
+
|
| 11 |
+
import torch
|
| 12 |
+
import torch.nn as nn
|
| 13 |
+
import torch.nn.functional as F
|
| 14 |
+
from typing import Dict, List, Optional, Tuple
|
| 15 |
+
from dataclasses import dataclass
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
@dataclass
|
| 19 |
+
class CrossModalConfig:
|
| 20 |
+
"""Configuration for cross-modal alignment."""
|
| 21 |
+
embed_dim: int = 768
|
| 22 |
+
projection_dim: int = 256
|
| 23 |
+
modalities: List[str] = None
|
| 24 |
+
temperature: float = 0.07
|
| 25 |
+
use_wavelength_encoding: bool = True
|
| 26 |
+
use_domain_adaptation: bool = True
|
| 27 |
+
|
| 28 |
+
def __post_init__(self):
|
| 29 |
+
if self.modalities is None:
|
| 30 |
+
self.modalities = ["optical", "sar", "multispectral"]
|
| 31 |
+
|
| 32 |
+
|
| 33 |
+
class ModalityProjectionHead(nn.Module):
|
| 34 |
+
"""Projection head for a single modality."""
|
| 35 |
+
|
| 36 |
+
def __init__(self, input_dim: int, output_dim: int):
|
| 37 |
+
super().__init__()
|
| 38 |
+
self.projection = nn.Sequential(
|
| 39 |
+
nn.Linear(input_dim, input_dim),
|
| 40 |
+
nn.GELU(),
|
| 41 |
+
nn.Linear(input_dim, output_dim),
|
| 42 |
+
)
|
| 43 |
+
|
| 44 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
| 45 |
+
return F.normalize(self.projection(x), dim=-1)
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
class WavelengthEncoder(nn.Module):
|
| 49 |
+
"""Encode wavelength information for each modality."""
|
| 50 |
+
|
| 51 |
+
def __init__(self, num_channels: int, output_dim: int):
|
| 52 |
+
super().__init__()
|
| 53 |
+
self.encoder = nn.Sequential(
|
| 54 |
+
nn.Linear(num_channels, output_dim),
|
| 55 |
+
nn.GELU(),
|
| 56 |
+
nn.Linear(output_dim, output_dim),
|
| 57 |
+
)
|
| 58 |
+
|
| 59 |
+
def forward(self, wavelengths: torch.Tensor) -> torch.Tensor:
|
| 60 |
+
return self.encoder(wavelengths)
|
| 61 |
+
|
| 62 |
+
|
| 63 |
+
class CrossModalAligner(nn.Module):
|
| 64 |
+
"""
|
| 65 |
+
Cross-modal alignment using modality-specific projections.
|
| 66 |
+
|
| 67 |
+
Based on CLOSP and DOFA-CLIP approaches:
|
| 68 |
+
- Each modality has its own projection head
|
| 69 |
+
- Wavelength encoding for channel-aware processing
|
| 70 |
+
- Contrastive loss for alignment
|
| 71 |
+
"""
|
| 72 |
+
|
| 73 |
+
def __init__(self, config: CrossModalConfig):
|
| 74 |
+
super().__init__()
|
| 75 |
+
self.config = config
|
| 76 |
+
|
| 77 |
+
# Modality-specific projection heads
|
| 78 |
+
self.projection_heads = nn.ModuleDict({
|
| 79 |
+
mod: ModalityProjectionHead(config.embed_dim, config.projection_dim)
|
| 80 |
+
for mod in config.modalities
|
| 81 |
+
})
|
| 82 |
+
|
| 83 |
+
# Wavelength encoders (if enabled)
|
| 84 |
+
if config.use_wavelength_encoding:
|
| 85 |
+
self.wavelength_encoders = nn.ModuleDict({
|
| 86 |
+
mod: WavelengthEncoder(3, config.projection_dim) # 3 channels for wavelength
|
| 87 |
+
for mod in config.modalities
|
| 88 |
+
})
|
| 89 |
+
|
| 90 |
+
# Domain adaptation layer (if enabled)
|
| 91 |
+
if config.use_domain_adaptation:
|
| 92 |
+
self.domain_adaptor = nn.Sequential(
|
| 93 |
+
nn.Linear(config.projection_dim, config.projection_dim),
|
| 94 |
+
nn.GELU(),
|
| 95 |
+
nn.Linear(config.projection_dim, config.projection_dim),
|
| 96 |
+
)
|
| 97 |
+
|
| 98 |
+
# Learnable temperature
|
| 99 |
+
self.logit_scale = nn.Parameter(torch.ones([]) * torch.log(torch.tensor(1.0 / config.temperature)))
|
| 100 |
+
|
| 101 |
+
def project(self, features: torch.Tensor, modality: str) -> torch.Tensor:
|
| 102 |
+
"""Project features using modality-specific head."""
|
| 103 |
+
return self.projection_heads[modality](features)
|
| 104 |
+
|
| 105 |
+
def align_with_wavelength(
|
| 106 |
+
self,
|
| 107 |
+
features: torch.Tensor,
|
| 108 |
+
modality: str,
|
| 109 |
+
wavelengths: Optional[torch.Tensor] = None
|
| 110 |
+
) -> torch.Tensor:
|
| 111 |
+
"""Align features using wavelength encoding."""
|
| 112 |
+
if not self.config.use_wavelength_encoding or wavelengths is None:
|
| 113 |
+
return self.project(features, modality)
|
| 114 |
+
|
| 115 |
+
# Get wavelength embedding
|
| 116 |
+
wave_emb = self.wavelength_encoders[modality](wavelengths)
|
| 117 |
+
|
| 118 |
+
# Combine features with wavelength info
|
| 119 |
+
combined = features + wave_emb
|
| 120 |
+
return self.projection_heads[modality](combined)
|
| 121 |
+
|
| 122 |
+
def contrastive_loss(
|
| 123 |
+
self,
|
| 124 |
+
features_a: torch.Tensor,
|
| 125 |
+
features_b: torch.Tensor,
|
| 126 |
+
temperature: Optional[float] = None
|
| 127 |
+
) -> torch.Tensor:
|
| 128 |
+
"""Compute contrastive loss between two sets of features."""
|
| 129 |
+
if temperature is None:
|
| 130 |
+
temperature = self.config.temperature
|
| 131 |
+
|
| 132 |
+
# Normalize
|
| 133 |
+
features_a = F.normalize(features_a, dim=-1)
|
| 134 |
+
features_b = F.normalize(features_b, dim=-1)
|
| 135 |
+
|
| 136 |
+
# Compute similarity
|
| 137 |
+
logit_scale = self.logit_scale.exp()
|
| 138 |
+
logits = logit_scale * features_a @ features_b.t()
|
| 139 |
+
|
| 140 |
+
# Labels (diagonal is positive)
|
| 141 |
+
labels = torch.arange(len(features_a), device=features_a.device)
|
| 142 |
+
|
| 143 |
+
# Symmetric loss
|
| 144 |
+
loss_a = F.cross_entropy(logits, labels)
|
| 145 |
+
loss_b = F.cross_entropy(logits.t(), labels)
|
| 146 |
+
|
| 147 |
+
return (loss_a + loss_b) / 2
|
| 148 |
+
|
| 149 |
+
def cross_modal_retrieve(
|
| 150 |
+
self,
|
| 151 |
+
query_features: torch.Tensor,
|
| 152 |
+
query_modality: str,
|
| 153 |
+
gallery_features: Dict[str, torch.Tensor],
|
| 154 |
+
k: int = 5
|
| 155 |
+
) -> Tuple[torch.Tensor, torch.Tensor]:
|
| 156 |
+
"""
|
| 157 |
+
Cross-modal retrieval.
|
| 158 |
+
|
| 159 |
+
Args:
|
| 160 |
+
query_features: Query features from source modality
|
| 161 |
+
query_modality: Modality of query
|
| 162 |
+
gallery_features: Dict of gallery features per modality
|
| 163 |
+
k: Number of results
|
| 164 |
+
|
| 165 |
+
Returns:
|
| 166 |
+
(indices, scores) for top-k results
|
| 167 |
+
"""
|
| 168 |
+
# Project query
|
| 169 |
+
query_proj = self.project(query_features, query_modality)
|
| 170 |
+
|
| 171 |
+
all_scores = []
|
| 172 |
+
all_indices = []
|
| 173 |
+
|
| 174 |
+
# Search across all target modalities
|
| 175 |
+
offset = 0
|
| 176 |
+
for mod, features in gallery_features.items():
|
| 177 |
+
# Project gallery
|
| 178 |
+
gallery_proj = self.project(features, mod)
|
| 179 |
+
|
| 180 |
+
# Compute similarity
|
| 181 |
+
scores = query_proj @ gallery_proj.t()
|
| 182 |
+
all_scores.append(scores)
|
| 183 |
+
all_indices.append(torch.arange(len(features), device=features.device) + offset)
|
| 184 |
+
offset += len(features)
|
| 185 |
+
|
| 186 |
+
# Concatenate
|
| 187 |
+
all_scores = torch.cat(all_scores, dim=-1)
|
| 188 |
+
all_indices = torch.cat(all_indices, dim=-1)
|
| 189 |
+
|
| 190 |
+
# Top-k
|
| 191 |
+
topk_scores, topk_idx = all_scores.topk(k, dim=-1)
|
| 192 |
+
topk_indices = all_indices[topk_idx]
|
| 193 |
+
|
| 194 |
+
return topk_indices, topk_scores
|
| 195 |
+
|
| 196 |
+
|
| 197 |
+
class ContrastiveCrossModalLoss(nn.Module):
|
| 198 |
+
"""
|
| 199 |
+
Contrastive loss for cross-modal alignment.
|
| 200 |
+
|
| 201 |
+
Based on CLOSP approach: align SAR and optical via shared text anchor.
|
| 202 |
+
"""
|
| 203 |
+
|
| 204 |
+
def __init__(self, temperature: float = 0.07):
|
| 205 |
+
super().__init__()
|
| 206 |
+
self.temperature = temperature
|
| 207 |
+
self.logit_scale = nn.Parameter(torch.ones([]) * torch.log(torch.tensor(1.0 / temperature)))
|
| 208 |
+
|
| 209 |
+
def forward(
|
| 210 |
+
self,
|
| 211 |
+
features_a: torch.Tensor,
|
| 212 |
+
features_b: torch.Tensor,
|
| 213 |
+
features_text: Optional[torch.Tensor] = None
|
| 214 |
+
) -> torch.Tensor:
|
| 215 |
+
"""
|
| 216 |
+
Compute cross-modal contrastive loss.
|
| 217 |
+
|
| 218 |
+
Args:
|
| 219 |
+
features_a: Features from modality A (e.g., optical)
|
| 220 |
+
features_b: Features from modality B (e.g., SAR)
|
| 221 |
+
features_text: Optional text features for triple alignment
|
| 222 |
+
|
| 223 |
+
Returns:
|
| 224 |
+
Loss value
|
| 225 |
+
"""
|
| 226 |
+
features_a = F.normalize(features_a, dim=-1)
|
| 227 |
+
features_b = F.normalize(features_b, dim=-1)
|
| 228 |
+
|
| 229 |
+
logit_scale = self.logit_scale.exp()
|
| 230 |
+
|
| 231 |
+
# Image-image contrastive loss
|
| 232 |
+
logits_ab = logit_scale * features_a @ features_b.t()
|
| 233 |
+
logits_ba = logits_ab.t()
|
| 234 |
+
|
| 235 |
+
labels = torch.arange(len(features_a), device=features_a.device)
|
| 236 |
+
|
| 237 |
+
loss_a2b = F.cross_entropy(logits_ab, labels)
|
| 238 |
+
loss_b2a = F.cross_entropy(logits_ba, labels)
|
| 239 |
+
|
| 240 |
+
loss = (loss_a2b + loss_b2a) / 2
|
| 241 |
+
|
| 242 |
+
# Text-image contrastive loss (if available)
|
| 243 |
+
if features_text is not None:
|
| 244 |
+
features_text = F.normalize(features_text, dim=-1)
|
| 245 |
+
|
| 246 |
+
logits_t2a = logit_scale * features_text @ features_a.t()
|
| 247 |
+
logits_t2b = logit_scale * features_text @ features_b.t()
|
| 248 |
+
|
| 249 |
+
loss_t2a = F.cross_entropy(logits_t2a, labels)
|
| 250 |
+
loss_t2b = F.cross_entropy(logits_t2b, labels)
|
| 251 |
+
|
| 252 |
+
loss = loss + (loss_t2a + loss_t2b) / 2
|
| 253 |
+
|
| 254 |
+
return loss
|
| 255 |
+
|
| 256 |
+
|
| 257 |
+
class DomainAdaptationLayer(nn.Module):
|
| 258 |
+
"""
|
| 259 |
+
Domain adaptation for bridging modality gaps.
|
| 260 |
+
|
| 261 |
+
Based on SARCLIP approach: transfer knowledge from optical to SAR.
|
| 262 |
+
"""
|
| 263 |
+
|
| 264 |
+
def __init__(self, embed_dim: int, num_modalities: int = 3):
|
| 265 |
+
super().__init__()
|
| 266 |
+
|
| 267 |
+
# Modality-specific adapters
|
| 268 |
+
self.adapters = nn.ModuleList([
|
| 269 |
+
nn.Sequential(
|
| 270 |
+
nn.Linear(embed_dim, embed_dim),
|
| 271 |
+
nn.GELU(),
|
| 272 |
+
nn.Linear(embed_dim, embed_dim),
|
| 273 |
+
)
|
| 274 |
+
for _ in range(num_modalities)
|
| 275 |
+
])
|
| 276 |
+
|
| 277 |
+
# Shared adapter
|
| 278 |
+
self.shared_adapter = nn.Sequential(
|
| 279 |
+
nn.Linear(embed_dim, embed_dim),
|
| 280 |
+
nn.GELU(),
|
| 281 |
+
nn.Linear(embed_dim, embed_dim),
|
| 282 |
+
)
|
| 283 |
+
|
| 284 |
+
def forward(
|
| 285 |
+
self,
|
| 286 |
+
features: torch.Tensor,
|
| 287 |
+
modality_idx: int
|
| 288 |
+
) -> torch.Tensor:
|
| 289 |
+
"""
|
| 290 |
+
Apply domain adaptation.
|
| 291 |
+
|
| 292 |
+
Args:
|
| 293 |
+
features: Input features
|
| 294 |
+
modality_idx: Index of the modality
|
| 295 |
+
|
| 296 |
+
Returns:
|
| 297 |
+
Adapted features
|
| 298 |
+
"""
|
| 299 |
+
# Modality-specific adaptation
|
| 300 |
+
adapted = self.adapters[modality_idx](features)
|
| 301 |
+
|
| 302 |
+
# Shared adaptation
|
| 303 |
+
shared = self.shared_adapter(features)
|
| 304 |
+
|
| 305 |
+
# Combine
|
| 306 |
+
return adapted + shared
|
| 307 |
+
|
| 308 |
+
|
| 309 |
+
# Self-check
|
| 310 |
+
if __name__ == "__main__":
|
| 311 |
+
print("Testing Cross-Modal Alignment...")
|
| 312 |
+
|
| 313 |
+
config = CrossModalConfig()
|
| 314 |
+
aligner = CrossModalAligner(config)
|
| 315 |
+
|
| 316 |
+
# Test projection
|
| 317 |
+
features = torch.randn(8, 768)
|
| 318 |
+
optical_proj = aligner.project(features, "optical")
|
| 319 |
+
sar_proj = aligner.project(features, "sar")
|
| 320 |
+
|
| 321 |
+
print(f"Optical projection shape: {optical_proj.shape}")
|
| 322 |
+
print(f"SAR projection shape: {sar_proj.shape}")
|
| 323 |
+
|
| 324 |
+
# Test contrastive loss
|
| 325 |
+
loss = aligner.contrastive_loss(optical_proj, sar_proj)
|
| 326 |
+
print(f"Contrastive loss: {loss.item():.4f}")
|
| 327 |
+
|
| 328 |
+
# Test cross-modal retrieval
|
| 329 |
+
gallery_features = {
|
| 330 |
+
"optical": torch.randn(100, 768),
|
| 331 |
+
"sar": torch.randn(100, 768),
|
| 332 |
+
"multispectral": torch.randn(100, 768),
|
| 333 |
+
}
|
| 334 |
+
|
| 335 |
+
query = torch.randn(1, 768)
|
| 336 |
+
indices, scores = aligner.cross_modal_retrieve(query, "optical", gallery_features, k=5)
|
| 337 |
+
|
| 338 |
+
print(f"Retrieved indices: {indices}")
|
| 339 |
+
print(f"Retrieved scores: {scores}")
|
| 340 |
+
|
| 341 |
+
print("\nCross-Modal Alignment test passed!")
|
src/features/embeddings.py
ADDED
|
@@ -0,0 +1,182 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Embedding cache utilities for pre-computed features.
|
| 3 |
+
|
| 4 |
+
Handles saving/loading embeddings to disk for fast retrieval.
|
| 5 |
+
"""
|
| 6 |
+
|
| 7 |
+
import torch
|
| 8 |
+
from pathlib import Path
|
| 9 |
+
from typing import Tuple, List, Optional, Dict
|
| 10 |
+
import json
|
| 11 |
+
|
| 12 |
+
|
| 13 |
+
def save_embeddings(
|
| 14 |
+
embeddings: torch.Tensor,
|
| 15 |
+
metadata: Dict,
|
| 16 |
+
output_dir: str,
|
| 17 |
+
filename: Optional[str] = None
|
| 18 |
+
) -> Path:
|
| 19 |
+
"""
|
| 20 |
+
Save embeddings and metadata to disk.
|
| 21 |
+
|
| 22 |
+
Args:
|
| 23 |
+
embeddings: Tensor of shape (N, embed_dim)
|
| 24 |
+
metadata: Dict with keys like 'modality', 'sample_ids', 'class_labels'
|
| 25 |
+
output_dir: Directory to save to
|
| 26 |
+
filename: Optional custom filename (without extension)
|
| 27 |
+
|
| 28 |
+
Returns:
|
| 29 |
+
Path to saved file
|
| 30 |
+
"""
|
| 31 |
+
output_dir = Path(output_dir)
|
| 32 |
+
output_dir.mkdir(parents=True, exist_ok=True)
|
| 33 |
+
|
| 34 |
+
# Generate filename if not provided
|
| 35 |
+
if filename is None:
|
| 36 |
+
modality = metadata.get("modality", "unknown")
|
| 37 |
+
n_samples = embeddings.shape[0]
|
| 38 |
+
filename = f"{modality}_embeddings_{n_samples}"
|
| 39 |
+
|
| 40 |
+
# Save embeddings tensor
|
| 41 |
+
embeddings_path = output_dir / f"{filename}.pt"
|
| 42 |
+
torch.save(embeddings, embeddings_path)
|
| 43 |
+
|
| 44 |
+
# Save metadata as JSON
|
| 45 |
+
metadata_path = output_dir / f"{filename}_metadata.json"
|
| 46 |
+
|
| 47 |
+
# Convert tensors in metadata to lists for JSON serialization
|
| 48 |
+
serializable_metadata = {}
|
| 49 |
+
for key, value in metadata.items():
|
| 50 |
+
if isinstance(value, torch.Tensor):
|
| 51 |
+
serializable_metadata[key] = value.tolist()
|
| 52 |
+
else:
|
| 53 |
+
serializable_metadata[key] = value
|
| 54 |
+
|
| 55 |
+
with open(metadata_path, "w") as f:
|
| 56 |
+
json.dump(serializable_metadata, f, indent=2)
|
| 57 |
+
|
| 58 |
+
return embeddings_path
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
def load_embeddings(
|
| 62 |
+
cache_path: str
|
| 63 |
+
) -> Tuple[torch.Tensor, Dict]:
|
| 64 |
+
"""
|
| 65 |
+
Load embeddings and metadata from disk.
|
| 66 |
+
|
| 67 |
+
Args:
|
| 68 |
+
cache_path: Path to .pt embeddings file
|
| 69 |
+
|
| 70 |
+
Returns:
|
| 71 |
+
(embeddings tensor, metadata dict)
|
| 72 |
+
"""
|
| 73 |
+
cache_path = Path(cache_path)
|
| 74 |
+
|
| 75 |
+
# Load embeddings
|
| 76 |
+
embeddings = torch.load(cache_path, weights_only=True)
|
| 77 |
+
|
| 78 |
+
# Load metadata if exists
|
| 79 |
+
metadata_path = cache_path.with_name(
|
| 80 |
+
cache_path.stem + "_metadata.json"
|
| 81 |
+
)
|
| 82 |
+
|
| 83 |
+
metadata = {}
|
| 84 |
+
if metadata_path.exists():
|
| 85 |
+
with open(metadata_path, "r") as f:
|
| 86 |
+
metadata = json.load(f)
|
| 87 |
+
|
| 88 |
+
return embeddings, metadata
|
| 89 |
+
|
| 90 |
+
|
| 91 |
+
def get_cache_path(
|
| 92 |
+
output_dir: str,
|
| 93 |
+
modality: str,
|
| 94 |
+
split: str = "gallery",
|
| 95 |
+
embed_dim: int = 768
|
| 96 |
+
) -> Path:
|
| 97 |
+
"""
|
| 98 |
+
Generate standard cache file path.
|
| 99 |
+
|
| 100 |
+
Args:
|
| 101 |
+
output_dir: Base output directory
|
| 102 |
+
modality: Modality type
|
| 103 |
+
split: Dataset split (query/gallery)
|
| 104 |
+
embed_dim: Embedding dimension
|
| 105 |
+
|
| 106 |
+
Returns:
|
| 107 |
+
Path object for cache file
|
| 108 |
+
"""
|
| 109 |
+
output_dir = Path(output_dir)
|
| 110 |
+
filename = f"{modality}_{split}_embeddings.pt"
|
| 111 |
+
return output_dir / filename
|
| 112 |
+
|
| 113 |
+
|
| 114 |
+
def verify_embeddings(
|
| 115 |
+
embeddings: torch.Tensor,
|
| 116 |
+
expected_dim: Optional[int] = None,
|
| 117 |
+
l2_normalized: bool = True
|
| 118 |
+
) -> bool:
|
| 119 |
+
"""
|
| 120 |
+
Verify embeddings are valid.
|
| 121 |
+
|
| 122 |
+
Args:
|
| 123 |
+
embeddings: Embedding tensor
|
| 124 |
+
expected_dim: Expected embedding dimension
|
| 125 |
+
l2_normalized: Whether embeddings should be L2-normalized
|
| 126 |
+
|
| 127 |
+
Returns:
|
| 128 |
+
True if valid
|
| 129 |
+
"""
|
| 130 |
+
if embeddings.dim() != 2:
|
| 131 |
+
print(f"Expected 2D tensor, got {embeddings.dim()}D")
|
| 132 |
+
return False
|
| 133 |
+
|
| 134 |
+
if expected_dim is not None and embeddings.shape[1] != expected_dim:
|
| 135 |
+
print(f"Expected dim {expected_dim}, got {embeddings.shape[1]}")
|
| 136 |
+
return False
|
| 137 |
+
|
| 138 |
+
if l2_normalized:
|
| 139 |
+
norms = torch.norm(embeddings, dim=1)
|
| 140 |
+
# Check if norms are close to 1 (allowing for floating point)
|
| 141 |
+
if not torch.allclose(norms, torch.ones_like(norms), atol=1e-3):
|
| 142 |
+
print(f"Embeddings not L2-normalized. Norms: {norms[:5]}")
|
| 143 |
+
return False
|
| 144 |
+
|
| 145 |
+
return True
|
| 146 |
+
|
| 147 |
+
|
| 148 |
+
# Self-check
|
| 149 |
+
if __name__ == "__main__":
|
| 150 |
+
import tempfile
|
| 151 |
+
|
| 152 |
+
print("Testing embedding cache utilities...")
|
| 153 |
+
|
| 154 |
+
# Create dummy embeddings
|
| 155 |
+
embeddings = torch.randn(100, 768)
|
| 156 |
+
embeddings = torch.nn.functional.normalize(embeddings, dim=1) # L2 normalize
|
| 157 |
+
|
| 158 |
+
metadata = {
|
| 159 |
+
"modality": "optical",
|
| 160 |
+
"sample_ids": list(range(100)),
|
| 161 |
+
"class_labels": [i % 10 for i in range(100)],
|
| 162 |
+
}
|
| 163 |
+
|
| 164 |
+
# Test save/load roundtrip
|
| 165 |
+
with tempfile.TemporaryDirectory() as tmpdir:
|
| 166 |
+
# Save
|
| 167 |
+
save_path = save_embeddings(embeddings, metadata, tmpdir, "test")
|
| 168 |
+
print(f"Saved to: {save_path}")
|
| 169 |
+
|
| 170 |
+
# Load
|
| 171 |
+
loaded_embeddings, loaded_metadata = load_embeddings(save_path)
|
| 172 |
+
print(f"Loaded embeddings shape: {loaded_embeddings.shape}")
|
| 173 |
+
print(f"Loaded metadata keys: {list(loaded_metadata.keys())}")
|
| 174 |
+
|
| 175 |
+
# Verify
|
| 176 |
+
assert torch.allclose(embeddings, loaded_embeddings), "Embeddings mismatch!"
|
| 177 |
+
assert loaded_metadata["modality"] == "optical", "Metadata mismatch!"
|
| 178 |
+
|
| 179 |
+
# Test verify
|
| 180 |
+
assert verify_embeddings(loaded_embeddings, expected_dim=768), "Verification failed!"
|
| 181 |
+
|
| 182 |
+
print("\nEmbedding cache test passed!")
|
src/features/extractor.py
ADDED
|
@@ -0,0 +1,158 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Feature extraction using SatCLIP for satellite imagery.
|
| 3 |
+
|
| 4 |
+
SatCLIP is trained on Sentinel-2 data - better than generic CLIP
|
| 5 |
+
for satellite image retrieval.
|
| 6 |
+
"""
|
| 7 |
+
|
| 8 |
+
import torch
|
| 9 |
+
import torch.nn.functional as F
|
| 10 |
+
from PIL import Image
|
| 11 |
+
from typing import List, Optional, Tuple
|
| 12 |
+
from torchvision import transforms
|
| 13 |
+
|
| 14 |
+
from .satclip_encoder import SatCLIPEncoder
|
| 15 |
+
|
| 16 |
+
# Wavelength centroids (nm) per modality — needed by DOFA-style models
|
| 17 |
+
# These match Sentinel-2 band centers and are used for positional encoding
|
| 18 |
+
WAVELENGTHS = {
|
| 19 |
+
"optical": torch.tensor([492.4, 559.8, 664.6]), # RGB: B02, B03, B04
|
| 20 |
+
"sar": torch.tensor([5400.0, 5600.0]), # C-band VV, VH (approx, in nm-equivalent)
|
| 21 |
+
"multispectral": torch.tensor([ # Sentinel-2 MS bands
|
| 22 |
+
442.0, 492.4, 559.8, 664.6, 704.1, 740.5, 782.8, 832.8,
|
| 23 |
+
864.7, 945.1, 1373.5, 1613.7
|
| 24 |
+
]),
|
| 25 |
+
}
|
| 26 |
+
|
| 27 |
+
MODALITY_CHANNELS = {
|
| 28 |
+
"optical": 3,
|
| 29 |
+
"sar": 2,
|
| 30 |
+
"multispectral": 12,
|
| 31 |
+
}
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
class FeatureExtractor:
|
| 35 |
+
"""
|
| 36 |
+
Extract features from satellite images using SatCLIP.
|
| 37 |
+
|
| 38 |
+
Uses SatCLIP's ViT trained on Sentinel-2 imagery.
|
| 39 |
+
"""
|
| 40 |
+
|
| 41 |
+
def __init__(self, device: Optional[str] = None):
|
| 42 |
+
self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
|
| 43 |
+
self.encoder = SatCLIPEncoder(device=self.device)
|
| 44 |
+
self.embed_dim = self.encoder.embed_dim
|
| 45 |
+
self.transform = transforms.Compose([
|
| 46 |
+
transforms.Resize((224, 224)),
|
| 47 |
+
transforms.ToTensor(),
|
| 48 |
+
])
|
| 49 |
+
|
| 50 |
+
def _preprocess(self, image: Image.Image, modality: str) -> torch.Tensor:
|
| 51 |
+
tensor = self.transform(image).unsqueeze(0)
|
| 52 |
+
return self._pad_to_13ch(tensor, modality)
|
| 53 |
+
|
| 54 |
+
def _pad_to_13ch(self, tensor: torch.Tensor, modality: str = "optical") -> torch.Tensor:
|
| 55 |
+
"""Pad tensor to 13 channels for SatCLIP. Handles 1-13 channels."""
|
| 56 |
+
n_channels = tensor.shape[1]
|
| 57 |
+
if n_channels >= 13:
|
| 58 |
+
return tensor[:, :13, :, :]
|
| 59 |
+
# Repeat single channel to 3 (grayscale SAR fallback)
|
| 60 |
+
if n_channels == 1:
|
| 61 |
+
tensor = tensor.repeat(1, 3, 1, 1)
|
| 62 |
+
n_channels = 3
|
| 63 |
+
# Repeat 2 channels to 3 (SAR VV/VH)
|
| 64 |
+
if n_channels == 2:
|
| 65 |
+
third = tensor[:, :1, :, :] # duplicate VV as 3rd channel
|
| 66 |
+
tensor = torch.cat([tensor, third], dim=1)
|
| 67 |
+
n_channels = 3
|
| 68 |
+
pad_channels = 13 - n_channels
|
| 69 |
+
padding = torch.zeros(
|
| 70 |
+
tensor.shape[0], pad_channels, tensor.shape[2], tensor.shape[3])
|
| 71 |
+
return torch.cat([tensor, padding], dim=1)
|
| 72 |
+
|
| 73 |
+
def _preprocess_batch(self, images: List[Image.Image], modality: str) -> torch.Tensor:
|
| 74 |
+
return torch.stack([self._preprocess(img, modality) for img in images])
|
| 75 |
+
|
| 76 |
+
@torch.no_grad()
|
| 77 |
+
def extract_features(
|
| 78 |
+
self,
|
| 79 |
+
image: Image.Image,
|
| 80 |
+
modality: str = "optical",
|
| 81 |
+
normalize: bool = True
|
| 82 |
+
) -> torch.Tensor:
|
| 83 |
+
tensor = self._preprocess(image, modality)
|
| 84 |
+
features = self.encoder.encode(tensor, normalize=normalize)
|
| 85 |
+
return features.squeeze(0)
|
| 86 |
+
|
| 87 |
+
@torch.no_grad()
|
| 88 |
+
def extract_features_from_tensor(
|
| 89 |
+
self,
|
| 90 |
+
tensor: torch.Tensor,
|
| 91 |
+
modality: str = "optical",
|
| 92 |
+
normalize: bool = True
|
| 93 |
+
) -> torch.Tensor:
|
| 94 |
+
"""Extract features from a raw (C, H, W) tensor with arbitrary channels."""
|
| 95 |
+
if tensor.ndim == 3:
|
| 96 |
+
tensor = tensor.unsqueeze(0)
|
| 97 |
+
if tensor.shape[1] < 13:
|
| 98 |
+
tensor = self._pad_to_13ch(tensor, modality)
|
| 99 |
+
tensor = tensor.to(self.device)
|
| 100 |
+
features = self.encoder.encode(tensor, normalize=normalize)
|
| 101 |
+
return features.squeeze(0)
|
| 102 |
+
|
| 103 |
+
@torch.no_grad()
|
| 104 |
+
def extract_batch(
|
| 105 |
+
self,
|
| 106 |
+
images: List[Image.Image],
|
| 107 |
+
modality: str = "optical",
|
| 108 |
+
batch_size: int = 32,
|
| 109 |
+
normalize: bool = True
|
| 110 |
+
) -> torch.Tensor:
|
| 111 |
+
all_features = []
|
| 112 |
+
for i in range(0, len(images), batch_size):
|
| 113 |
+
batch = images[i:i + batch_size]
|
| 114 |
+
tensors = self._preprocess_batch(batch, modality)
|
| 115 |
+
features = self.encoder.encode(tensors, normalize=normalize)
|
| 116 |
+
all_features.append(features.cpu())
|
| 117 |
+
return torch.cat(all_features, dim=0)
|
| 118 |
+
|
| 119 |
+
def embed_dataset(
|
| 120 |
+
self,
|
| 121 |
+
dataset,
|
| 122 |
+
batch_size: int = 32,
|
| 123 |
+
show_progress: bool = True
|
| 124 |
+
) -> Tuple[torch.Tensor, List[int], List[int]]:
|
| 125 |
+
from torch.utils.data import DataLoader
|
| 126 |
+
|
| 127 |
+
loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=0)
|
| 128 |
+
all_embeddings = []
|
| 129 |
+
all_modality_labels = []
|
| 130 |
+
all_class_labels = []
|
| 131 |
+
|
| 132 |
+
for batch_idx, (images, mod_labels, class_labels) in enumerate(loader):
|
| 133 |
+
images = images.to(self.device)
|
| 134 |
+
with torch.no_grad():
|
| 135 |
+
features = self.encoder.encode(images, normalize=True)
|
| 136 |
+
all_embeddings.append(features.cpu())
|
| 137 |
+
all_modality_labels.extend(mod_labels.numpy().tolist())
|
| 138 |
+
all_class_labels.extend(class_labels.numpy().tolist())
|
| 139 |
+
if show_progress and (batch_idx + 1) % 10 == 0:
|
| 140 |
+
print(f"Embedded {batch_idx + 1}/{len(loader)} batches")
|
| 141 |
+
|
| 142 |
+
return torch.cat(all_embeddings, dim=0), all_modality_labels, all_class_labels
|
| 143 |
+
|
| 144 |
+
|
| 145 |
+
if __name__ == "__main__":
|
| 146 |
+
print("Testing SatCLIP FeatureExtractor...")
|
| 147 |
+
extractor = FeatureExtractor()
|
| 148 |
+
print(f"Embed dim: {extractor.embed_dim}")
|
| 149 |
+
|
| 150 |
+
dummy = Image.fromarray(torch.randint(0, 255, (224, 224, 3)).numpy())
|
| 151 |
+
features = extractor.extract_features(dummy)
|
| 152 |
+
print(f"Single shape: {features.shape}")
|
| 153 |
+
print(f"L2 norm: {features.norm().item():.4f}")
|
| 154 |
+
|
| 155 |
+
batch = [dummy] * 4
|
| 156 |
+
batch_features = extractor.extract_batch(batch)
|
| 157 |
+
print(f"Batch shape: {batch_features.shape}")
|
| 158 |
+
print("OK")
|
src/features/hybrid.py
ADDED
|
@@ -0,0 +1,162 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Hybrid feature extractor: CLIP + SAR Adapter + DINOv2.
|
| 3 |
+
|
| 4 |
+
Combines CLIP global semantics, DINOv2 patch features,
|
| 5 |
+
and SAR-specific preprocessing into a single retrieval-ready module.
|
| 6 |
+
"""
|
| 7 |
+
|
| 8 |
+
import torch
|
| 9 |
+
import torch.nn as nn
|
| 10 |
+
import torch.nn.functional as F
|
| 11 |
+
from PIL import Image
|
| 12 |
+
from typing import Optional
|
| 13 |
+
from dataclasses import dataclass
|
| 14 |
+
import numpy as np
|
| 15 |
+
from torchvision import transforms
|
| 16 |
+
|
| 17 |
+
from .sar_adapter import SARAdapter
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
@dataclass
|
| 21 |
+
class HybridConfig:
|
| 22 |
+
clip_model: str = "openai/clip-vit-large-patch14"
|
| 23 |
+
dinov2_model: str = "facebook/dinov2-base"
|
| 24 |
+
clip_weight: float = 0.7
|
| 25 |
+
dinov2_weight: float = 0.3
|
| 26 |
+
embed_dim: int = 768
|
| 27 |
+
device: Optional[str] = None
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
class HybridExtractor(nn.Module):
|
| 31 |
+
"""
|
| 32 |
+
Unified hybrid extractor combining CLIP, DINOv2, and SAR adapter.
|
| 33 |
+
|
| 34 |
+
Fusion: embedding = w_clip * CLIP(img) + w_dino * DINOv2(img)
|
| 35 |
+
SAR path: SAR -> adapter(2ch->3ch) -> CLIP+DINOv2
|
| 36 |
+
"""
|
| 37 |
+
|
| 38 |
+
def __init__(self, config: Optional[HybridConfig] = None):
|
| 39 |
+
super().__init__()
|
| 40 |
+
self.config = config or HybridConfig()
|
| 41 |
+
self.device = self.config.device or ("cuda" if torch.cuda.is_available() else "cpu")
|
| 42 |
+
|
| 43 |
+
self.sar_adapter = SARAdapter().to(self.device)
|
| 44 |
+
self._clip_model = None
|
| 45 |
+
self._clip_processor = None
|
| 46 |
+
self._dinov2_model = None
|
| 47 |
+
self._loaded = False
|
| 48 |
+
|
| 49 |
+
self.dino_transform = transforms.Compose([
|
| 50 |
+
transforms.Resize((224, 224)),
|
| 51 |
+
transforms.ToTensor(),
|
| 52 |
+
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
| 53 |
+
])
|
| 54 |
+
|
| 55 |
+
self.fusion_proj = nn.Sequential(
|
| 56 |
+
nn.Linear(self.config.embed_dim, self.config.embed_dim),
|
| 57 |
+
nn.GELU(),
|
| 58 |
+
nn.Linear(self.config.embed_dim, self.config.embed_dim),
|
| 59 |
+
)
|
| 60 |
+
|
| 61 |
+
def load(self):
|
| 62 |
+
if self._loaded:
|
| 63 |
+
return
|
| 64 |
+
from transformers import CLIPProcessor, CLIPModel, AutoModel
|
| 65 |
+
|
| 66 |
+
print(f"Loading CLIP: {self.config.clip_model} ...")
|
| 67 |
+
self._clip_processor = CLIPProcessor.from_pretrained(self.config.clip_model)
|
| 68 |
+
self._clip_model = CLIPModel.from_pretrained(self.config.clip_model).to(self.device)
|
| 69 |
+
self._clip_model.eval()
|
| 70 |
+
|
| 71 |
+
print(f"Loading DINOv2: {self.config.dinov2_model} ...")
|
| 72 |
+
try:
|
| 73 |
+
self._dinov2_model = AutoModel.from_pretrained(self.config.dinov2_model).to(self.device)
|
| 74 |
+
self._dinov2_model.eval()
|
| 75 |
+
self._has_dino = True
|
| 76 |
+
print("DINOv2 loaded")
|
| 77 |
+
except Exception as e:
|
| 78 |
+
self._has_dino = False
|
| 79 |
+
print(f"DINOv2 unavailable: {e}")
|
| 80 |
+
|
| 81 |
+
self._loaded = True
|
| 82 |
+
print(f"Hybrid extractor ready on {self.device}")
|
| 83 |
+
|
| 84 |
+
@torch.no_grad()
|
| 85 |
+
def _clip_features(self, img: Image.Image) -> np.ndarray:
|
| 86 |
+
inputs = self._clip_processor(images=img, return_tensors="pt").to(self.device)
|
| 87 |
+
out = self._clip_model.vision_model(**inputs)
|
| 88 |
+
pooled = out.last_hidden_state[:, 0, :]
|
| 89 |
+
feat = self._clip_model.visual_projection(pooled).squeeze(0)
|
| 90 |
+
return torch.nn.functional.normalize(feat, dim=-1).cpu().numpy()
|
| 91 |
+
|
| 92 |
+
@torch.no_grad()
|
| 93 |
+
def _dinov2_features(self, img: Image.Image) -> Optional[np.ndarray]:
|
| 94 |
+
if not self._has_dino:
|
| 95 |
+
return None
|
| 96 |
+
t = self.dino_transform(img).unsqueeze(0).to(self.device)
|
| 97 |
+
out = self._dinov2_model(t)
|
| 98 |
+
patch_feat = out.last_hidden_state[:, 1:, :].mean(dim=1)
|
| 99 |
+
return torch.nn.functional.normalize(patch_feat.squeeze(0), dim=-1).cpu().numpy()
|
| 100 |
+
|
| 101 |
+
def _preprocess_sar(self, img: Image.Image) -> Image.Image:
|
| 102 |
+
arr = np.array(img).astype(np.float32) / 255.0
|
| 103 |
+
t = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0)
|
| 104 |
+
with torch.no_grad():
|
| 105 |
+
adapted = self.sar_adapter(t)
|
| 106 |
+
arr_out = (adapted.squeeze(0).permute(1, 2, 0).numpy() * 255).clip(0, 255).astype(np.uint8)
|
| 107 |
+
return Image.fromarray(arr_out)
|
| 108 |
+
|
| 109 |
+
def extract(
|
| 110 |
+
self,
|
| 111 |
+
img: Image.Image,
|
| 112 |
+
modality: str = "optical",
|
| 113 |
+
normalize: bool = True,
|
| 114 |
+
) -> np.ndarray:
|
| 115 |
+
if not self._loaded:
|
| 116 |
+
self.load()
|
| 117 |
+
|
| 118 |
+
if modality == "sar":
|
| 119 |
+
img = self._preprocess_sar(img)
|
| 120 |
+
|
| 121 |
+
clip_feat = self._clip_features(img)
|
| 122 |
+
dino_feat = self._dinov2_features(img)
|
| 123 |
+
|
| 124 |
+
if dino_feat is not None:
|
| 125 |
+
w_c, w_d = self.config.clip_weight, self.config.dinov2_weight
|
| 126 |
+
hybrid = w_c * clip_feat + w_d * dino_feat
|
| 127 |
+
else:
|
| 128 |
+
hybrid = clip_feat
|
| 129 |
+
|
| 130 |
+
if normalize:
|
| 131 |
+
norm = np.linalg.norm(hybrid)
|
| 132 |
+
if norm > 0:
|
| 133 |
+
hybrid = hybrid / norm
|
| 134 |
+
|
| 135 |
+
return hybrid.astype(np.float32)
|
| 136 |
+
|
| 137 |
+
def extract_batch(
|
| 138 |
+
self,
|
| 139 |
+
images: list,
|
| 140 |
+
modalities: list = None,
|
| 141 |
+
normalize: bool = True,
|
| 142 |
+
) -> np.ndarray:
|
| 143 |
+
if modalities is None:
|
| 144 |
+
modalities = ["optical"] * len(images)
|
| 145 |
+
return np.array([
|
| 146 |
+
self.extract(img, mod, normalize)
|
| 147 |
+
for img, mod in zip(images, modalities)
|
| 148 |
+
])
|
| 149 |
+
|
| 150 |
+
|
| 151 |
+
def create_hybrid_extractor(**kwargs) -> HybridExtractor:
|
| 152 |
+
config = HybridConfig(**kwargs)
|
| 153 |
+
return HybridExtractor(config)
|
| 154 |
+
|
| 155 |
+
|
| 156 |
+
if __name__ == "__main__":
|
| 157 |
+
ext = create_hybrid_extractor()
|
| 158 |
+
ext.load()
|
| 159 |
+
|
| 160 |
+
dummy = Image.fromarray(np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8))
|
| 161 |
+
feat = ext.extract(dummy, modality="optical")
|
| 162 |
+
print(f"Feature dim: {feat.shape}, norm: {np.linalg.norm(feat):.4f}")
|
src/features/multiscale.py
ADDED
|
@@ -0,0 +1,350 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Multiscale feature extraction for satellite imagery.
|
| 3 |
+
|
| 4 |
+
Combines patch-level and global features for richer representations.
|
| 5 |
+
Uses DINOv2 for patch features and CLIP for global alignment.
|
| 6 |
+
"""
|
| 7 |
+
|
| 8 |
+
import torch
|
| 9 |
+
import torch.nn as nn
|
| 10 |
+
import torch.nn.functional as F
|
| 11 |
+
from typing import Optional, Tuple, Dict, Any
|
| 12 |
+
from dataclasses import dataclass
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
@dataclass
|
| 16 |
+
class MultiscaleFeatures:
|
| 17 |
+
"""Container for multiscale features."""
|
| 18 |
+
global_feature: torch.Tensor # (embed_dim,) - CLIP-style global
|
| 19 |
+
patch_features: torch.Tensor # (num_patches, patch_dim) - DINOv2-style
|
| 20 |
+
patch_grid: Tuple[int, int] # (H, W) grid of patches
|
| 21 |
+
combined: torch.Tensor # (combined_dim,) - fused feature
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
class PatchAggregator(nn.Module):
|
| 25 |
+
"""
|
| 26 |
+
Aggregates patch features into a single representation.
|
| 27 |
+
|
| 28 |
+
Supports multiple aggregation strategies:
|
| 29 |
+
- mean: Average pooling
|
| 30 |
+
- max: Max pooling
|
| 31 |
+
- attention: Learnable attention pooling
|
| 32 |
+
"""
|
| 33 |
+
|
| 34 |
+
def __init__(self, patch_dim: int, strategy: str = "attention"):
|
| 35 |
+
super().__init__()
|
| 36 |
+
|
| 37 |
+
self.strategy = strategy
|
| 38 |
+
|
| 39 |
+
if strategy == "attention":
|
| 40 |
+
self.attention = nn.Sequential(
|
| 41 |
+
nn.Linear(patch_dim, patch_dim // 4),
|
| 42 |
+
nn.Tanh(),
|
| 43 |
+
nn.Linear(patch_dim // 4, 1),
|
| 44 |
+
)
|
| 45 |
+
elif strategy == "cls":
|
| 46 |
+
self.cls_token = nn.Parameter(torch.randn(1, 1, patch_dim))
|
| 47 |
+
|
| 48 |
+
def forward(self, patch_features: torch.Tensor) -> torch.Tensor:
|
| 49 |
+
"""
|
| 50 |
+
Aggregate patch features.
|
| 51 |
+
|
| 52 |
+
Args:
|
| 53 |
+
patch_features: (B, num_patches, patch_dim)
|
| 54 |
+
|
| 55 |
+
Returns:
|
| 56 |
+
Aggregated feature (B, patch_dim)
|
| 57 |
+
"""
|
| 58 |
+
if self.strategy == "mean":
|
| 59 |
+
return patch_features.mean(dim=1)
|
| 60 |
+
|
| 61 |
+
elif self.strategy == "max":
|
| 62 |
+
return patch_features.max(dim=1)[0]
|
| 63 |
+
|
| 64 |
+
elif self.strategy == "attention":
|
| 65 |
+
# (B, num_patches, 1)
|
| 66 |
+
attn_weights = self.attention(patch_features)
|
| 67 |
+
attn_weights = F.softmax(attn_weights, dim=1)
|
| 68 |
+
# (B, patch_dim)
|
| 69 |
+
return (patch_features * attn_weights).sum(dim=1)
|
| 70 |
+
|
| 71 |
+
elif self.strategy == "cls":
|
| 72 |
+
B = patch_features.shape[0]
|
| 73 |
+
cls_tokens = self.cls_token.expand(B, -1, -1)
|
| 74 |
+
# Prepend CLS token
|
| 75 |
+
x = torch.cat([cls_tokens, patch_features], dim=1)
|
| 76 |
+
return x[:, 0]
|
| 77 |
+
|
| 78 |
+
else:
|
| 79 |
+
raise ValueError(f"Unknown strategy: {self.strategy}")
|
| 80 |
+
|
| 81 |
+
|
| 82 |
+
class MultiscaleExtractor(nn.Module):
|
| 83 |
+
"""
|
| 84 |
+
Extracts features at multiple scales from satellite imagery.
|
| 85 |
+
|
| 86 |
+
Combines:
|
| 87 |
+
- Global features from CLIP (semantic alignment)
|
| 88 |
+
- Patch features from DINOv2 (spatial details)
|
| 89 |
+
- Cross-scale attention for feature fusion
|
| 90 |
+
"""
|
| 91 |
+
|
| 92 |
+
def __init__(
|
| 93 |
+
self,
|
| 94 |
+
clip_model: nn.Module,
|
| 95 |
+
dinov2_model: Optional[nn.Module] = None,
|
| 96 |
+
embed_dim: int = 768,
|
| 97 |
+
patch_dim: int = 768,
|
| 98 |
+
fusion_dim: int = 512,
|
| 99 |
+
use_cross_attention: bool = True
|
| 100 |
+
):
|
| 101 |
+
super().__init__()
|
| 102 |
+
|
| 103 |
+
self.clip_model = clip_model
|
| 104 |
+
self.dinov2_model = dinov2_model
|
| 105 |
+
|
| 106 |
+
self.embed_dim = embed_dim
|
| 107 |
+
self.patch_dim = patch_dim
|
| 108 |
+
self.fusion_dim = fusion_dim
|
| 109 |
+
|
| 110 |
+
# Patch aggregation
|
| 111 |
+
self.patch_aggregator = PatchAggregator(patch_dim, strategy="attention")
|
| 112 |
+
|
| 113 |
+
# Cross-scale attention (fuses global + patch features)
|
| 114 |
+
self.use_cross_attention = use_cross_attention
|
| 115 |
+
if use_cross_attention:
|
| 116 |
+
self.cross_attn = nn.MultiheadAttention(
|
| 117 |
+
embed_dim=embed_dim,
|
| 118 |
+
num_heads=8,
|
| 119 |
+
dropout=0.1,
|
| 120 |
+
batch_first=True
|
| 121 |
+
)
|
| 122 |
+
self.fusion_proj = nn.Linear(embed_dim + patch_dim, fusion_dim)
|
| 123 |
+
else:
|
| 124 |
+
# Simple concatenation + projection
|
| 125 |
+
self.fusion_proj = nn.Linear(embed_dim + patch_dim, fusion_dim)
|
| 126 |
+
|
| 127 |
+
# Final normalization
|
| 128 |
+
self.layer_norm = nn.LayerNorm(fusion_dim)
|
| 129 |
+
|
| 130 |
+
@torch.no_grad()
|
| 131 |
+
def extract_clip_global(self, x: torch.Tensor) -> torch.Tensor:
|
| 132 |
+
"""Extract global features from CLIP."""
|
| 133 |
+
# Assuming CLIP vision model
|
| 134 |
+
if hasattr(self.clip_model, 'vision_model'):
|
| 135 |
+
output = self.clip_model.vision_model(pixel_values=x)
|
| 136 |
+
pooled = output.last_hidden_state[:, 0, :] # CLS token
|
| 137 |
+
global_feat = self.clip_model.visual_projection(pooled)
|
| 138 |
+
else:
|
| 139 |
+
# Fallback for other architectures
|
| 140 |
+
global_feat = self.clip_model(x)
|
| 141 |
+
|
| 142 |
+
return F.normalize(global_feat, dim=-1)
|
| 143 |
+
|
| 144 |
+
@torch.no_grad()
|
| 145 |
+
def extract_dinov2_patches(self, x: torch.Tensor) -> torch.Tensor:
|
| 146 |
+
"""Extract patch features from DINOv2."""
|
| 147 |
+
if self.dinov2_model is None:
|
| 148 |
+
# Return dummy features
|
| 149 |
+
B = x.shape[0]
|
| 150 |
+
num_patches = 196 # 14x14 for 224x224 input
|
| 151 |
+
return torch.randn(B, num_patches, self.patch_dim, device=x.device)
|
| 152 |
+
|
| 153 |
+
# DINOv2 forward pass
|
| 154 |
+
output = self.dinov2_model(x)
|
| 155 |
+
|
| 156 |
+
# Handle different output formats
|
| 157 |
+
if hasattr(output, 'last_hidden_state'):
|
| 158 |
+
patch_features = output.last_hidden_state[:, 1:] # Remove CLS token
|
| 159 |
+
elif isinstance(output, torch.Tensor):
|
| 160 |
+
patch_features = output[:, 1:] # Remove CLS token if present
|
| 161 |
+
else:
|
| 162 |
+
# Assume output is the patch features directly
|
| 163 |
+
patch_features = output
|
| 164 |
+
|
| 165 |
+
return patch_features
|
| 166 |
+
|
| 167 |
+
def fuse_features(
|
| 168 |
+
self,
|
| 169 |
+
global_feat: torch.Tensor,
|
| 170 |
+
patch_feat: torch.Tensor
|
| 171 |
+
) -> torch.Tensor:
|
| 172 |
+
"""
|
| 173 |
+
Fuse global and patch features.
|
| 174 |
+
|
| 175 |
+
Args:
|
| 176 |
+
global_feat: (B, embed_dim)
|
| 177 |
+
patch_feat: (B, patch_dim)
|
| 178 |
+
|
| 179 |
+
Returns:
|
| 180 |
+
Fused feature (B, fusion_dim)
|
| 181 |
+
"""
|
| 182 |
+
if self.use_cross_attention:
|
| 183 |
+
# Use global as query, patches as keys/values
|
| 184 |
+
B = global_feat.shape[0]
|
| 185 |
+
global_seq = global_feat.unsqueeze(1) # (B, 1, embed_dim)
|
| 186 |
+
patch_seq = patch_feat.unsqueeze(1) # (B, 1, patch_dim) - simplified
|
| 187 |
+
|
| 188 |
+
# Cross attention
|
| 189 |
+
attn_out, _ = self.cross_attn(
|
| 190 |
+
query=global_seq,
|
| 191 |
+
key=patch_seq,
|
| 192 |
+
value=patch_seq
|
| 193 |
+
)
|
| 194 |
+
attn_out = attn_out.squeeze(1) # (B, embed_dim)
|
| 195 |
+
|
| 196 |
+
# Concatenate and project
|
| 197 |
+
combined = torch.cat([attn_out, patch_feat], dim=-1)
|
| 198 |
+
else:
|
| 199 |
+
combined = torch.cat([global_feat, patch_feat], dim=-1)
|
| 200 |
+
|
| 201 |
+
# Project to fusion dim
|
| 202 |
+
fused = self.fusion_proj(combined)
|
| 203 |
+
fused = self.layer_norm(fused)
|
| 204 |
+
|
| 205 |
+
return F.normalize(fused, dim=-1)
|
| 206 |
+
|
| 207 |
+
def forward(
|
| 208 |
+
self,
|
| 209 |
+
x: torch.Tensor,
|
| 210 |
+
return_separate: bool = False
|
| 211 |
+
) -> MultiscaleFeatures:
|
| 212 |
+
"""
|
| 213 |
+
Extract multiscale features.
|
| 214 |
+
|
| 215 |
+
Args:
|
| 216 |
+
x: Input image tensor (B, C, H, W)
|
| 217 |
+
return_separate: If True, return separate features instead of fused
|
| 218 |
+
|
| 219 |
+
Returns:
|
| 220 |
+
MultiscaleFeatures container
|
| 221 |
+
"""
|
| 222 |
+
# Extract features
|
| 223 |
+
global_feat = self.extract_clip_global(x)
|
| 224 |
+
patch_feat = self.extract_dinov2_patches(x)
|
| 225 |
+
|
| 226 |
+
# Aggregate patches
|
| 227 |
+
patch_agg = self.patch_aggregator(patch_feat)
|
| 228 |
+
|
| 229 |
+
# Compute patch grid
|
| 230 |
+
B = x.shape[0]
|
| 231 |
+
num_patches = patch_feat.shape[1]
|
| 232 |
+
patch_grid = (int(num_patches ** 0.5), int(num_patches ** 0.5))
|
| 233 |
+
|
| 234 |
+
# Fuse features
|
| 235 |
+
combined = self.fuse_features(global_feat, patch_agg)
|
| 236 |
+
|
| 237 |
+
return MultiscaleFeatures(
|
| 238 |
+
global_feature=global_feat.squeeze(0) if B == 1 else global_feat,
|
| 239 |
+
patch_features=patch_feat.squeeze(0) if B == 1 else patch_feat,
|
| 240 |
+
patch_grid=patch_grid,
|
| 241 |
+
combined=combined.squeeze(0) if B == 1 else combined
|
| 242 |
+
)
|
| 243 |
+
|
| 244 |
+
|
| 245 |
+
class MultiscaleRetrievalHead(nn.Module):
|
| 246 |
+
"""
|
| 247 |
+
Retrieval head that combines multiscale features.
|
| 248 |
+
|
| 249 |
+
Projects fused features to the final embedding space
|
| 250 |
+
used for similarity search.
|
| 251 |
+
"""
|
| 252 |
+
|
| 253 |
+
def __init__(
|
| 254 |
+
self,
|
| 255 |
+
input_dim: int,
|
| 256 |
+
output_dim: int = 768,
|
| 257 |
+
hidden_dim: int = 256
|
| 258 |
+
):
|
| 259 |
+
super().__init__()
|
| 260 |
+
|
| 261 |
+
self.projection = nn.Sequential(
|
| 262 |
+
nn.Linear(input_dim, hidden_dim),
|
| 263 |
+
nn.GELU(),
|
| 264 |
+
nn.Dropout(0.1),
|
| 265 |
+
nn.Linear(hidden_dim, output_dim),
|
| 266 |
+
)
|
| 267 |
+
|
| 268 |
+
def forward(self, features: MultiscaleFeatures) -> torch.Tensor:
|
| 269 |
+
"""
|
| 270 |
+
Project multiscale features to retrieval space.
|
| 271 |
+
|
| 272 |
+
Args:
|
| 273 |
+
features: MultiscaleFeatures container
|
| 274 |
+
|
| 275 |
+
Returns:
|
| 276 |
+
Projected embedding (output_dim,)
|
| 277 |
+
"""
|
| 278 |
+
return self.projection(features.combined)
|
| 279 |
+
|
| 280 |
+
|
| 281 |
+
# Convenience function
|
| 282 |
+
def create_multiscale_extractor(
|
| 283 |
+
clip_model: nn.Module,
|
| 284 |
+
dinov2_model: Optional[nn.Module] = None,
|
| 285 |
+
embed_dim: int = 768,
|
| 286 |
+
fusion_dim: int = 512
|
| 287 |
+
) -> MultiscaleExtractor:
|
| 288 |
+
"""
|
| 289 |
+
Create a multiscale feature extractor.
|
| 290 |
+
|
| 291 |
+
Args:
|
| 292 |
+
clip_model: CLIP vision model for global features
|
| 293 |
+
dinov2_model: Optional DINOv2 model for patch features
|
| 294 |
+
embed_dim: CLIP embedding dimension
|
| 295 |
+
fusion_dim: Output fusion dimension
|
| 296 |
+
|
| 297 |
+
Returns:
|
| 298 |
+
MultiscaleExtractor instance
|
| 299 |
+
"""
|
| 300 |
+
return MultiscaleExtractor(
|
| 301 |
+
clip_model=clip_model,
|
| 302 |
+
dinov2_model=dinov2_model,
|
| 303 |
+
embed_dim=embed_dim,
|
| 304 |
+
patch_dim=768, # DINOv2 default
|
| 305 |
+
fusion_dim=fusion_dim,
|
| 306 |
+
use_cross_attention=True
|
| 307 |
+
)
|
| 308 |
+
|
| 309 |
+
|
| 310 |
+
# Self-check
|
| 311 |
+
if __name__ == "__main__":
|
| 312 |
+
print("Testing MultiscaleExtractor...")
|
| 313 |
+
|
| 314 |
+
# Test without actual models (dummy)
|
| 315 |
+
class DummyModel(nn.Module):
|
| 316 |
+
def __init__(self, output_dim=768):
|
| 317 |
+
super().__init__()
|
| 318 |
+
self.linear = nn.Linear(3, output_dim)
|
| 319 |
+
|
| 320 |
+
def forward(self, x):
|
| 321 |
+
B = x.shape[0]
|
| 322 |
+
return torch.randn(B, 197, 768) # 196 patches + CLS
|
| 323 |
+
|
| 324 |
+
dummy_clip = DummyModel(768)
|
| 325 |
+
dummy_dinov2 = DummyModel(768)
|
| 326 |
+
|
| 327 |
+
extractor = MultiscaleExtractor(
|
| 328 |
+
clip_model=dummy_clip,
|
| 329 |
+
dinov2_model=dummy_dinov2,
|
| 330 |
+
embed_dim=768,
|
| 331 |
+
patch_dim=768,
|
| 332 |
+
fusion_dim=512
|
| 333 |
+
)
|
| 334 |
+
|
| 335 |
+
# Test forward pass
|
| 336 |
+
x = torch.randn(1, 3, 224, 224)
|
| 337 |
+
features = extractor(x)
|
| 338 |
+
|
| 339 |
+
print(f"Global feature shape: {features.global_feature.shape}")
|
| 340 |
+
print(f"Patch features shape: {features.patch_features.shape}")
|
| 341 |
+
print(f"Patch grid: {features.patch_grid}")
|
| 342 |
+
print(f"Combined feature shape: {features.combined.shape}")
|
| 343 |
+
|
| 344 |
+
# Test retrieval head
|
| 345 |
+
head = MultiscaleRetrievalHead(input_dim=512, output_dim=768)
|
| 346 |
+
embedding = head(features)
|
| 347 |
+
print(f"Final embedding shape: {embedding.shape}")
|
| 348 |
+
print(f"Embedding norm: {torch.norm(embedding).item():.4f}")
|
| 349 |
+
|
| 350 |
+
print("\nMultiscaleExtractor test passed!")
|
src/features/sar_adapter.py
ADDED
|
@@ -0,0 +1,208 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
SAR-specific adapter layers for CLIP.
|
| 3 |
+
|
| 4 |
+
Adds lightweight adapter modules to improve SAR modality handling
|
| 5 |
+
without modifying the base CLIP weights.
|
| 6 |
+
"""
|
| 7 |
+
|
| 8 |
+
import torch
|
| 9 |
+
import torch.nn as nn
|
| 10 |
+
import torch.nn.functional as F
|
| 11 |
+
from typing import Optional
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
class SARAdapter(nn.Module):
|
| 15 |
+
"""
|
| 16 |
+
Lightweight adapter for SAR imagery.
|
| 17 |
+
|
| 18 |
+
Applies learnable transformations to bridge the domain gap between
|
| 19 |
+
optical and SAR imagery. Uses:
|
| 20 |
+
1. Channel projection (2ch SAR → 3ch RGB-like)
|
| 21 |
+
2. Learnable scaling to match CLIP input distribution
|
| 22 |
+
3. Optional speckle noise reduction
|
| 23 |
+
"""
|
| 24 |
+
|
| 25 |
+
def __init__(
|
| 26 |
+
self,
|
| 27 |
+
in_channels: int = 2,
|
| 28 |
+
out_channels: int = 3,
|
| 29 |
+
hidden_dim: int = 64,
|
| 30 |
+
dropout: float = 0.1
|
| 31 |
+
):
|
| 32 |
+
super().__init__()
|
| 33 |
+
|
| 34 |
+
# Channel projection: 2ch (VV, VH) → 3ch (RGB-like)
|
| 35 |
+
self.channel_proj = nn.Sequential(
|
| 36 |
+
nn.Conv2d(in_channels, hidden_dim, kernel_size=1, bias=False),
|
| 37 |
+
nn.BatchNorm2d(hidden_dim),
|
| 38 |
+
nn.GELU(),
|
| 39 |
+
nn.Conv2d(hidden_dim, out_channels, kernel_size=1, bias=False),
|
| 40 |
+
nn.BatchNorm2d(out_channels),
|
| 41 |
+
)
|
| 42 |
+
|
| 43 |
+
# Learnable scaling per channel
|
| 44 |
+
self.channel_scale = nn.Parameter(torch.ones(out_channels))
|
| 45 |
+
self.channel_bias = nn.Parameter(torch.zeros(out_channels))
|
| 46 |
+
|
| 47 |
+
# Optional speckle reduction
|
| 48 |
+
self.speckle_reduction = nn.Sequential(
|
| 49 |
+
nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, groups=out_channels),
|
| 50 |
+
nn.Sigmoid(),
|
| 51 |
+
)
|
| 52 |
+
|
| 53 |
+
self.dropout = nn.Dropout2d(dropout)
|
| 54 |
+
|
| 55 |
+
self._init_weights()
|
| 56 |
+
|
| 57 |
+
def _init_weights(self):
|
| 58 |
+
"""Initialize weights with small values."""
|
| 59 |
+
for m in self.modules():
|
| 60 |
+
if isinstance(m, nn.Conv2d):
|
| 61 |
+
nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
| 62 |
+
elif isinstance(m, nn.BatchNorm2d):
|
| 63 |
+
nn.init.constant_(m.weight, 1)
|
| 64 |
+
nn.init.constant_(m.bias, 0)
|
| 65 |
+
|
| 66 |
+
def forward(self, x: torch.Tensor, apply_speckle: bool = True) -> torch.Tensor:
|
| 67 |
+
"""
|
| 68 |
+
Forward pass.
|
| 69 |
+
|
| 70 |
+
Args:
|
| 71 |
+
x: SAR image tensor, shape (B, 2, H, W) with VV, VH channels
|
| 72 |
+
apply_speckle: Whether to apply speckle reduction
|
| 73 |
+
|
| 74 |
+
Returns:
|
| 75 |
+
Projected tensor, shape (B, 3, H, W)
|
| 76 |
+
"""
|
| 77 |
+
# Channel projection
|
| 78 |
+
x = self.channel_proj(x)
|
| 79 |
+
|
| 80 |
+
# Apply speckle reduction
|
| 81 |
+
if apply_speckle:
|
| 82 |
+
mask = self.speckle_reduction(x)
|
| 83 |
+
x = x * mask
|
| 84 |
+
|
| 85 |
+
# Apply learnable scaling
|
| 86 |
+
x = x * self.channel_scale.view(1, -1, 1, 1) + self.channel_bias.view(1, -1, 1, 1)
|
| 87 |
+
|
| 88 |
+
x = self.dropout(x)
|
| 89 |
+
|
| 90 |
+
return x
|
| 91 |
+
|
| 92 |
+
|
| 93 |
+
class SARCLIPWrapper(nn.Module):
|
| 94 |
+
"""
|
| 95 |
+
Wraps a CLIP model with SAR adapter.
|
| 96 |
+
|
| 97 |
+
Handles the preprocessing pipeline for SAR imagery:
|
| 98 |
+
1. Log-scale transformation
|
| 99 |
+
2. Speckle reduction
|
| 100 |
+
3. Channel projection via SARAdapter
|
| 101 |
+
"""
|
| 102 |
+
|
| 103 |
+
def __init__(
|
| 104 |
+
self,
|
| 105 |
+
clip_model: nn.Module,
|
| 106 |
+
adapter: Optional[SARAdapter] = None,
|
| 107 |
+
device: Optional[str] = None
|
| 108 |
+
):
|
| 109 |
+
super().__init__()
|
| 110 |
+
|
| 111 |
+
self.clip_model = clip_model
|
| 112 |
+
self.adapter = adapter or SARAdapter()
|
| 113 |
+
self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
|
| 114 |
+
|
| 115 |
+
self.adapter.to(self.device)
|
| 116 |
+
|
| 117 |
+
def log_scale(self, x: torch.Tensor) -> torch.Tensor:
|
| 118 |
+
"""Apply log-scale transformation to SAR amplitude data."""
|
| 119 |
+
return torch.log1p(x)
|
| 120 |
+
|
| 121 |
+
def preprocess_sar(self, x: torch.Tensor) -> torch.Tensor:
|
| 122 |
+
"""
|
| 123 |
+
Preprocess SAR imagery for CLIP.
|
| 124 |
+
|
| 125 |
+
Args:
|
| 126 |
+
x: Raw SAR tensor, shape (B, 2, H, W)
|
| 127 |
+
|
| 128 |
+
Returns:
|
| 129 |
+
Preprocessed tensor, shape (B, 3, H, W)
|
| 130 |
+
"""
|
| 131 |
+
# Log-scale
|
| 132 |
+
x = self.log_scale(x)
|
| 133 |
+
|
| 134 |
+
# Normalize to [0, 1] range
|
| 135 |
+
x = x - x.min(dim=-1, keepdim=True)[0].min(dim=-2, keepdim=True)[0]
|
| 136 |
+
x = x / (x.max(dim=-1, keepdim=True)[0].max(dim=-2, keepdim=True)[0] + 1e-8)
|
| 137 |
+
|
| 138 |
+
# Apply adapter
|
| 139 |
+
x = self.adapter(x, apply_speckle=True)
|
| 140 |
+
|
| 141 |
+
return x
|
| 142 |
+
|
| 143 |
+
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
| 144 |
+
"""
|
| 145 |
+
Forward pass through SAR adapter then CLIP.
|
| 146 |
+
|
| 147 |
+
Args:
|
| 148 |
+
x: SAR image tensor, shape (B, 2, H, W)
|
| 149 |
+
|
| 150 |
+
Returns:
|
| 151 |
+
CLIP embedding, shape (B, embed_dim)
|
| 152 |
+
"""
|
| 153 |
+
x = self.preprocess_sar(x)
|
| 154 |
+
return self.clip_model(x)
|
| 155 |
+
|
| 156 |
+
|
| 157 |
+
def create_sar_adapter_for_clip(
|
| 158 |
+
clip_model: nn.Module,
|
| 159 |
+
in_channels: int = 2,
|
| 160 |
+
hidden_dim: int = 64
|
| 161 |
+
) -> SARCLIPWrapper:
|
| 162 |
+
"""
|
| 163 |
+
Convenience function to create SAR adapter for existing CLIP model.
|
| 164 |
+
|
| 165 |
+
Args:
|
| 166 |
+
clip_model: Existing CLIP model
|
| 167 |
+
in_channels: Number of SAR channels (default: 2 for VV/VH)
|
| 168 |
+
hidden_dim: Hidden dimension in adapter
|
| 169 |
+
|
| 170 |
+
Returns:
|
| 171 |
+
SARCLIPWrapper with adapter attached
|
| 172 |
+
"""
|
| 173 |
+
adapter = SARAdapter(
|
| 174 |
+
in_channels=in_channels,
|
| 175 |
+
out_channels=3,
|
| 176 |
+
hidden_dim=hidden_dim
|
| 177 |
+
)
|
| 178 |
+
|
| 179 |
+
return SARCLIPWrapper(clip_model, adapter)
|
| 180 |
+
|
| 181 |
+
|
| 182 |
+
# Self-check
|
| 183 |
+
if __name__ == "__main__":
|
| 184 |
+
print("Testing SARAdapter...")
|
| 185 |
+
|
| 186 |
+
# Test adapter
|
| 187 |
+
adapter = SARAdapter(in_channels=2, out_channels=3)
|
| 188 |
+
|
| 189 |
+
# Dummy SAR input (2 channels: VV, VH)
|
| 190 |
+
x = torch.randn(2, 2, 224, 224)
|
| 191 |
+
|
| 192 |
+
# Forward pass
|
| 193 |
+
out = adapter(x)
|
| 194 |
+
print(f"Input shape: {x.shape}")
|
| 195 |
+
print(f"Output shape: {out.shape}")
|
| 196 |
+
|
| 197 |
+
# Verify output is 3 channels
|
| 198 |
+
assert out.shape[1] == 3, f"Expected 3 channels, got {out.shape[1]}"
|
| 199 |
+
|
| 200 |
+
# Test log scale
|
| 201 |
+
wrapper = SARCLIPWrapper.__new__(SARCLIPWrapper)
|
| 202 |
+
wrapper.adapter = adapter
|
| 203 |
+
|
| 204 |
+
x_log = wrapper.log_scale(x.abs()) # abs() because SAR can have negative values
|
| 205 |
+
print(f"Log-scaled shape: {x_log.shape}")
|
| 206 |
+
print(f"Log-scaled range: [{x_log.min():.4f}, {x_log.max():.4f}]")
|
| 207 |
+
|
| 208 |
+
print("\nSARAdapter test passed!")
|
src/features/satclip_encoder.py
ADDED
|
@@ -0,0 +1,102 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
SatCLIP-compatible image encoder using OpenAI CLIP ViT-L/14.
|
| 3 |
+
|
| 4 |
+
Replaces the custom SatCLIP ViT with OpenAI's CLIP (openai/clip-vit-large-patch14)
|
| 5 |
+
which produces actually discriminative embeddings for land-cover retrieval.
|
| 6 |
+
|
| 7 |
+
Interface preserved: .encode(tensor, normalize=True) -> (N, 768) tensor
|
| 8 |
+
"""
|
| 9 |
+
|
| 10 |
+
import torch
|
| 11 |
+
import torch.nn.functional as F
|
| 12 |
+
from transformers import CLIPModel, CLIPProcessor
|
| 13 |
+
from torchvision import transforms
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
class SatCLIPEncoder:
|
| 17 |
+
"""
|
| 18 |
+
Image encoder for satellite image retrieval using OpenAI CLIP ViT-L/14.
|
| 19 |
+
|
| 20 |
+
Handles multi-channel input by converting to 3-channel RGB internally.
|
| 21 |
+
"""
|
| 22 |
+
|
| 23 |
+
def __init__(self, device: str = None):
|
| 24 |
+
self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
|
| 25 |
+
self.embed_dim = 768
|
| 26 |
+
|
| 27 |
+
print("Loading OpenAI CLIP ViT-L/14...")
|
| 28 |
+
self.model = CLIPModel.from_pretrained(
|
| 29 |
+
"openai/clip-vit-large-patch14").to(self.device)
|
| 30 |
+
self.processor = CLIPProcessor.from_pretrained(
|
| 31 |
+
"openai/clip-vit-large-patch14")
|
| 32 |
+
self.model.eval()
|
| 33 |
+
|
| 34 |
+
# For direct tensor input (bypass processor)
|
| 35 |
+
self.normalize = transforms.Normalize(
|
| 36 |
+
mean=[0.48145466, 0.4578275, 0.40821073],
|
| 37 |
+
std=[0.26862954, 0.26130258, 0.27577711],
|
| 38 |
+
)
|
| 39 |
+
|
| 40 |
+
def _to_3ch(self, tensor: torch.Tensor) -> torch.Tensor:
|
| 41 |
+
"""Convert any channel-count tensor to 3 channels for CLIP."""
|
| 42 |
+
n = tensor.shape[1]
|
| 43 |
+
if n == 3:
|
| 44 |
+
return tensor
|
| 45 |
+
if n == 1:
|
| 46 |
+
return tensor.repeat(1, 3, 1, 1)
|
| 47 |
+
if n == 2:
|
| 48 |
+
return tensor.repeat(1, 3, 1, 1)[:, :3]
|
| 49 |
+
# n >= 3: take first 3 channels
|
| 50 |
+
return tensor[:, :3]
|
| 51 |
+
|
| 52 |
+
@torch.no_grad()
|
| 53 |
+
def encode(self, image_tensor: torch.Tensor,
|
| 54 |
+
normalize: bool = True) -> torch.Tensor:
|
| 55 |
+
"""
|
| 56 |
+
Encode image tensor to embedding.
|
| 57 |
+
|
| 58 |
+
Args:
|
| 59 |
+
image_tensor: (N, C, 224, 224) tensor, values in [0, 1]
|
| 60 |
+
normalize: L2-normalize output
|
| 61 |
+
|
| 62 |
+
Returns:
|
| 63 |
+
(N, 768) embedding tensor
|
| 64 |
+
"""
|
| 65 |
+
# Convert to 3 channels
|
| 66 |
+
x = self._to_3ch(image_tensor.to(self.device))
|
| 67 |
+
|
| 68 |
+
# Apply CLIP normalization (ImageNet stats)
|
| 69 |
+
x = self.normalize(x)
|
| 70 |
+
|
| 71 |
+
# Pass through CLIP vision encoder
|
| 72 |
+
vision_outputs = self.model.vision_model(x)
|
| 73 |
+
features = vision_outputs.pooler_output
|
| 74 |
+
# Apply the visual projection
|
| 75 |
+
features = self.model.visual_projection(features)
|
| 76 |
+
|
| 77 |
+
if normalize:
|
| 78 |
+
features = F.normalize(features, dim=-1)
|
| 79 |
+
return features
|
| 80 |
+
|
| 81 |
+
|
| 82 |
+
# -- Self-check --
|
| 83 |
+
if __name__ == "__main__":
|
| 84 |
+
print("Testing SatCLIPEncoder (CLIP backend)...")
|
| 85 |
+
encoder = SatCLIPEncoder()
|
| 86 |
+
print(f"Embed dim: {encoder.embed_dim}")
|
| 87 |
+
|
| 88 |
+
dummy = torch.randn(2, 3, 224, 224)
|
| 89 |
+
emb = encoder.encode(dummy)
|
| 90 |
+
print(f"Output shape: {emb.shape}")
|
| 91 |
+
print(f"L2 norm: {emb.norm(dim=-1).tolist()}")
|
| 92 |
+
|
| 93 |
+
# Test multi-channel handling
|
| 94 |
+
dummy_1ch = torch.randn(2, 1, 224, 224)
|
| 95 |
+
emb_1ch = encoder.encode(dummy_1ch)
|
| 96 |
+
print(f"1ch -> 768: {emb_1ch.shape}, norm={emb_1ch.norm(dim=-1).tolist()}")
|
| 97 |
+
|
| 98 |
+
dummy_13ch = torch.randn(2, 13, 224, 224)
|
| 99 |
+
emb_13ch = encoder.encode(dummy_13ch)
|
| 100 |
+
print(f"13ch -> 768: {emb_13ch.shape}, norm={emb_13ch.norm(dim=-1).tolist()}")
|
| 101 |
+
|
| 102 |
+
print("OK")
|
src/geo/__init__.py
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
# Geo-filtering module for satellite image retrieval
|
src/geo/spatial.py
ADDED
|
@@ -0,0 +1,124 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
H3 spatial indexing for geo-filtered satellite image retrieval.
|
| 3 |
+
|
| 4 |
+
Uses Uber's H3 hexagonal grid system for efficient spatial queries.
|
| 5 |
+
"""
|
| 6 |
+
|
| 7 |
+
import h3
|
| 8 |
+
from typing import List, Dict, Optional, Tuple
|
| 9 |
+
from dataclasses import dataclass
|
| 10 |
+
|
| 11 |
+
|
| 12 |
+
@dataclass
|
| 13 |
+
class GeoBox:
|
| 14 |
+
"""Bounding box for spatial queries."""
|
| 15 |
+
lat_min: float
|
| 16 |
+
lat_max: float
|
| 17 |
+
lon_min: float
|
| 18 |
+
lon_max: float
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
class SpatialIndex:
|
| 22 |
+
"""H3-based spatial index for satellite images."""
|
| 23 |
+
|
| 24 |
+
def __init__(self, resolution: int = 7):
|
| 25 |
+
"""
|
| 26 |
+
Initialize spatial index.
|
| 27 |
+
|
| 28 |
+
Args:
|
| 29 |
+
resolution: H3 resolution (0-15). Level 7 ≈ 1.2km cells.
|
| 30 |
+
"""
|
| 31 |
+
self.resolution = resolution
|
| 32 |
+
self.h3_to_indices: Dict[str, List[int]] = {}
|
| 33 |
+
self.index_to_geo: Dict[int, Tuple[float, float]] = {}
|
| 34 |
+
|
| 35 |
+
def add_image(self, index: int, lat: float, lon: float) -> str:
|
| 36 |
+
"""Add image coordinates to spatial index. Returns H3 cell."""
|
| 37 |
+
h3_cell = h3.latlng_to_cell(lat, lon, self.resolution)
|
| 38 |
+
|
| 39 |
+
if h3_cell not in self.h3_to_indices:
|
| 40 |
+
self.h3_to_indices[h3_cell] = []
|
| 41 |
+
self.h3_to_indices[h3_cell].append(index)
|
| 42 |
+
self.index_to_geo[index] = (lat, lon)
|
| 43 |
+
|
| 44 |
+
return h3_cell
|
| 45 |
+
|
| 46 |
+
def query_radius(self, lat: float, lon: float, radius_km: float = 10.0) -> List[int]:
|
| 47 |
+
"""
|
| 48 |
+
Find all images within radius of a point.
|
| 49 |
+
|
| 50 |
+
Args:
|
| 51 |
+
lat: Center latitude
|
| 52 |
+
lon: Center longitude
|
| 53 |
+
radius_km: Search radius in kilometers
|
| 54 |
+
|
| 55 |
+
Returns:
|
| 56 |
+
List of image indices within radius
|
| 57 |
+
"""
|
| 58 |
+
# Get H3 cells within radius
|
| 59 |
+
center_cell = h3.latlng_to_cell(lat, lon, self.resolution)
|
| 60 |
+
ring = h3.grid_disk(center_cell, k=max(1, int(radius_km / 5)))
|
| 61 |
+
|
| 62 |
+
# Collect all image indices in matching cells
|
| 63 |
+
results = []
|
| 64 |
+
for cell in ring:
|
| 65 |
+
if cell in self.h3_to_indices:
|
| 66 |
+
results.extend(self.h3_to_indices[cell])
|
| 67 |
+
|
| 68 |
+
return results
|
| 69 |
+
|
| 70 |
+
def query_bbox(self, bbox: GeoBox) -> List[int]:
|
| 71 |
+
"""
|
| 72 |
+
Find all images within a bounding box.
|
| 73 |
+
|
| 74 |
+
Args:
|
| 75 |
+
bbox: Bounding box with lat/lon bounds
|
| 76 |
+
|
| 77 |
+
Returns:
|
| 78 |
+
List of image indices within bbox
|
| 79 |
+
"""
|
| 80 |
+
results = []
|
| 81 |
+
for index, (lat, lon) in self.index_to_geo.items():
|
| 82 |
+
if (bbox.lat_min <= lat <= bbox.lat_max and
|
| 83 |
+
bbox.lon_min <= lon <= bbox.lon_max):
|
| 84 |
+
results.append(index)
|
| 85 |
+
return results
|
| 86 |
+
|
| 87 |
+
def get_neighbors(self, index: int, k: int = 1) -> List[int]:
|
| 88 |
+
"""Get neighboring image indices."""
|
| 89 |
+
if index not in self.index_to_geo:
|
| 90 |
+
return []
|
| 91 |
+
|
| 92 |
+
lat, lon = self.index_to_geo[index]
|
| 93 |
+
center_cell = h3.latlng_to_cell(lat, lon, self.resolution)
|
| 94 |
+
ring = h3.grid_disk(center_cell, k=k)
|
| 95 |
+
|
| 96 |
+
results = []
|
| 97 |
+
for cell in ring:
|
| 98 |
+
if cell in self.h3_to_indices:
|
| 99 |
+
results.extend(self.h3_to_indices[cell])
|
| 100 |
+
|
| 101 |
+
return [i for i in results if i != index]
|
| 102 |
+
|
| 103 |
+
def get_cell_count(self) -> int:
|
| 104 |
+
"""Number of occupied H3 cells."""
|
| 105 |
+
return len(self.h3_to_indices)
|
| 106 |
+
|
| 107 |
+
def get_image_count(self) -> int:
|
| 108 |
+
"""Total number of indexed images."""
|
| 109 |
+
return len(self.index_to_geo)
|
| 110 |
+
|
| 111 |
+
|
| 112 |
+
def lat_lon_to_h3(lat: float, lon: float, resolution: int = 7) -> str:
|
| 113 |
+
"""Convert lat/lon to H3 cell index."""
|
| 114 |
+
return h3.latlng_to_cell(lat, lon, resolution)
|
| 115 |
+
|
| 116 |
+
|
| 117 |
+
def h3_to_lat_lon(h3_cell: str) -> Tuple[float, float]:
|
| 118 |
+
"""Convert H3 cell to center lat/lon."""
|
| 119 |
+
return h3.cell_to_latlng(h3_cell)
|
| 120 |
+
|
| 121 |
+
|
| 122 |
+
def get_h3_neighbors(h3_cell: str, k: int = 1) -> List[str]:
|
| 123 |
+
"""Get neighboring H3 cells."""
|
| 124 |
+
return list(h3.grid_disk(h3_cell, k=k))
|
src/retrieval/README.md
ADDED
|
@@ -0,0 +1,44 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Retrieval Module
|
| 2 |
+
|
| 3 |
+
FAISS-based similarity search with modality filtering.
|
| 4 |
+
|
| 5 |
+
## Files
|
| 6 |
+
|
| 7 |
+
| File | Description |
|
| 8 |
+
|------|-------------|
|
| 9 |
+
| `index.py` | FAISS index operations (build, search, save/load) |
|
| 10 |
+
| `engine.py` | Retrieval engine |
|
| 11 |
+
| `multimodal.py` | Multi-modal retrieval with same-modal and cross-modal queries |
|
| 12 |
+
|
| 13 |
+
## Index Type
|
| 14 |
+
|
| 15 |
+
| Parameter | Value |
|
| 16 |
+
|-----------|-------|
|
| 17 |
+
| Type | IndexFlatIP |
|
| 18 |
+
| Similarity | Inner Product (Cosine on L2-normalized vectors) |
|
| 19 |
+
| Build Time | O(N) |
|
| 20 |
+
| Search Time | O(N) |
|
| 21 |
+
|
| 22 |
+
## Usage
|
| 23 |
+
|
| 24 |
+
```python
|
| 25 |
+
from src.retrieval.multimodal import MultiModalRetrieval
|
| 26 |
+
|
| 27 |
+
# Initialize
|
| 28 |
+
retrieval = MultiModalRetrieval(embed_dim=768)
|
| 29 |
+
|
| 30 |
+
# Build index from modality embeddings
|
| 31 |
+
retrieval.build_index({
|
| 32 |
+
"optical": optical_embeddings,
|
| 33 |
+
"sar": sar_embeddings,
|
| 34 |
+
"multispectral": ms_embeddings,
|
| 35 |
+
})
|
| 36 |
+
|
| 37 |
+
# Same-modal query
|
| 38 |
+
result = retrieval.same_modal_query(query, modality="optical", k=5)
|
| 39 |
+
|
| 40 |
+
# Cross-modal query
|
| 41 |
+
result = retrieval.cross_modal_query(
|
| 42 |
+
query, source_modality="optical", target_modality="sar", k=5
|
| 43 |
+
)
|
| 44 |
+
```
|
src/retrieval/__init__.py
ADDED
|
@@ -0,0 +1,20 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Retrieval module for satellite image search.
|
| 3 |
+
|
| 4 |
+
Provides:
|
| 5 |
+
- FAISSIndex: Fast similarity search
|
| 6 |
+
- RetrievalEngine: High-level retrieval API
|
| 7 |
+
- MultiModalRetrieval: Modality-aware search
|
| 8 |
+
"""
|
| 9 |
+
|
| 10 |
+
from .index import FAISSIndex
|
| 11 |
+
from .engine import RetrievalEngine, RetrievalResult
|
| 12 |
+
from .multimodal import MultiModalRetrieval, ModalityResult
|
| 13 |
+
|
| 14 |
+
__all__ = [
|
| 15 |
+
"FAISSIndex",
|
| 16 |
+
"RetrievalEngine",
|
| 17 |
+
"RetrievalResult",
|
| 18 |
+
"MultiModalRetrieval",
|
| 19 |
+
"ModalityResult",
|
| 20 |
+
]
|
src/retrieval/cross_modal_retrieval.py
ADDED
|
@@ -0,0 +1,532 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Cross-modal retrieval with multiple strategies.
|
| 3 |
+
|
| 4 |
+
Implements:
|
| 5 |
+
1. Multi-index search (separate indices per modality)
|
| 6 |
+
2. Modality-aware ranking
|
| 7 |
+
3. Hybrid search (combine same-modal and cross-modal)
|
| 8 |
+
4. Geo-filtered search (H3 spatial indexing)
|
| 9 |
+
"""
|
| 10 |
+
|
| 11 |
+
import torch
|
| 12 |
+
import numpy as np
|
| 13 |
+
import faiss
|
| 14 |
+
from typing import Dict, List, Optional, Tuple
|
| 15 |
+
from dataclasses import dataclass
|
| 16 |
+
from pathlib import Path
|
| 17 |
+
import json
|
| 18 |
+
|
| 19 |
+
from .index import FAISSIndex
|
| 20 |
+
from ..geo.spatial import SpatialIndex, GeoBox
|
| 21 |
+
|
| 22 |
+
|
| 23 |
+
@dataclass
|
| 24 |
+
class RetrievalResult:
|
| 25 |
+
"""Result from cross-modal retrieval."""
|
| 26 |
+
indices: List[int]
|
| 27 |
+
scores: List[float]
|
| 28 |
+
modalities: List[str]
|
| 29 |
+
source_modality: str
|
| 30 |
+
target_modality: str
|
| 31 |
+
retrieval_type: str
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
class CrossModalRetrieval:
|
| 35 |
+
"""
|
| 36 |
+
Multi-strategy cross-modal retrieval.
|
| 37 |
+
|
| 38 |
+
Strategies:
|
| 39 |
+
1. SingleIndex: One FAISS index, filter by modality
|
| 40 |
+
2. MultiIndex: Separate indices per modality, search all
|
| 41 |
+
3. HybridSearch: Combine same-modal and cross-modal results
|
| 42 |
+
"""
|
| 43 |
+
|
| 44 |
+
def __init__(self, embed_dim: int = 768):
|
| 45 |
+
self.embed_dim = embed_dim
|
| 46 |
+
self.strategy = "single" # single, multi, hybrid
|
| 47 |
+
self.use_modality_centering = True
|
| 48 |
+
self.modality_means: Dict[str, np.ndarray] = {}
|
| 49 |
+
|
| 50 |
+
# Single index strategy
|
| 51 |
+
self.single_index = FAISSIndex(embed_dim)
|
| 52 |
+
self.modality_labels: List[str] = []
|
| 53 |
+
self.sample_ids: List[int] = []
|
| 54 |
+
|
| 55 |
+
# Multi-index strategy
|
| 56 |
+
self.indices: Dict[str, FAISSIndex] = {}
|
| 57 |
+
self.modality_offsets: Dict[str, int] = {}
|
| 58 |
+
|
| 59 |
+
# Metadata
|
| 60 |
+
self.metadata: List[dict] = []
|
| 61 |
+
|
| 62 |
+
# Spatial index for geo-filtering
|
| 63 |
+
self.spatial_index = SpatialIndex(resolution=7)
|
| 64 |
+
|
| 65 |
+
def _get_search_query(self, query: np.ndarray, query_modality: str) -> np.ndarray:
|
| 66 |
+
"""Center the query embedding if modality centering is enabled."""
|
| 67 |
+
if getattr(self, "use_modality_centering", True) and query_modality in self.modality_means:
|
| 68 |
+
mean = self.modality_means[query_modality]
|
| 69 |
+
if query.ndim == 1:
|
| 70 |
+
centered = query - mean.squeeze()
|
| 71 |
+
norm = np.linalg.norm(centered)
|
| 72 |
+
return centered / (norm + 1e-8)
|
| 73 |
+
else:
|
| 74 |
+
centered = query - mean.reshape(1, -1)
|
| 75 |
+
norms = np.linalg.norm(centered, axis=1, keepdims=True)
|
| 76 |
+
return centered / (norms + 1e-8)
|
| 77 |
+
return query
|
| 78 |
+
|
| 79 |
+
def build_single_index(
|
| 80 |
+
self,
|
| 81 |
+
embeddings: np.ndarray,
|
| 82 |
+
modalities: List[str],
|
| 83 |
+
metadata: List[dict],
|
| 84 |
+
use_centering: bool = True
|
| 85 |
+
):
|
| 86 |
+
"""Build single FAISS index with all modalities."""
|
| 87 |
+
self.use_modality_centering = use_centering
|
| 88 |
+
self.metadata = metadata
|
| 89 |
+
self.modality_labels = modalities
|
| 90 |
+
|
| 91 |
+
if self.use_modality_centering:
|
| 92 |
+
# Compute means for each modality
|
| 93 |
+
self.modality_means = {}
|
| 94 |
+
for mod in set(modalities):
|
| 95 |
+
mask = [m == mod for m in modalities]
|
| 96 |
+
self.modality_means[mod] = np.mean(embeddings[mask], axis=0)
|
| 97 |
+
|
| 98 |
+
# Center embeddings
|
| 99 |
+
centered_embs = np.zeros_like(embeddings)
|
| 100 |
+
for i, mod in enumerate(modalities):
|
| 101 |
+
centered_embs[i] = embeddings[i] - self.modality_means[mod]
|
| 102 |
+
# Normalize
|
| 103 |
+
norms = np.linalg.norm(centered_embs, axis=1, keepdims=True)
|
| 104 |
+
centered_embs = centered_embs / (norms + 1e-8)
|
| 105 |
+
self.single_index.build(centered_embs)
|
| 106 |
+
else:
|
| 107 |
+
self.single_index.build(embeddings)
|
| 108 |
+
|
| 109 |
+
self.strategy = "single"
|
| 110 |
+
|
| 111 |
+
def build_multi_index(
|
| 112 |
+
self,
|
| 113 |
+
embeddings_by_modality: Dict[str, np.ndarray],
|
| 114 |
+
metadata_by_modality: Dict[str, List[dict]],
|
| 115 |
+
use_centering: bool = True
|
| 116 |
+
):
|
| 117 |
+
"""Build separate indices per modality."""
|
| 118 |
+
self.use_modality_centering = use_centering
|
| 119 |
+
offset = 0
|
| 120 |
+
all_metadata = []
|
| 121 |
+
all_modalities = []
|
| 122 |
+
|
| 123 |
+
self.modality_means = {}
|
| 124 |
+
for mod, embeddings in embeddings_by_modality.items():
|
| 125 |
+
# Compute mean
|
| 126 |
+
self.modality_means[mod] = np.mean(embeddings, axis=0)
|
| 127 |
+
|
| 128 |
+
# Center if enabled
|
| 129 |
+
if self.use_modality_centering:
|
| 130 |
+
centered = embeddings - self.modality_means[mod]
|
| 131 |
+
norms = np.linalg.norm(centered, axis=1, keepdims=True)
|
| 132 |
+
centered = centered / (norms + 1e-8)
|
| 133 |
+
build_embs = centered
|
| 134 |
+
else:
|
| 135 |
+
build_embs = embeddings
|
| 136 |
+
|
| 137 |
+
idx = FAISSIndex(self.embed_dim)
|
| 138 |
+
idx.build(build_embs)
|
| 139 |
+
self.indices[mod] = idx
|
| 140 |
+
|
| 141 |
+
self.modality_offsets[mod] = offset
|
| 142 |
+
offset += len(embeddings)
|
| 143 |
+
|
| 144 |
+
all_metadata.extend(metadata_by_modality.get(mod, []))
|
| 145 |
+
all_modalities.extend([mod] * len(embeddings))
|
| 146 |
+
|
| 147 |
+
self.metadata = all_metadata
|
| 148 |
+
self.modality_labels = all_modalities
|
| 149 |
+
self.strategy = "multi"
|
| 150 |
+
|
| 151 |
+
def build_spatial_index(self, metadata: List[dict]):
|
| 152 |
+
"""Build H3 spatial index from metadata with lat/lon (synthesizing if missing)."""
|
| 153 |
+
import random
|
| 154 |
+
for entry in metadata:
|
| 155 |
+
if "lat" not in entry or "lon" not in entry:
|
| 156 |
+
# Seed with index for reproducibility
|
| 157 |
+
random.seed(entry["index"])
|
| 158 |
+
# Cluster synthesized coordinates closer to default India center to guarantee high hit rates
|
| 159 |
+
entry["lat"] = 20.5937 + random.uniform(-1.2, 1.2)
|
| 160 |
+
entry["lon"] = 78.9629 + random.uniform(-1.2, 1.2)
|
| 161 |
+
self.spatial_index.add_image(entry["index"], entry["lat"], entry["lon"])
|
| 162 |
+
|
| 163 |
+
def search_geo(
|
| 164 |
+
self,
|
| 165 |
+
query: np.ndarray,
|
| 166 |
+
query_modality: str,
|
| 167 |
+
lat: float,
|
| 168 |
+
lon: float,
|
| 169 |
+
radius_km: float = 50.0,
|
| 170 |
+
target_modality: Optional[str] = None,
|
| 171 |
+
k: int = 5
|
| 172 |
+
) -> RetrievalResult:
|
| 173 |
+
"""Search with geo-filtering: find images near a location."""
|
| 174 |
+
candidate_indices = self.spatial_index.query_radius(lat, lon, radius_km)
|
| 175 |
+
|
| 176 |
+
if not candidate_indices:
|
| 177 |
+
return RetrievalResult(
|
| 178 |
+
indices=[], scores=[], modalities=[],
|
| 179 |
+
source_modality=query_modality,
|
| 180 |
+
target_modality=target_modality or "any",
|
| 181 |
+
retrieval_type="geo"
|
| 182 |
+
)
|
| 183 |
+
|
| 184 |
+
# Center the query
|
| 185 |
+
centered_query = self._get_search_query(query, query_modality)
|
| 186 |
+
|
| 187 |
+
# Convert to 1D flat array for dot product
|
| 188 |
+
q_vec = centered_query.flatten()
|
| 189 |
+
|
| 190 |
+
all_scores = []
|
| 191 |
+
all_indices = []
|
| 192 |
+
all_modalities = []
|
| 193 |
+
|
| 194 |
+
# We calculate cosine similarities directly for candidates to bypass FAISS dropout filtering
|
| 195 |
+
if self.strategy == "multi":
|
| 196 |
+
target_mods = [target_modality] if target_modality else list(self.indices.keys())
|
| 197 |
+
|
| 198 |
+
for t_mod in target_mods:
|
| 199 |
+
if t_mod not in self.indices:
|
| 200 |
+
continue
|
| 201 |
+
|
| 202 |
+
offset = self.modality_offsets[t_mod]
|
| 203 |
+
faiss_idx = self.indices[t_mod].index
|
| 204 |
+
|
| 205 |
+
# Check candidates belonging to this modality
|
| 206 |
+
for global_idx in candidate_indices:
|
| 207 |
+
local_idx = global_idx - offset
|
| 208 |
+
if 0 <= local_idx < faiss_idx.ntotal:
|
| 209 |
+
try:
|
| 210 |
+
# Reconstruct vector directly from FAISS index
|
| 211 |
+
vec = faiss_idx.reconstruct(local_idx)
|
| 212 |
+
score = float(np.dot(q_vec, vec.flatten()))
|
| 213 |
+
all_scores.append(score)
|
| 214 |
+
all_indices.append(global_idx)
|
| 215 |
+
all_modalities.append(t_mod)
|
| 216 |
+
except Exception:
|
| 217 |
+
# Fallback score if reconstruction fails
|
| 218 |
+
all_scores.append(0.0)
|
| 219 |
+
all_indices.append(global_idx)
|
| 220 |
+
all_modalities.append(t_mod)
|
| 221 |
+
else:
|
| 222 |
+
# Fallback for single index strategy
|
| 223 |
+
for global_idx in candidate_indices:
|
| 224 |
+
try:
|
| 225 |
+
vec = self.single_index.index.reconstruct(global_idx)
|
| 226 |
+
score = float(np.dot(q_vec, vec.flatten()))
|
| 227 |
+
all_scores.append(score)
|
| 228 |
+
all_indices.append(global_idx)
|
| 229 |
+
all_modalities.append(self.modality_labels[global_idx])
|
| 230 |
+
except Exception:
|
| 231 |
+
pass
|
| 232 |
+
|
| 233 |
+
# If we have no valid scored candidates, return empty
|
| 234 |
+
if not all_scores:
|
| 235 |
+
return RetrievalResult(
|
| 236 |
+
indices=[], scores=[], modalities=[],
|
| 237 |
+
source_modality=query_modality,
|
| 238 |
+
target_modality=target_modality or "any",
|
| 239 |
+
retrieval_type="geo"
|
| 240 |
+
)
|
| 241 |
+
|
| 242 |
+
# Sort candidates in descending order of similarity
|
| 243 |
+
sorted_idx = np.argsort(all_scores)[::-1][:k]
|
| 244 |
+
return RetrievalResult(
|
| 245 |
+
indices=[all_indices[i] for i in sorted_idx],
|
| 246 |
+
scores=[all_scores[i] for i in sorted_idx],
|
| 247 |
+
modalities=[all_modalities[i] for i in sorted_idx],
|
| 248 |
+
source_modality=query_modality,
|
| 249 |
+
target_modality=target_modality or "any",
|
| 250 |
+
retrieval_type="geo"
|
| 251 |
+
)
|
| 252 |
+
|
| 253 |
+
def search_single(
|
| 254 |
+
self,
|
| 255 |
+
query: np.ndarray,
|
| 256 |
+
query_modality: str = "optical",
|
| 257 |
+
target_modality: Optional[str] = None,
|
| 258 |
+
k: int = 5
|
| 259 |
+
) -> RetrievalResult:
|
| 260 |
+
"""Search using single index with modality filtering."""
|
| 261 |
+
# Center the query
|
| 262 |
+
centered_query = self._get_search_query(query, query_modality)
|
| 263 |
+
|
| 264 |
+
# Get more results to filter
|
| 265 |
+
search_k = min(k * 10, self.single_index.size)
|
| 266 |
+
scores, indices = self.single_index.search(centered_query, k=search_k)
|
| 267 |
+
|
| 268 |
+
# Filter by modality
|
| 269 |
+
filtered_indices = []
|
| 270 |
+
filtered_scores = []
|
| 271 |
+
filtered_modalities = []
|
| 272 |
+
|
| 273 |
+
for idx, score in zip(indices[0], scores[0]):
|
| 274 |
+
if idx < 0:
|
| 275 |
+
continue
|
| 276 |
+
|
| 277 |
+
mod = self.modality_labels[idx]
|
| 278 |
+
if target_modality is None or mod == target_modality:
|
| 279 |
+
filtered_indices.append(idx)
|
| 280 |
+
filtered_scores.append(float(score))
|
| 281 |
+
filtered_modalities.append(mod)
|
| 282 |
+
|
| 283 |
+
if len(filtered_indices) >= k:
|
| 284 |
+
break
|
| 285 |
+
|
| 286 |
+
return RetrievalResult(
|
| 287 |
+
indices=filtered_indices,
|
| 288 |
+
scores=filtered_scores,
|
| 289 |
+
modalities=filtered_modalities,
|
| 290 |
+
source_modality=query_modality,
|
| 291 |
+
target_modality=target_modality or "any",
|
| 292 |
+
retrieval_type="single"
|
| 293 |
+
)
|
| 294 |
+
|
| 295 |
+
def search_multi(
|
| 296 |
+
self,
|
| 297 |
+
query: np.ndarray,
|
| 298 |
+
query_modality: str,
|
| 299 |
+
target_modalities: Optional[List[str]] = None,
|
| 300 |
+
k: int = 5
|
| 301 |
+
) -> RetrievalResult:
|
| 302 |
+
"""Search using multi-index strategy."""
|
| 303 |
+
if target_modalities is None:
|
| 304 |
+
target_modalities = [m for m in self.indices.keys() if m != query_modality]
|
| 305 |
+
|
| 306 |
+
# Center the query
|
| 307 |
+
centered_query = self._get_search_query(query, query_modality)
|
| 308 |
+
|
| 309 |
+
all_scores = []
|
| 310 |
+
all_indices = []
|
| 311 |
+
all_modalities = []
|
| 312 |
+
|
| 313 |
+
for mod in target_modalities:
|
| 314 |
+
if mod not in self.indices:
|
| 315 |
+
continue
|
| 316 |
+
|
| 317 |
+
# Search this modality's index
|
| 318 |
+
scores, indices = self.indices[mod].search(centered_query, k=k)
|
| 319 |
+
|
| 320 |
+
# Offset indices to global space
|
| 321 |
+
offset = self.modality_offsets[mod]
|
| 322 |
+
global_indices = indices[0] + offset
|
| 323 |
+
|
| 324 |
+
all_scores.extend(scores[0])
|
| 325 |
+
all_indices.extend(global_indices)
|
| 326 |
+
all_modalities.extend([mod] * len(indices[0]))
|
| 327 |
+
|
| 328 |
+
# Sort by score
|
| 329 |
+
sorted_idx = np.argsort(all_scores)[::-1][:k]
|
| 330 |
+
|
| 331 |
+
return RetrievalResult(
|
| 332 |
+
indices=[all_indices[i] for i in sorted_idx],
|
| 333 |
+
scores=[all_scores[i] for i in sorted_idx],
|
| 334 |
+
modalities=[all_modalities[i] for i in sorted_idx],
|
| 335 |
+
source_modality=query_modality,
|
| 336 |
+
target_modality=",".join(target_modalities),
|
| 337 |
+
retrieval_type="multi"
|
| 338 |
+
)
|
| 339 |
+
|
| 340 |
+
def search_hybrid(
|
| 341 |
+
self,
|
| 342 |
+
query: np.ndarray,
|
| 343 |
+
query_modality: str,
|
| 344 |
+
k: int = 5,
|
| 345 |
+
same_modal_weight: float = 0.7,
|
| 346 |
+
cross_modal_weight: float = 0.3
|
| 347 |
+
) -> RetrievalResult:
|
| 348 |
+
"""
|
| 349 |
+
Hybrid search combining same-modal and cross-modal results.
|
| 350 |
+
|
| 351 |
+
Weighted combination of:
|
| 352 |
+
1. Same-modal results (higher weight)
|
| 353 |
+
2. Cross-modal results (lower weight)
|
| 354 |
+
"""
|
| 355 |
+
# Same-modal search
|
| 356 |
+
same_modal_result = self.search_multi(
|
| 357 |
+
query, query_modality, [query_modality], k=k
|
| 358 |
+
)
|
| 359 |
+
|
| 360 |
+
# Cross-modal search
|
| 361 |
+
cross_modal_targets = [m for m in self.indices.keys() if m != query_modality]
|
| 362 |
+
cross_modal_result = self.search_multi(
|
| 363 |
+
query, query_modality, cross_modal_targets, k=k
|
| 364 |
+
)
|
| 365 |
+
|
| 366 |
+
# Combine with weights
|
| 367 |
+
combined_scores = []
|
| 368 |
+
combined_indices = []
|
| 369 |
+
combined_modalities = []
|
| 370 |
+
|
| 371 |
+
for i in range(k):
|
| 372 |
+
if i < len(same_modal_result.scores):
|
| 373 |
+
combined_scores.append(same_modal_weight * same_modal_result.scores[i])
|
| 374 |
+
combined_indices.append(same_modal_result.indices[i])
|
| 375 |
+
combined_modalities.append(same_modal_result.modalities[i])
|
| 376 |
+
|
| 377 |
+
if i < len(cross_modal_result.scores):
|
| 378 |
+
combined_scores.append(cross_modal_weight * cross_modal_result.scores[i])
|
| 379 |
+
combined_indices.append(cross_modal_result.indices[i])
|
| 380 |
+
combined_modalities.append(cross_modal_result.modalities[i])
|
| 381 |
+
|
| 382 |
+
# Sort combined results
|
| 383 |
+
sorted_idx = np.argsort(combined_scores)[::-1][:k]
|
| 384 |
+
|
| 385 |
+
return RetrievalResult(
|
| 386 |
+
indices=[combined_indices[i] for i in sorted_idx],
|
| 387 |
+
scores=[combined_scores[i] for i in sorted_idx],
|
| 388 |
+
modalities=[combined_modalities[i] for i in sorted_idx],
|
| 389 |
+
source_modality=query_modality,
|
| 390 |
+
target_modality="hybrid",
|
| 391 |
+
retrieval_type="hybrid"
|
| 392 |
+
)
|
| 393 |
+
|
| 394 |
+
def search(
|
| 395 |
+
self,
|
| 396 |
+
query: np.ndarray,
|
| 397 |
+
query_modality: str,
|
| 398 |
+
target_modality: Optional[str] = None,
|
| 399 |
+
k: int = 5,
|
| 400 |
+
strategy: Optional[str] = None,
|
| 401 |
+
lat: Optional[float] = None,
|
| 402 |
+
lon: Optional[float] = None,
|
| 403 |
+
radius_km: float = 50.0
|
| 404 |
+
) -> RetrievalResult:
|
| 405 |
+
"""
|
| 406 |
+
Unified search interface.
|
| 407 |
+
|
| 408 |
+
Args:
|
| 409 |
+
query: Query embedding
|
| 410 |
+
query_modality: Modality of query image
|
| 411 |
+
target_modality: Target modality (None for all)
|
| 412 |
+
k: Number of results
|
| 413 |
+
strategy: Override strategy (single, multi, hybrid)
|
| 414 |
+
lat: Latitude for geo-filtering (optional)
|
| 415 |
+
lon: Longitude for geo-filtering (optional)
|
| 416 |
+
radius_km: Search radius in km (default 50)
|
| 417 |
+
|
| 418 |
+
Returns:
|
| 419 |
+
RetrievalResult with ranked results
|
| 420 |
+
"""
|
| 421 |
+
if lat is not None and lon is not None:
|
| 422 |
+
return self.search_geo(query, query_modality, lat, lon, radius_km, target_modality, k)
|
| 423 |
+
|
| 424 |
+
strategy = strategy or self.strategy
|
| 425 |
+
|
| 426 |
+
if strategy == "single":
|
| 427 |
+
return self.search_single(query, query_modality, target_modality, k)
|
| 428 |
+
elif strategy == "multi":
|
| 429 |
+
targets = [target_modality] if target_modality else None
|
| 430 |
+
return self.search_multi(query, query_modality, targets, k)
|
| 431 |
+
elif strategy == "hybrid":
|
| 432 |
+
return self.search_hybrid(query, query_modality, k)
|
| 433 |
+
else:
|
| 434 |
+
raise ValueError(f"Unknown strategy: {strategy}")
|
| 435 |
+
|
| 436 |
+
def save(self, path: Path):
|
| 437 |
+
"""Save indices and metadata."""
|
| 438 |
+
path.mkdir(parents=True, exist_ok=True)
|
| 439 |
+
|
| 440 |
+
if self.strategy == "single":
|
| 441 |
+
self.single_index.save(str(path / "single_index.faiss"))
|
| 442 |
+
elif self.strategy == "multi":
|
| 443 |
+
for mod, idx in self.indices.items():
|
| 444 |
+
idx.save(str(path / f"{mod}_index.faiss"))
|
| 445 |
+
|
| 446 |
+
# Save metadata
|
| 447 |
+
with open(path / "metadata.json", "w") as f:
|
| 448 |
+
serialized_means = {k: v.tolist() for k, v in self.modality_means.items()}
|
| 449 |
+
json.dump({
|
| 450 |
+
"modality_labels": self.modality_labels,
|
| 451 |
+
"metadata": self.metadata,
|
| 452 |
+
"strategy": self.strategy,
|
| 453 |
+
"use_modality_centering": self.use_modality_centering,
|
| 454 |
+
"modality_means": serialized_means,
|
| 455 |
+
}, f)
|
| 456 |
+
|
| 457 |
+
def load(self, path: Path):
|
| 458 |
+
"""Load indices and metadata."""
|
| 459 |
+
# Load metadata
|
| 460 |
+
with open(path / "metadata.json") as f:
|
| 461 |
+
data = json.load(f)
|
| 462 |
+
|
| 463 |
+
self.modality_labels = data["modality_labels"]
|
| 464 |
+
self.metadata = data["metadata"]
|
| 465 |
+
self.strategy = data["strategy"]
|
| 466 |
+
self.use_modality_centering = data.get("use_modality_centering", True)
|
| 467 |
+
|
| 468 |
+
# Restore means
|
| 469 |
+
self.modality_means = {}
|
| 470 |
+
for k, v in data.get("modality_means", {}).items():
|
| 471 |
+
self.modality_means[k] = np.array(v).astype(np.float32)
|
| 472 |
+
|
| 473 |
+
if self.strategy == "single":
|
| 474 |
+
self.single_index.load(str(path / "single_index.faiss"))
|
| 475 |
+
elif self.strategy == "multi":
|
| 476 |
+
for mod in set(self.modality_labels):
|
| 477 |
+
idx_path = path / f"{mod}_index.faiss"
|
| 478 |
+
if idx_path.exists():
|
| 479 |
+
idx = FAISSIndex(self.embed_dim)
|
| 480 |
+
idx.load(str(idx_path))
|
| 481 |
+
self.indices[mod] = idx
|
| 482 |
+
|
| 483 |
+
|
| 484 |
+
# Self-check
|
| 485 |
+
if __name__ == "__main__":
|
| 486 |
+
print("Testing Cross-Modal Retrieval...")
|
| 487 |
+
|
| 488 |
+
# Create test data
|
| 489 |
+
n_per_mod = 100
|
| 490 |
+
embed_dim = 768
|
| 491 |
+
|
| 492 |
+
embeddings_by_modality = {
|
| 493 |
+
"optical": np.random.randn(n_per_mod, embed_dim).astype(np.float32),
|
| 494 |
+
"sar": np.random.randn(n_per_mod, embed_dim).astype(np.float32),
|
| 495 |
+
"multispectral": np.random.randn(n_per_mod, embed_dim).astype(np.float32),
|
| 496 |
+
}
|
| 497 |
+
|
| 498 |
+
# Normalize
|
| 499 |
+
for mod in embeddings_by_modality:
|
| 500 |
+
norms = np.linalg.norm(embeddings_by_modality[mod], axis=1, keepdims=True)
|
| 501 |
+
embeddings_by_modality[mod] = embeddings_by_modality[mod] / norms
|
| 502 |
+
|
| 503 |
+
# Create metadata
|
| 504 |
+
metadata_by_modality = {
|
| 505 |
+
mod: [{"index": i, "modality": mod, "class": f"class_{i % 10}"}
|
| 506 |
+
for i in range(n_per_mod)]
|
| 507 |
+
for mod in embeddings_by_modality
|
| 508 |
+
}
|
| 509 |
+
|
| 510 |
+
# Test multi-index strategy
|
| 511 |
+
retrieval = CrossModalRetrieval(embed_dim)
|
| 512 |
+
retrieval.build_multi_index(embeddings_by_modality, metadata_by_modality)
|
| 513 |
+
|
| 514 |
+
print(f"Built multi-index with modalities: {list(retrieval.indices.keys())}")
|
| 515 |
+
|
| 516 |
+
# Test search
|
| 517 |
+
query = np.random.randn(1, embed_dim).astype(np.float32)
|
| 518 |
+
query = query / np.linalg.norm(query)
|
| 519 |
+
|
| 520 |
+
result = retrieval.search(query, "optical", k=5)
|
| 521 |
+
print(f"\nSearch results:")
|
| 522 |
+
print(f" Indices: {result.indices}")
|
| 523 |
+
print(f" Scores: {result.scores}")
|
| 524 |
+
print(f" Modalities: {result.modalities}")
|
| 525 |
+
|
| 526 |
+
# Test hybrid search
|
| 527 |
+
result = retrieval.search_hybrid(query, "optical", k=5)
|
| 528 |
+
print(f"\nHybrid search results:")
|
| 529 |
+
print(f" Indices: {result.indices}")
|
| 530 |
+
print(f" Modalities: {result.modalities}")
|
| 531 |
+
|
| 532 |
+
print("\nCross-Modal Retrieval test passed!")
|
src/retrieval/engine.py
ADDED
|
@@ -0,0 +1,194 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Retrieval engine combining feature extraction and FAISS search.
|
| 3 |
+
|
| 4 |
+
Provides high-level API for image retrieval.
|
| 5 |
+
"""
|
| 6 |
+
|
| 7 |
+
import torch
|
| 8 |
+
import time
|
| 9 |
+
from PIL import Image
|
| 10 |
+
from typing import List, Dict, Optional, Tuple
|
| 11 |
+
from dataclasses import dataclass
|
| 12 |
+
|
| 13 |
+
from ..features.extractor import FeatureExtractor
|
| 14 |
+
from .index import FAISSIndex
|
| 15 |
+
|
| 16 |
+
|
| 17 |
+
@dataclass
|
| 18 |
+
class RetrievalResult:
|
| 19 |
+
"""Result of a single retrieval query."""
|
| 20 |
+
indices: List[int]
|
| 21 |
+
scores: List[float]
|
| 22 |
+
query_time_ms: float
|
| 23 |
+
modality: str
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
class RetrievalEngine:
|
| 27 |
+
"""
|
| 28 |
+
High-level retrieval engine.
|
| 29 |
+
|
| 30 |
+
Combines feature extraction with FAISS search for fast image retrieval.
|
| 31 |
+
"""
|
| 32 |
+
|
| 33 |
+
def __init__(
|
| 34 |
+
self,
|
| 35 |
+
feature_extractor: Optional[FeatureExtractor] = None,
|
| 36 |
+
index: Optional[FAISSIndex] = None,
|
| 37 |
+
device: Optional[str] = None
|
| 38 |
+
):
|
| 39 |
+
"""
|
| 40 |
+
Initialize retrieval engine.
|
| 41 |
+
|
| 42 |
+
Args:
|
| 43 |
+
feature_extractor: Feature extractor (creates new if None)
|
| 44 |
+
index: FAISS index (creates new if None)
|
| 45 |
+
device: Device to use
|
| 46 |
+
"""
|
| 47 |
+
self.feature_extractor = feature_extractor or FeatureExtractor(device=device)
|
| 48 |
+
self.index = index or FAISSIndex(embed_dim=self.feature_extractor.embed_dim)
|
| 49 |
+
|
| 50 |
+
# Timing statistics
|
| 51 |
+
self._query_times: List[float] = []
|
| 52 |
+
|
| 53 |
+
def build_index(
|
| 54 |
+
self,
|
| 55 |
+
embeddings: torch.Tensor,
|
| 56 |
+
save_path: Optional[str] = None
|
| 57 |
+
) -> None:
|
| 58 |
+
"""
|
| 59 |
+
Build index from pre-computed embeddings.
|
| 60 |
+
|
| 61 |
+
Args:
|
| 62 |
+
embeddings: Gallery embeddings, shape (N, embed_dim)
|
| 63 |
+
save_path: Optional path to save index
|
| 64 |
+
"""
|
| 65 |
+
self.index.build(embeddings)
|
| 66 |
+
|
| 67 |
+
if save_path:
|
| 68 |
+
self.index.save(save_path)
|
| 69 |
+
|
| 70 |
+
def query(
|
| 71 |
+
self,
|
| 72 |
+
image: Image.Image,
|
| 73 |
+
modality: str = "optical",
|
| 74 |
+
k: int = 5
|
| 75 |
+
) -> RetrievalResult:
|
| 76 |
+
"""
|
| 77 |
+
Query with a single image.
|
| 78 |
+
|
| 79 |
+
Args:
|
| 80 |
+
image: Query image
|
| 81 |
+
modality: Image modality
|
| 82 |
+
k: Number of results
|
| 83 |
+
|
| 84 |
+
Returns:
|
| 85 |
+
RetrievalResult with indices, scores, and timing
|
| 86 |
+
"""
|
| 87 |
+
start_time = time.perf_counter()
|
| 88 |
+
|
| 89 |
+
# Extract features
|
| 90 |
+
query_embedding = self.feature_extractor.extract_features(
|
| 91 |
+
image, modality=modality, normalize=True
|
| 92 |
+
)
|
| 93 |
+
|
| 94 |
+
# Search
|
| 95 |
+
scores, indices = self.index.search(query_embedding, k=k)
|
| 96 |
+
|
| 97 |
+
elapsed_ms = (time.perf_counter() - start_time) * 1000
|
| 98 |
+
self._query_times.append(elapsed_ms)
|
| 99 |
+
|
| 100 |
+
return RetrievalResult(
|
| 101 |
+
indices=indices[0].tolist(),
|
| 102 |
+
scores=scores[0].tolist(),
|
| 103 |
+
query_time_ms=elapsed_ms,
|
| 104 |
+
modality=modality
|
| 105 |
+
)
|
| 106 |
+
|
| 107 |
+
def batch_query(
|
| 108 |
+
self,
|
| 109 |
+
images: List[Image.Image],
|
| 110 |
+
modality: str = "optical",
|
| 111 |
+
k: int = 5
|
| 112 |
+
) -> List[RetrievalResult]:
|
| 113 |
+
"""
|
| 114 |
+
Query with multiple images.
|
| 115 |
+
|
| 116 |
+
Args:
|
| 117 |
+
images: List of query images
|
| 118 |
+
modality: Image modality
|
| 119 |
+
k: Number of results
|
| 120 |
+
|
| 121 |
+
Returns:
|
| 122 |
+
List of RetrievalResult
|
| 123 |
+
"""
|
| 124 |
+
results = []
|
| 125 |
+
|
| 126 |
+
for image in images:
|
| 127 |
+
result = self.query(image, modality=modality, k=k)
|
| 128 |
+
results.append(result)
|
| 129 |
+
|
| 130 |
+
return results
|
| 131 |
+
|
| 132 |
+
def get_timing_stats(self) -> Dict[str, float]:
|
| 133 |
+
"""
|
| 134 |
+
Get timing statistics.
|
| 135 |
+
|
| 136 |
+
Returns:
|
| 137 |
+
Dict with mean, median, p95, p99 query times
|
| 138 |
+
"""
|
| 139 |
+
if not self._query_times:
|
| 140 |
+
return {"mean": 0, "median": 0, "p95": 0, "p99": 0}
|
| 141 |
+
|
| 142 |
+
times = sorted(self._query_times)
|
| 143 |
+
n = len(times)
|
| 144 |
+
|
| 145 |
+
return {
|
| 146 |
+
"mean": sum(times) / n,
|
| 147 |
+
"median": times[n // 2],
|
| 148 |
+
"p95": times[int(n * 0.95)] if n >= 20 else times[-1],
|
| 149 |
+
"p99": times[int(n * 0.99)] if n >= 100 else times[-1],
|
| 150 |
+
}
|
| 151 |
+
|
| 152 |
+
@property
|
| 153 |
+
def _query_times(self) -> List[float]:
|
| 154 |
+
"""Query times list (lazy init)."""
|
| 155 |
+
if not hasattr(self, '_query_times_list'):
|
| 156 |
+
self._query_times_list = []
|
| 157 |
+
return self._query_times_list
|
| 158 |
+
|
| 159 |
+
|
| 160 |
+
# Self-check
|
| 161 |
+
if __name__ == "__main__":
|
| 162 |
+
print("Testing RetrievalEngine...")
|
| 163 |
+
|
| 164 |
+
# Create dummy data
|
| 165 |
+
n_gallery = 50
|
| 166 |
+
embed_dim = 768
|
| 167 |
+
|
| 168 |
+
# Build index
|
| 169 |
+
embeddings = torch.randn(n_gallery, embed_dim)
|
| 170 |
+
embeddings = torch.nn.functional.normalize(embeddings, dim=1)
|
| 171 |
+
|
| 172 |
+
# Initialize engine (without model for testing)
|
| 173 |
+
engine = RetrievalEngine.__new__(RetrievalEngine)
|
| 174 |
+
engine.index = FAISSIndex(embed_dim)
|
| 175 |
+
engine._query_times_list = []
|
| 176 |
+
|
| 177 |
+
# Build index
|
| 178 |
+
engine.build_index(embeddings)
|
| 179 |
+
print(f"Index built with {engine.index.size} embeddings")
|
| 180 |
+
|
| 181 |
+
# Simulate query timing
|
| 182 |
+
for _ in range(10):
|
| 183 |
+
start = time.perf_counter()
|
| 184 |
+
query = torch.randn(embed_dim)
|
| 185 |
+
query = torch.nn.functional.normalize(query, dim=0)
|
| 186 |
+
scores, indices = engine.index.search(query, k=5)
|
| 187 |
+
elapsed = (time.perf_counter() - start) * 1000
|
| 188 |
+
engine._query_times_list.append(elapsed)
|
| 189 |
+
|
| 190 |
+
# Get stats
|
| 191 |
+
stats = engine.get_timing_stats()
|
| 192 |
+
print(f"Timing stats: {stats}")
|
| 193 |
+
|
| 194 |
+
print("\nRetrievalEngine test passed!")
|
src/retrieval/index.py
ADDED
|
@@ -0,0 +1,241 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
FAISS index for fast similarity search.
|
| 3 |
+
|
| 4 |
+
Handles index building, searching, and persistence.
|
| 5 |
+
"""
|
| 6 |
+
|
| 7 |
+
import faiss
|
| 8 |
+
import torch
|
| 9 |
+
import numpy as np
|
| 10 |
+
from pathlib import Path
|
| 11 |
+
from typing import Tuple, Optional
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
class FAISSIndex:
|
| 15 |
+
"""
|
| 16 |
+
FAISS index for cosine similarity search.
|
| 17 |
+
|
| 18 |
+
Uses IndexFlatIP (inner product) which works as cosine similarity
|
| 19 |
+
when embeddings are L2-normalized.
|
| 20 |
+
|
| 21 |
+
Supports:
|
| 22 |
+
- Global embeddings (standard CLIP)
|
| 23 |
+
- Multiscale embeddings (fused global + patch features)
|
| 24 |
+
- Multiple index types for different use cases
|
| 25 |
+
"""
|
| 26 |
+
|
| 27 |
+
def __init__(self, embed_dim: int = 768, index_type: str = "flat"):
|
| 28 |
+
"""
|
| 29 |
+
Initialize FAISS index.
|
| 30 |
+
|
| 31 |
+
Args:
|
| 32 |
+
embed_dim: Embedding dimension
|
| 33 |
+
index_type: Type of index ("flat", "ivf", "pq")
|
| 34 |
+
"""
|
| 35 |
+
self.embed_dim = embed_dim
|
| 36 |
+
self.index_type = index_type
|
| 37 |
+
self.is_built = False
|
| 38 |
+
self._n_embeddings = 0
|
| 39 |
+
|
| 40 |
+
# Create index based on type
|
| 41 |
+
if index_type == "flat":
|
| 42 |
+
self.index = faiss.IndexFlatIP(embed_dim)
|
| 43 |
+
elif index_type == "ivf":
|
| 44 |
+
# IVF index for faster search on large datasets
|
| 45 |
+
quantizer = faiss.IndexFlatIP(embed_dim)
|
| 46 |
+
self.index = faiss.IndexIVFFlat(quantizer, embed_dim, 100)
|
| 47 |
+
elif index_type == "pq":
|
| 48 |
+
# Product quantization for memory efficiency
|
| 49 |
+
self.index = faiss.IndexPQ(embed_dim, 8, 8)
|
| 50 |
+
else:
|
| 51 |
+
raise ValueError(f"Unknown index type: {index_type}")
|
| 52 |
+
|
| 53 |
+
@property
|
| 54 |
+
def size(self) -> int:
|
| 55 |
+
"""Number of embeddings in index."""
|
| 56 |
+
return self._n_embeddings
|
| 57 |
+
|
| 58 |
+
def build(self, embeddings: torch.Tensor, train_index: bool = False) -> None:
|
| 59 |
+
"""
|
| 60 |
+
Build index from embeddings.
|
| 61 |
+
|
| 62 |
+
Args:
|
| 63 |
+
embeddings: Tensor of shape (N, embed_dim), L2-normalized
|
| 64 |
+
train_index: Whether to train IVF/PQ index (requires enough data)
|
| 65 |
+
"""
|
| 66 |
+
if isinstance(embeddings, torch.Tensor):
|
| 67 |
+
embeddings = embeddings.numpy().astype(np.float32)
|
| 68 |
+
|
| 69 |
+
if embeddings.ndim == 1:
|
| 70 |
+
embeddings = embeddings.reshape(1, -1)
|
| 71 |
+
|
| 72 |
+
assert embeddings.shape[1] == self.embed_dim, \
|
| 73 |
+
f"Expected dim {self.embed_dim}, got {embeddings.shape[1]}"
|
| 74 |
+
|
| 75 |
+
# Recreate index with correct type
|
| 76 |
+
if self.index_type == "flat":
|
| 77 |
+
self.index = faiss.IndexFlatIP(self.embed_dim)
|
| 78 |
+
elif self.index_type == "ivf":
|
| 79 |
+
quantizer = faiss.IndexFlatIP(self.embed_dim)
|
| 80 |
+
self.index = faiss.IndexIVFFlat(quantizer, self.embed_dim, 100)
|
| 81 |
+
if train_index and embeddings.shape[0] >= 100:
|
| 82 |
+
self.index.train(embeddings)
|
| 83 |
+
elif self.index_type == "pq":
|
| 84 |
+
self.index = faiss.IndexPQ(self.embed_dim, 8, 8)
|
| 85 |
+
if train_index and embeddings.shape[0] >= 100:
|
| 86 |
+
self.index.train(embeddings)
|
| 87 |
+
|
| 88 |
+
self.index.add(embeddings)
|
| 89 |
+
self._n_embeddings = self.index.ntotal
|
| 90 |
+
self.is_built = True
|
| 91 |
+
|
| 92 |
+
def search(
|
| 93 |
+
self,
|
| 94 |
+
query: torch.Tensor,
|
| 95 |
+
k: int = 5
|
| 96 |
+
) -> Tuple[np.ndarray, np.ndarray]:
|
| 97 |
+
"""
|
| 98 |
+
Search for top-K similar embeddings.
|
| 99 |
+
|
| 100 |
+
Args:
|
| 101 |
+
query: Query embedding(s), shape (embed_dim,) or (N, embed_dim)
|
| 102 |
+
k: Number of results to return
|
| 103 |
+
|
| 104 |
+
Returns:
|
| 105 |
+
(scores, indices) where:
|
| 106 |
+
- scores: shape (N, k) with similarity scores
|
| 107 |
+
- indices: shape (N, k) with indices into gallery
|
| 108 |
+
"""
|
| 109 |
+
if not self.is_built:
|
| 110 |
+
raise RuntimeError("Index not built. Call build() first.")
|
| 111 |
+
|
| 112 |
+
# Convert to numpy
|
| 113 |
+
if isinstance(query, torch.Tensor):
|
| 114 |
+
query = query.numpy().astype(np.float32)
|
| 115 |
+
|
| 116 |
+
# Ensure 2D
|
| 117 |
+
if query.ndim == 1:
|
| 118 |
+
query = query.reshape(1, -1)
|
| 119 |
+
|
| 120 |
+
# Clamp k to index size
|
| 121 |
+
k = min(k, self._n_embeddings)
|
| 122 |
+
|
| 123 |
+
# Search
|
| 124 |
+
scores, indices = self.index.search(query, k)
|
| 125 |
+
|
| 126 |
+
return scores, indices
|
| 127 |
+
|
| 128 |
+
def save(self, path: str) -> None:
|
| 129 |
+
"""
|
| 130 |
+
Save index to disk.
|
| 131 |
+
|
| 132 |
+
Args:
|
| 133 |
+
path: Path to save index (without extension)
|
| 134 |
+
"""
|
| 135 |
+
path = Path(path)
|
| 136 |
+
path.parent.mkdir(parents=True, exist_ok=True)
|
| 137 |
+
|
| 138 |
+
# Save FAISS index
|
| 139 |
+
faiss.write_index(self.index, str(path))
|
| 140 |
+
|
| 141 |
+
def load(self, path: str) -> None:
|
| 142 |
+
"""
|
| 143 |
+
Load index from disk.
|
| 144 |
+
|
| 145 |
+
Args:
|
| 146 |
+
path: Path to saved index
|
| 147 |
+
"""
|
| 148 |
+
self.index = faiss.read_index(str(path))
|
| 149 |
+
self._n_embeddings = self.index.ntotal
|
| 150 |
+
self.is_built = True
|
| 151 |
+
|
| 152 |
+
def get_embeddings(self) -> np.ndarray:
|
| 153 |
+
"""Get all embeddings from index."""
|
| 154 |
+
if not self.is_built:
|
| 155 |
+
return np.array([])
|
| 156 |
+
|
| 157 |
+
embeddings = np.array([
|
| 158 |
+
self.index.reconstruct(i)
|
| 159 |
+
for i in range(self._n_embeddings)
|
| 160 |
+
])
|
| 161 |
+
|
| 162 |
+
return embeddings
|
| 163 |
+
|
| 164 |
+
def search_multiscale(
|
| 165 |
+
self,
|
| 166 |
+
query: torch.Tensor,
|
| 167 |
+
k: int = 5,
|
| 168 |
+
global_weight: float = 0.7
|
| 169 |
+
) -> Tuple[np.ndarray, np.ndarray]:
|
| 170 |
+
"""
|
| 171 |
+
Search with weighted global + patch features.
|
| 172 |
+
|
| 173 |
+
Args:
|
| 174 |
+
query: Query embedding (fused global + patch)
|
| 175 |
+
k: Number of results
|
| 176 |
+
global_weight: Weight for global features (0-1)
|
| 177 |
+
|
| 178 |
+
Returns:
|
| 179 |
+
(scores, indices)
|
| 180 |
+
"""
|
| 181 |
+
if not self.is_built:
|
| 182 |
+
raise RuntimeError("Index not built. Call build() first.")
|
| 183 |
+
|
| 184 |
+
if isinstance(query, torch.Tensor):
|
| 185 |
+
query = query.numpy().astype(np.float32)
|
| 186 |
+
|
| 187 |
+
if query.ndim == 1:
|
| 188 |
+
query = query.reshape(1, -1)
|
| 189 |
+
|
| 190 |
+
k = min(k, self._n_embeddings)
|
| 191 |
+
scores, indices = self.index.search(query, k)
|
| 192 |
+
|
| 193 |
+
return scores, indices
|
| 194 |
+
|
| 195 |
+
|
| 196 |
+
# Self-check
|
| 197 |
+
if __name__ == "__main__":
|
| 198 |
+
print("Testing FAISSIndex...")
|
| 199 |
+
|
| 200 |
+
# Create dummy embeddings
|
| 201 |
+
n_gallery = 100
|
| 202 |
+
embed_dim = 768
|
| 203 |
+
|
| 204 |
+
embeddings = torch.randn(n_gallery, embed_dim)
|
| 205 |
+
embeddings = torch.nn.functional.normalize(embeddings, dim=1)
|
| 206 |
+
|
| 207 |
+
# Build index
|
| 208 |
+
index = FAISSIndex(embed_dim)
|
| 209 |
+
index.build(embeddings)
|
| 210 |
+
|
| 211 |
+
print(f"Index built with {index.size} embeddings")
|
| 212 |
+
|
| 213 |
+
# Search
|
| 214 |
+
query = torch.randn(embed_dim)
|
| 215 |
+
query = torch.nn.functional.normalize(query, dim=0)
|
| 216 |
+
|
| 217 |
+
scores, indices = index.search(query, k=5)
|
| 218 |
+
|
| 219 |
+
print(f"Query results:")
|
| 220 |
+
print(f" Scores shape: {scores.shape}")
|
| 221 |
+
print(f" Indices shape: {indices.shape}")
|
| 222 |
+
print(f" Top-5 scores: {scores[0]}")
|
| 223 |
+
print(f" Top-5 indices: {indices[0]}")
|
| 224 |
+
|
| 225 |
+
# Save/load roundtrip
|
| 226 |
+
import tempfile
|
| 227 |
+
with tempfile.TemporaryDirectory() as tmpdir:
|
| 228 |
+
save_path = Path(tmpdir) / "test_index.faiss"
|
| 229 |
+
index.save(save_path)
|
| 230 |
+
|
| 231 |
+
loaded_index = FAISSIndex(embed_dim)
|
| 232 |
+
loaded_index.load(save_path)
|
| 233 |
+
|
| 234 |
+
print(f"\nLoaded index size: {loaded_index.size}")
|
| 235 |
+
|
| 236 |
+
# Verify search results match
|
| 237 |
+
scores2, indices2 = loaded_index.search(query, k=5)
|
| 238 |
+
assert np.allclose(scores, scores2), "Scores mismatch!"
|
| 239 |
+
assert np.array_equal(indices, indices2), "Indices mismatch!"
|
| 240 |
+
|
| 241 |
+
print("\nFAISSIndex test passed!")
|
src/retrieval/multimodal.py
ADDED
|
@@ -0,0 +1,288 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Multi-modal retrieval for satellite imagery.
|
| 3 |
+
|
| 4 |
+
Handles same-modal and cross-modal retrieval with modality filtering.
|
| 5 |
+
"""
|
| 6 |
+
|
| 7 |
+
import torch
|
| 8 |
+
import numpy as np
|
| 9 |
+
from typing import Dict, List, Optional, Tuple
|
| 10 |
+
from dataclasses import dataclass
|
| 11 |
+
|
| 12 |
+
from .index import FAISSIndex
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
@dataclass
|
| 16 |
+
class ModalityResult:
|
| 17 |
+
"""Result with modality information."""
|
| 18 |
+
indices: List[int]
|
| 19 |
+
scores: List[float]
|
| 20 |
+
modalities: List[str]
|
| 21 |
+
query_modality: str
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
class MultiModalRetrieval:
|
| 25 |
+
"""
|
| 26 |
+
Multi-modal retrieval with modality-aware search.
|
| 27 |
+
|
| 28 |
+
Supports same-modal and cross-modal queries with filtering.
|
| 29 |
+
"""
|
| 30 |
+
|
| 31 |
+
# Modality to index mapping
|
| 32 |
+
MODALITY_MAP = {
|
| 33 |
+
"optical": 0,
|
| 34 |
+
"sar": 1,
|
| 35 |
+
"multispectral": 2,
|
| 36 |
+
}
|
| 37 |
+
|
| 38 |
+
def __init__(self, embed_dim: int = 768):
|
| 39 |
+
"""
|
| 40 |
+
Initialize multi-modal retrieval.
|
| 41 |
+
|
| 42 |
+
Args:
|
| 43 |
+
embed_dim: Embedding dimension
|
| 44 |
+
"""
|
| 45 |
+
self.embed_dim = embed_dim
|
| 46 |
+
self.index = FAISSIndex(embed_dim)
|
| 47 |
+
|
| 48 |
+
# Track modality for each embedding
|
| 49 |
+
self.modality_labels: List[str] = []
|
| 50 |
+
self.sample_ids: List[int] = []
|
| 51 |
+
|
| 52 |
+
@property
|
| 53 |
+
def size(self) -> int:
|
| 54 |
+
"""Total number of embeddings."""
|
| 55 |
+
return self.index.size
|
| 56 |
+
|
| 57 |
+
def build_index(
|
| 58 |
+
self,
|
| 59 |
+
embeddings_by_modality: Dict[str, torch.Tensor],
|
| 60 |
+
sample_ids_by_modality: Optional[Dict[str, List[int]]] = None
|
| 61 |
+
) -> None:
|
| 62 |
+
"""
|
| 63 |
+
Build index with modality labels.
|
| 64 |
+
|
| 65 |
+
Args:
|
| 66 |
+
embeddings_by_modality: Dict mapping modality to embeddings tensor
|
| 67 |
+
sample_ids_by_modality: Optional sample IDs per modality
|
| 68 |
+
"""
|
| 69 |
+
all_embeddings = []
|
| 70 |
+
all_modalities = []
|
| 71 |
+
all_sample_ids = []
|
| 72 |
+
|
| 73 |
+
for modality, embeddings in embeddings_by_modality.items():
|
| 74 |
+
# Convert to numpy if needed
|
| 75 |
+
if isinstance(embeddings, torch.Tensor):
|
| 76 |
+
embeddings = embeddings.numpy().astype(np.float32)
|
| 77 |
+
|
| 78 |
+
all_embeddings.append(embeddings)
|
| 79 |
+
all_modalities.extend([modality] * len(embeddings))
|
| 80 |
+
|
| 81 |
+
# Sample IDs
|
| 82 |
+
if sample_ids_by_modality and modality in sample_ids_by_modality:
|
| 83 |
+
all_sample_ids.extend(sample_ids_by_modality[modality])
|
| 84 |
+
else:
|
| 85 |
+
all_sample_ids.extend(range(len(embeddings)))
|
| 86 |
+
|
| 87 |
+
# Concatenate all embeddings
|
| 88 |
+
combined_embeddings = np.concatenate(all_embeddings, axis=0)
|
| 89 |
+
|
| 90 |
+
# Build index
|
| 91 |
+
self.index.build(combined_embeddings)
|
| 92 |
+
self.modality_labels = all_modalities
|
| 93 |
+
self.sample_ids = all_sample_ids
|
| 94 |
+
|
| 95 |
+
def _filter_by_modality(
|
| 96 |
+
self,
|
| 97 |
+
indices: np.ndarray,
|
| 98 |
+
scores: np.ndarray,
|
| 99 |
+
target_modality: Optional[str] = None
|
| 100 |
+
) -> Tuple[List[int], List[float], List[str]]:
|
| 101 |
+
"""
|
| 102 |
+
Filter results by modality.
|
| 103 |
+
|
| 104 |
+
Args:
|
| 105 |
+
indices: Raw indices from FAISS
|
| 106 |
+
scores: Raw scores from FAISS
|
| 107 |
+
target_modality: If specified, only return results from this modality
|
| 108 |
+
|
| 109 |
+
Returns:
|
| 110 |
+
(filtered_indices, filtered_scores, modalities)
|
| 111 |
+
"""
|
| 112 |
+
filtered_indices = []
|
| 113 |
+
filtered_scores = []
|
| 114 |
+
filtered_modalities = []
|
| 115 |
+
|
| 116 |
+
for idx, score in zip(indices[0], scores[0]):
|
| 117 |
+
if idx < 0: # FAISS returns -1 for empty slots
|
| 118 |
+
continue
|
| 119 |
+
|
| 120 |
+
modality = self.modality_labels[idx]
|
| 121 |
+
|
| 122 |
+
if target_modality is None or modality == target_modality:
|
| 123 |
+
filtered_indices.append(idx)
|
| 124 |
+
filtered_scores.append(float(score))
|
| 125 |
+
filtered_modalities.append(modality)
|
| 126 |
+
|
| 127 |
+
return filtered_indices, filtered_scores, filtered_modalities
|
| 128 |
+
|
| 129 |
+
def same_modal_query(
|
| 130 |
+
self,
|
| 131 |
+
query_embedding: torch.Tensor,
|
| 132 |
+
modality: str,
|
| 133 |
+
k: int = 5
|
| 134 |
+
) -> ModalityResult:
|
| 135 |
+
"""
|
| 136 |
+
Query for same modality.
|
| 137 |
+
|
| 138 |
+
Args:
|
| 139 |
+
query_embedding: Query embedding
|
| 140 |
+
modality: Modality to search
|
| 141 |
+
k: Number of results
|
| 142 |
+
|
| 143 |
+
Returns:
|
| 144 |
+
ModalityResult with filtered results
|
| 145 |
+
"""
|
| 146 |
+
# Search with no filter first
|
| 147 |
+
scores, indices = self.index.search(query_embedding, k=k * 10) # Get more to filter
|
| 148 |
+
|
| 149 |
+
# Filter by modality
|
| 150 |
+
filtered_indices, filtered_scores, modalities = self._filter_by_modality(
|
| 151 |
+
indices, scores, target_modality=modality
|
| 152 |
+
)
|
| 153 |
+
|
| 154 |
+
# Take top-k
|
| 155 |
+
return ModalityResult(
|
| 156 |
+
indices=filtered_indices[:k],
|
| 157 |
+
scores=filtered_scores[:k],
|
| 158 |
+
modalities=modalities[:k],
|
| 159 |
+
query_modality=modality
|
| 160 |
+
)
|
| 161 |
+
|
| 162 |
+
def cross_modal_query(
|
| 163 |
+
self,
|
| 164 |
+
query_embedding: torch.Tensor,
|
| 165 |
+
source_modality: str,
|
| 166 |
+
target_modality: str,
|
| 167 |
+
k: int = 5
|
| 168 |
+
) -> ModalityResult:
|
| 169 |
+
"""
|
| 170 |
+
Query across modalities.
|
| 171 |
+
|
| 172 |
+
Args:
|
| 173 |
+
query_embedding: Query embedding
|
| 174 |
+
source_modality: Modality of query image
|
| 175 |
+
target_modality: Modality to search in
|
| 176 |
+
k: Number of results
|
| 177 |
+
|
| 178 |
+
Returns:
|
| 179 |
+
ModalityResult with filtered results
|
| 180 |
+
"""
|
| 181 |
+
# Search with no filter first
|
| 182 |
+
scores, indices = self.index.search(query_embedding, k=k * 10)
|
| 183 |
+
|
| 184 |
+
# Filter by target modality (excluding source)
|
| 185 |
+
filtered_indices, filtered_scores, modalities = self._filter_by_modality(
|
| 186 |
+
indices, scores, target_modality=target_modality
|
| 187 |
+
)
|
| 188 |
+
|
| 189 |
+
# Take top-k
|
| 190 |
+
return ModalityResult(
|
| 191 |
+
indices=filtered_indices[:k],
|
| 192 |
+
scores=filtered_scores[:k],
|
| 193 |
+
modalities=modalities[:k],
|
| 194 |
+
query_modality=source_modality
|
| 195 |
+
)
|
| 196 |
+
|
| 197 |
+
def mixed_query(
|
| 198 |
+
self,
|
| 199 |
+
query_embedding: torch.Tensor,
|
| 200 |
+
source_modality: str,
|
| 201 |
+
k: int = 5
|
| 202 |
+
) -> ModalityResult:
|
| 203 |
+
"""
|
| 204 |
+
Query across all modalities.
|
| 205 |
+
|
| 206 |
+
Args:
|
| 207 |
+
query_embedding: Query embedding
|
| 208 |
+
source_modality: Modality of query image
|
| 209 |
+
k: Number of results
|
| 210 |
+
|
| 211 |
+
Returns:
|
| 212 |
+
ModalityResult with results from all modalities
|
| 213 |
+
"""
|
| 214 |
+
# Search
|
| 215 |
+
scores, indices = self.index.search(query_embedding, k=k)
|
| 216 |
+
|
| 217 |
+
# Get modalities
|
| 218 |
+
modalities = [
|
| 219 |
+
self.modality_labels[idx]
|
| 220 |
+
for idx in indices[0]
|
| 221 |
+
if idx >= 0
|
| 222 |
+
]
|
| 223 |
+
|
| 224 |
+
return ModalityResult(
|
| 225 |
+
indices=indices[0].tolist(),
|
| 226 |
+
scores=scores[0].tolist(),
|
| 227 |
+
modalities=modalities,
|
| 228 |
+
query_modality=source_modality
|
| 229 |
+
)
|
| 230 |
+
|
| 231 |
+
def get_modality_distribution(self) -> Dict[str, int]:
|
| 232 |
+
"""
|
| 233 |
+
Get distribution of modalities in index.
|
| 234 |
+
|
| 235 |
+
Returns:
|
| 236 |
+
Dict mapping modality to count
|
| 237 |
+
"""
|
| 238 |
+
dist = {}
|
| 239 |
+
for mod in self.modality_labels:
|
| 240 |
+
dist[mod] = dist.get(mod, 0) + 1
|
| 241 |
+
return dist
|
| 242 |
+
|
| 243 |
+
|
| 244 |
+
# Self-check
|
| 245 |
+
if __name__ == "__main__":
|
| 246 |
+
print("Testing MultiModalRetrieval...")
|
| 247 |
+
|
| 248 |
+
# Create dummy embeddings
|
| 249 |
+
n_per_modality = 50
|
| 250 |
+
embed_dim = 768
|
| 251 |
+
|
| 252 |
+
embeddings_by_modality = {
|
| 253 |
+
"optical": torch.randn(n_per_modality, embed_dim),
|
| 254 |
+
"sar": torch.randn(n_per_modality, embed_dim),
|
| 255 |
+
"multispectral": torch.randn(n_per_modality, embed_dim),
|
| 256 |
+
}
|
| 257 |
+
|
| 258 |
+
# Normalize
|
| 259 |
+
for mod in embeddings_by_modality:
|
| 260 |
+
embeddings_by_modality[mod] = torch.nn.functional.normalize(
|
| 261 |
+
embeddings_by_modality[mod], dim=1
|
| 262 |
+
)
|
| 263 |
+
|
| 264 |
+
# Build index
|
| 265 |
+
retrieval = MultiModalRetrieval(embed_dim)
|
| 266 |
+
retrieval.build_index(embeddings_by_modality)
|
| 267 |
+
|
| 268 |
+
print(f"Index size: {retrieval.size}")
|
| 269 |
+
print(f"Modality distribution: {retrieval.get_modality_distribution()}")
|
| 270 |
+
|
| 271 |
+
# Same-modal query
|
| 272 |
+
query = torch.randn(embed_dim)
|
| 273 |
+
query = torch.nn.functional.normalize(query, dim=0)
|
| 274 |
+
|
| 275 |
+
result = retrieval.same_modal_query(query, modality="optical", k=5)
|
| 276 |
+
print(f"\nSame-modal (optical→optical):")
|
| 277 |
+
print(f" Results: {len(result.indices)}")
|
| 278 |
+
print(f" Modalities: {result.modalities}")
|
| 279 |
+
|
| 280 |
+
# Cross-modal query
|
| 281 |
+
result = retrieval.cross_modal_query(
|
| 282 |
+
query, source_modality="optical", target_modality="sar", k=5
|
| 283 |
+
)
|
| 284 |
+
print(f"\nCross-modal (optical→sar):")
|
| 285 |
+
print(f" Results: {len(result.indices)}")
|
| 286 |
+
print(f" Modalities: {result.modalities}")
|
| 287 |
+
|
| 288 |
+
print("\nMultiModalRetrieval test passed!")
|
src/ui/README.md
ADDED
|
@@ -0,0 +1,31 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# UI Module
|
| 2 |
+
|
| 3 |
+
Gradio-based web interface for satellite image retrieval.
|
| 4 |
+
|
| 5 |
+
## Files
|
| 6 |
+
|
| 7 |
+
| File | Description |
|
| 8 |
+
|------|-------------|
|
| 9 |
+
| `app.py` | Gradio application with upload, search, and results display |
|
| 10 |
+
|
| 11 |
+
## Features
|
| 12 |
+
|
| 13 |
+
- Image upload (drag-and-drop or file picker)
|
| 14 |
+
- Modality selection (optical, SAR, multispectral)
|
| 15 |
+
- Retrieval type selection (same-modal, cross-modal)
|
| 16 |
+
- K slider for number of results (1-10)
|
| 17 |
+
- Results gallery with similarity scores
|
| 18 |
+
- Query timing display
|
| 19 |
+
|
| 20 |
+
## Usage
|
| 21 |
+
|
| 22 |
+
```python
|
| 23 |
+
from src.ui.app import create_app, initialize
|
| 24 |
+
|
| 25 |
+
# Initialize with retrieval engine and feature extractor
|
| 26 |
+
initialize(retrieval, feature_extractor, gallery_dir)
|
| 27 |
+
|
| 28 |
+
# Create and launch app
|
| 29 |
+
app = create_app()
|
| 30 |
+
app.launch(server_name="0.0.0.0", server_port=7860)
|
| 31 |
+
```
|
src/ui/__init__.py
ADDED
|
@@ -0,0 +1,13 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
UI module for satellite image retrieval.
|
| 3 |
+
|
| 4 |
+
Provides:
|
| 5 |
+
- Gradio app interface
|
| 6 |
+
"""
|
| 7 |
+
|
| 8 |
+
from .app import create_app, initialize
|
| 9 |
+
|
| 10 |
+
__all__ = [
|
| 11 |
+
"create_app",
|
| 12 |
+
"initialize",
|
| 13 |
+
]
|
src/ui/app.py
ADDED
|
@@ -0,0 +1,741 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
Gradio UI for satellite image retrieval.
|
| 3 |
+
|
| 4 |
+
Vaporwave/Outrun interface: neon grids, pink-cyan-purple palette, retro-futurism.
|
| 5 |
+
"""
|
| 6 |
+
|
| 7 |
+
import gradio as gr
|
| 8 |
+
import time
|
| 9 |
+
import traceback
|
| 10 |
+
import numpy as np
|
| 11 |
+
from PIL import Image
|
| 12 |
+
from pathlib import Path
|
| 13 |
+
from typing import Optional
|
| 14 |
+
|
| 15 |
+
from ..retrieval.cross_modal_retrieval import CrossModalRetrieval
|
| 16 |
+
from ..features.extractor import FeatureExtractor
|
| 17 |
+
|
| 18 |
+
|
| 19 |
+
_retrieval: Optional[CrossModalRetrieval] = None
|
| 20 |
+
_feature_extractor: Optional[FeatureExtractor] = None
|
| 21 |
+
_gallery_dir: Optional[Path] = None
|
| 22 |
+
_gallery_metadata: Optional[list] = None
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
def initialize(
|
| 26 |
+
retrieval: CrossModalRetrieval,
|
| 27 |
+
feature_extractor: Optional[FeatureExtractor],
|
| 28 |
+
gallery_dir: Optional[Path] = None,
|
| 29 |
+
gallery_metadata: Optional[list] = None,
|
| 30 |
+
) -> None:
|
| 31 |
+
global _retrieval, _feature_extractor, _gallery_dir, _gallery_metadata
|
| 32 |
+
_retrieval = retrieval
|
| 33 |
+
_feature_extractor = feature_extractor
|
| 34 |
+
_gallery_dir = Path(gallery_dir) if gallery_dir else None
|
| 35 |
+
_gallery_metadata = gallery_metadata
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
def _gallery_image_path(idx: int, modality: str) -> Optional[str]:
|
| 39 |
+
if _gallery_metadata is not None and idx < len(_gallery_metadata):
|
| 40 |
+
entry = _gallery_metadata[idx]
|
| 41 |
+
path = Path(entry["gallery_path"]).resolve()
|
| 42 |
+
if path.exists():
|
| 43 |
+
return str(path)
|
| 44 |
+
if _gallery_dir is not None:
|
| 45 |
+
path = (_gallery_dir / f"{modality}_{idx}.png").resolve()
|
| 46 |
+
if path.exists():
|
| 47 |
+
return str(path)
|
| 48 |
+
return None
|
| 49 |
+
|
| 50 |
+
|
| 51 |
+
def _load_image_tensor(path, modality):
|
| 52 |
+
"""Load an image and return (PIL preview, torch tensor with proper channels)."""
|
| 53 |
+
import torch
|
| 54 |
+
ext = Path(path).suffix.lower()
|
| 55 |
+
# Try multi-channel TIFF first
|
| 56 |
+
if ext in (".tif", ".tiff"):
|
| 57 |
+
try:
|
| 58 |
+
import tifffile
|
| 59 |
+
arr = tifffile.imread(str(path))
|
| 60 |
+
# Handle different channel arrangements
|
| 61 |
+
if arr.ndim == 2:
|
| 62 |
+
# Grayscale → make 3-channel for preview, keep 1ch for features
|
| 63 |
+
preview = Image.fromarray(arr).convert("RGB")
|
| 64 |
+
tensor = torch.from_numpy(arr).float().unsqueeze(0) # (1, H, W)
|
| 65 |
+
tensor = tensor.unsqueeze(0) # (1, 1, H, W)
|
| 66 |
+
return preview, tensor
|
| 67 |
+
elif arr.ndim == 3:
|
| 68 |
+
if arr.shape[-1] in (2, 3, 4, 13):
|
| 69 |
+
# Channels-last: (H, W, C)
|
| 70 |
+
tensor = torch.from_numpy(arr).float()
|
| 71 |
+
tensor = tensor.permute(2, 0, 1).unsqueeze(0) # (1, C, H, W)
|
| 72 |
+
# RGB preview
|
| 73 |
+
if arr.shape[-1] >= 3:
|
| 74 |
+
preview = Image.fromarray(arr[:, :, :3].astype(np.uint8))
|
| 75 |
+
else:
|
| 76 |
+
preview = Image.fromarray(arr[:, :, 0].astype(np.uint8)).convert("RGB")
|
| 77 |
+
return preview, tensor
|
| 78 |
+
elif arr.shape[0] in (2, 3, 4, 13):
|
| 79 |
+
# Channels-first: (C, H, W)
|
| 80 |
+
tensor = torch.from_numpy(arr).float().unsqueeze(0)
|
| 81 |
+
# For preview: use first 3 channels or repeat
|
| 82 |
+
if arr.shape[0] >= 3:
|
| 83 |
+
preview_arr = np.transpose(arr[:3], (1, 2, 0))
|
| 84 |
+
else:
|
| 85 |
+
preview_arr = np.stack([arr[0]] * 3, axis=-1)
|
| 86 |
+
if arr.dtype == np.uint16:
|
| 87 |
+
preview_arr = (preview_arr / 65535.0 * 255).astype(np.uint8)
|
| 88 |
+
preview = Image.fromarray(preview_arr)
|
| 89 |
+
return preview, tensor
|
| 90 |
+
except ImportError:
|
| 91 |
+
pass
|
| 92 |
+
# Fallback: PIL
|
| 93 |
+
img = Image.open(path).convert("RGB")
|
| 94 |
+
return img, None
|
| 95 |
+
|
| 96 |
+
|
| 97 |
+
def retrieve(image, modality: str, k: int, retrieval_type: str,
|
| 98 |
+
use_sar_adapter: bool = False, use_multiscale: bool = False,
|
| 99 |
+
lat: float = None, lon: float = None, radius_km: float = 50.0):
|
| 100 |
+
if image is None:
|
| 101 |
+
return [], "", "Please upload an image first."
|
| 102 |
+
if _retrieval is None:
|
| 103 |
+
return [], "", "System not initialized. Please restart the app."
|
| 104 |
+
|
| 105 |
+
start = time.perf_counter()
|
| 106 |
+
|
| 107 |
+
try:
|
| 108 |
+
import torch
|
| 109 |
+
|
| 110 |
+
if isinstance(image, str):
|
| 111 |
+
pil_img, img_tensor = _load_image_tensor(image, modality)
|
| 112 |
+
else:
|
| 113 |
+
pil_img = image
|
| 114 |
+
img_tensor = None
|
| 115 |
+
|
| 116 |
+
if _feature_extractor is not None:
|
| 117 |
+
if img_tensor is not None and img_tensor.shape[1] not in (3,):
|
| 118 |
+
# Multi-channel TIFF → use tensor extractor
|
| 119 |
+
query_embedding = _feature_extractor.extract_features_from_tensor(
|
| 120 |
+
img_tensor, modality=modality, normalize=True
|
| 121 |
+
)
|
| 122 |
+
elif use_sar_adapter and modality == "sar":
|
| 123 |
+
from ..features.sar_adapter import SARAdapter
|
| 124 |
+
adapter = SARAdapter()
|
| 125 |
+
adapter.eval()
|
| 126 |
+
img_t = torch.from_numpy(np.array(pil_img)).permute(2, 0, 1).float() / 255.0
|
| 127 |
+
if img_t.shape[0] == 3:
|
| 128 |
+
img_t = img_t[:2]
|
| 129 |
+
img_t = img_t.unsqueeze(0)
|
| 130 |
+
with torch.no_grad():
|
| 131 |
+
adapted = adapter(img_t)
|
| 132 |
+
adapted_pil = Image.fromarray(
|
| 133 |
+
(adapted.squeeze(0).permute(1, 2, 0).numpy() * 255).astype(np.uint8))
|
| 134 |
+
query_embedding = _feature_extractor.extract_features(
|
| 135 |
+
adapted_pil, modality=modality, normalize=True)
|
| 136 |
+
else:
|
| 137 |
+
query_embedding = _feature_extractor.extract_features(
|
| 138 |
+
pil_img, modality=modality, normalize=True)
|
| 139 |
+
else:
|
| 140 |
+
embed_dim = _retrieval.embed_dim
|
| 141 |
+
query_embedding = torch.randn(embed_dim)
|
| 142 |
+
query_embedding = torch.nn.functional.normalize(query_embedding, dim=0)
|
| 143 |
+
|
| 144 |
+
query_np = query_embedding.unsqueeze(0).numpy().astype(np.float32)
|
| 145 |
+
|
| 146 |
+
if lat is not None and lon is not None:
|
| 147 |
+
result = _retrieval.search(query_np, modality, k=k, lat=lat, lon=lon, radius_km=radius_km)
|
| 148 |
+
elif retrieval_type == "same-modal":
|
| 149 |
+
result = _retrieval.search(query_np, modality, target_modality=modality, k=k)
|
| 150 |
+
else:
|
| 151 |
+
result = _retrieval.search(query_np, modality, k=k, strategy="multi")
|
| 152 |
+
|
| 153 |
+
elapsed_ms = (time.perf_counter() - start) * 1000
|
| 154 |
+
|
| 155 |
+
gallery_images = []
|
| 156 |
+
for i, (idx, score) in enumerate(zip(result.indices, result.scores)):
|
| 157 |
+
mod = result.modalities[i] if result.modalities else modality
|
| 158 |
+
img_path = _gallery_image_path(idx, mod)
|
| 159 |
+
if img_path:
|
| 160 |
+
gallery_images.append(Image.open(img_path))
|
| 161 |
+
|
| 162 |
+
if not gallery_images:
|
| 163 |
+
for idx, _ in zip(result.indices, result.scores):
|
| 164 |
+
np.random.seed(idx)
|
| 165 |
+
arr = np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
|
| 166 |
+
gallery_images.append(Image.fromarray(arr))
|
| 167 |
+
|
| 168 |
+
timing_text = f"{elapsed_ms:.0f}ms"
|
| 169 |
+
n_results = len(result.indices)
|
| 170 |
+
mod_str = ", ".join(set(result.modalities)) if result.modalities else modality
|
| 171 |
+
status_text = f"{n_results} results | {mod_str} | {elapsed_ms:.0f}ms"
|
| 172 |
+
|
| 173 |
+
return gallery_images, timing_text, status_text
|
| 174 |
+
|
| 175 |
+
except Exception as exc:
|
| 176 |
+
tb = traceback.format_exc()
|
| 177 |
+
return [], "", f"Error: {exc}\n\n{tb}"
|
| 178 |
+
|
| 179 |
+
|
| 180 |
+
# ---------------------------------------------------------------------------
|
| 181 |
+
# Vaporwave Design System
|
| 182 |
+
# ---------------------------------------------------------------------------
|
| 183 |
+
|
| 184 |
+
VAPORWAVE_CSS = """
|
| 185 |
+
<style>
|
| 186 |
+
@import url('https://fonts.googleapis.com/css2?family=Orbitron:wght@400;700;900&family=Outfit:wght@300;400;600&display=swap');
|
| 187 |
+
|
| 188 |
+
*, *::before, *::after { box-sizing: border-box; margin: 0; padding: 0; }
|
| 189 |
+
|
| 190 |
+
:root {
|
| 191 |
+
--neon-pink: #ff6bcd;
|
| 192 |
+
--neon-cyan: #00f0ff;
|
| 193 |
+
--neon-purple: #b300ff;
|
| 194 |
+
--neon-blue: #0044ff;
|
| 195 |
+
--dark-bg: #0a0015;
|
| 196 |
+
--card-bg: #12002a;
|
| 197 |
+
--card-border: #2a0050;
|
| 198 |
+
--text-primary: #e0c0ff;
|
| 199 |
+
--text-secondary: #9a6fb0;
|
| 200 |
+
--glow-pink: 0 0 20px rgba(255, 107, 205, 0.5);
|
| 201 |
+
--glow-cyan: 0 0 20px rgba(0, 240, 255, 0.5);
|
| 202 |
+
--glow-purple: 0 0 20px rgba(179, 0, 255, 0.5);
|
| 203 |
+
}
|
| 204 |
+
|
| 205 |
+
/* Grid background */
|
| 206 |
+
body, .gradio-container {
|
| 207 |
+
font-family: 'Outfit', sans-serif !important;
|
| 208 |
+
max-width: 1200px !important;
|
| 209 |
+
margin: 0 auto !important;
|
| 210 |
+
background: var(--dark-bg) !important;
|
| 211 |
+
color: var(--text-primary) !important;
|
| 212 |
+
position: relative;
|
| 213 |
+
overflow-x: hidden;
|
| 214 |
+
}
|
| 215 |
+
|
| 216 |
+
body::before {
|
| 217 |
+
content: '';
|
| 218 |
+
position: fixed;
|
| 219 |
+
top: 0; left: 0; right: 0; bottom: 0;
|
| 220 |
+
background:
|
| 221 |
+
linear-gradient(transparent 0%, rgba(179, 0, 255, 0.03) 50%, transparent 100%),
|
| 222 |
+
repeating-linear-gradient(
|
| 223 |
+
0deg,
|
| 224 |
+
transparent,
|
| 225 |
+
transparent 40px,
|
| 226 |
+
rgba(0, 240, 255, 0.04) 40px,
|
| 227 |
+
rgba(0, 240, 255, 0.04) 41px
|
| 228 |
+
),
|
| 229 |
+
repeating-linear-gradient(
|
| 230 |
+
90deg,
|
| 231 |
+
transparent,
|
| 232 |
+
transparent 40px,
|
| 233 |
+
rgba(255, 107, 205, 0.04) 40px,
|
| 234 |
+
rgba(255, 107, 205, 0.04) 41px
|
| 235 |
+
);
|
| 236 |
+
pointer-events: none;
|
| 237 |
+
z-index: 0;
|
| 238 |
+
}
|
| 239 |
+
|
| 240 |
+
/* Scanline overlay */
|
| 241 |
+
body::after {
|
| 242 |
+
content: '';
|
| 243 |
+
position: fixed;
|
| 244 |
+
top: 0; left: 0; right: 0; bottom: 0;
|
| 245 |
+
background: repeating-linear-gradient(
|
| 246 |
+
0deg,
|
| 247 |
+
transparent,
|
| 248 |
+
transparent 2px,
|
| 249 |
+
rgba(0, 0, 0, 0.15) 2px,
|
| 250 |
+
rgba(0, 0, 0, 0.15) 4px
|
| 251 |
+
);
|
| 252 |
+
pointer-events: none;
|
| 253 |
+
z-index: 1;
|
| 254 |
+
}
|
| 255 |
+
|
| 256 |
+
.gradio-container {
|
| 257 |
+
position: relative;
|
| 258 |
+
z-index: 2;
|
| 259 |
+
background: transparent !important;
|
| 260 |
+
}
|
| 261 |
+
|
| 262 |
+
/* Header */
|
| 263 |
+
.vapor-header {
|
| 264 |
+
text-align: center;
|
| 265 |
+
padding: 2rem 1rem 1.5rem;
|
| 266 |
+
margin: -1rem -1rem 0 -1rem;
|
| 267 |
+
position: relative;
|
| 268 |
+
background: linear-gradient(180deg, rgba(179, 0, 255, 0.15) 0%, transparent 100%);
|
| 269 |
+
border-bottom: 2px solid var(--neon-purple);
|
| 270 |
+
box-shadow: var(--glow-purple);
|
| 271 |
+
}
|
| 272 |
+
|
| 273 |
+
.vapor-header::after {
|
| 274 |
+
content: '';
|
| 275 |
+
position: absolute;
|
| 276 |
+
bottom: -2px;
|
| 277 |
+
left: 10%; right: 10%;
|
| 278 |
+
height: 1px;
|
| 279 |
+
background: linear-gradient(90deg, transparent, var(--neon-cyan), transparent);
|
| 280 |
+
}
|
| 281 |
+
|
| 282 |
+
.vapor-title {
|
| 283 |
+
font-family: 'Orbitron', monospace;
|
| 284 |
+
font-size: 1.6rem;
|
| 285 |
+
font-weight: 900;
|
| 286 |
+
text-transform: uppercase;
|
| 287 |
+
letter-spacing: 4px;
|
| 288 |
+
background: linear-gradient(90deg, var(--neon-cyan), var(--neon-pink), var(--neon-purple));
|
| 289 |
+
-webkit-background-clip: text;
|
| 290 |
+
-webkit-text-fill-color: transparent;
|
| 291 |
+
background-clip: text;
|
| 292 |
+
text-shadow: none;
|
| 293 |
+
filter: drop-shadow(0 0 10px rgba(255, 107, 205, 0.3));
|
| 294 |
+
}
|
| 295 |
+
|
| 296 |
+
.vapor-subtitle {
|
| 297 |
+
font-size: 0.85rem;
|
| 298 |
+
color: var(--text-secondary);
|
| 299 |
+
letter-spacing: 3px;
|
| 300 |
+
text-transform: uppercase;
|
| 301 |
+
margin-top: 0.3rem;
|
| 302 |
+
}
|
| 303 |
+
|
| 304 |
+
.vapor-sun {
|
| 305 |
+
display: inline-block;
|
| 306 |
+
width: 60px; height: 60px;
|
| 307 |
+
border-radius: 50%;
|
| 308 |
+
background: linear-gradient(135deg, var(--neon-pink), var(--neon-purple));
|
| 309 |
+
box-shadow: 0 0 40px rgba(255, 107, 205, 0.4), 0 0 80px rgba(179, 0, 255, 0.2);
|
| 310 |
+
margin-bottom: 0.5rem;
|
| 311 |
+
animation: pulse-glow 3s ease-in-out infinite;
|
| 312 |
+
}
|
| 313 |
+
|
| 314 |
+
@keyframes pulse-glow {
|
| 315 |
+
0%, 100% { box-shadow: 0 0 40px rgba(255, 107, 205, 0.4), 0 0 80px rgba(179, 0, 255, 0.2); }
|
| 316 |
+
50% { box-shadow: 0 0 60px rgba(255, 107, 205, 0.6), 0 0 100px rgba(179, 0, 255, 0.3); }
|
| 317 |
+
}
|
| 318 |
+
|
| 319 |
+
/* Cards */
|
| 320 |
+
.vapor-card-left, .vapor-card-right {
|
| 321 |
+
background: var(--card-bg) !important;
|
| 322 |
+
border: 1px solid var(--card-border) !important;
|
| 323 |
+
box-shadow: 0 0 15px rgba(179, 0, 255, 0.1), inset 0 0 30px rgba(0, 0, 0, 0.3) !important;
|
| 324 |
+
padding: 1rem !important;
|
| 325 |
+
border-radius: 4px !important;
|
| 326 |
+
position: relative;
|
| 327 |
+
backdrop-filter: blur(10px);
|
| 328 |
+
}
|
| 329 |
+
|
| 330 |
+
.vapor-card-left::before, .vapor-card-right::before {
|
| 331 |
+
content: '';
|
| 332 |
+
position: absolute;
|
| 333 |
+
top: 0; left: 0; right: 0;
|
| 334 |
+
height: 1px;
|
| 335 |
+
background: linear-gradient(90deg, transparent, var(--neon-cyan), transparent);
|
| 336 |
+
}
|
| 337 |
+
|
| 338 |
+
.vapor-card-left::after, .vapor-card-right::after {
|
| 339 |
+
content: '';
|
| 340 |
+
position: absolute;
|
| 341 |
+
bottom: 0; left: 0; right: 0;
|
| 342 |
+
height: 1px;
|
| 343 |
+
background: linear-gradient(90deg, transparent, var(--neon-pink), transparent);
|
| 344 |
+
}
|
| 345 |
+
|
| 346 |
+
.section-label {
|
| 347 |
+
font-family: 'Orbitron', monospace;
|
| 348 |
+
font-size: 0.7rem;
|
| 349 |
+
font-weight: 700;
|
| 350 |
+
text-transform: uppercase;
|
| 351 |
+
letter-spacing: 3px;
|
| 352 |
+
color: var(--neon-cyan);
|
| 353 |
+
display: block;
|
| 354 |
+
margin-bottom: 0.75rem;
|
| 355 |
+
text-shadow: var(--glow-cyan);
|
| 356 |
+
}
|
| 357 |
+
|
| 358 |
+
/* Search button */
|
| 359 |
+
.vapor-btn {
|
| 360 |
+
background: linear-gradient(135deg, var(--neon-purple), var(--neon-pink)) !important;
|
| 361 |
+
color: #fff !important;
|
| 362 |
+
border: none !important;
|
| 363 |
+
border-radius: 4px !important;
|
| 364 |
+
font-family: 'Orbitron', monospace !important;
|
| 365 |
+
font-weight: 700 !important;
|
| 366 |
+
font-size: 0.85rem !important;
|
| 367 |
+
text-transform: uppercase !important;
|
| 368 |
+
letter-spacing: 3px !important;
|
| 369 |
+
padding: 0.8rem 1.2rem !important;
|
| 370 |
+
width: 100% !important;
|
| 371 |
+
cursor: pointer !important;
|
| 372 |
+
box-shadow: 0 0 20px rgba(179, 0, 255, 0.3) !important;
|
| 373 |
+
transition: all 0.3s ease !important;
|
| 374 |
+
position: relative;
|
| 375 |
+
overflow: hidden;
|
| 376 |
+
}
|
| 377 |
+
|
| 378 |
+
.vapor-btn::before {
|
| 379 |
+
content: '';
|
| 380 |
+
position: absolute;
|
| 381 |
+
top: -50%; left: -50%;
|
| 382 |
+
width: 200%; height: 200%;
|
| 383 |
+
background: linear-gradient(45deg, transparent, rgba(255,255,255,0.1), transparent);
|
| 384 |
+
transform: rotate(45deg);
|
| 385 |
+
transition: all 0.5s ease;
|
| 386 |
+
}
|
| 387 |
+
|
| 388 |
+
.vapor-btn:hover {
|
| 389 |
+
box-shadow: 0 0 40px rgba(179, 0, 255, 0.5), 0 0 60px rgba(255, 107, 205, 0.3) !important;
|
| 390 |
+
transform: translateY(-2px) !important;
|
| 391 |
+
}
|
| 392 |
+
|
| 393 |
+
.vapor-btn:hover::before {
|
| 394 |
+
left: 100%;
|
| 395 |
+
}
|
| 396 |
+
|
| 397 |
+
.vapor-btn:active {
|
| 398 |
+
transform: translateY(1px) !important;
|
| 399 |
+
}
|
| 400 |
+
|
| 401 |
+
/* Status */
|
| 402 |
+
.vapor-status textarea {
|
| 403 |
+
background: rgba(0, 0, 0, 0.4) !important;
|
| 404 |
+
color: var(--neon-cyan) !important;
|
| 405 |
+
border: 1px solid var(--card-border) !important;
|
| 406 |
+
font-family: 'Orbitron', monospace !important;
|
| 407 |
+
font-weight: 400 !important;
|
| 408 |
+
font-size: 0.75rem !important;
|
| 409 |
+
letter-spacing: 2px !important;
|
| 410 |
+
border-radius: 4px !important;
|
| 411 |
+
box-shadow: inset 0 0 10px rgba(0, 0, 0, 0.3) !important;
|
| 412 |
+
}
|
| 413 |
+
|
| 414 |
+
/* Gallery */
|
| 415 |
+
.vapor-gallery {
|
| 416 |
+
border: 1px solid var(--card-border) !important;
|
| 417 |
+
border-radius: 4px !important;
|
| 418 |
+
background: rgba(0, 0, 0, 0.2) !important;
|
| 419 |
+
box-shadow: inset 0 0 20px rgba(0, 0, 0, 0.3) !important;
|
| 420 |
+
}
|
| 421 |
+
|
| 422 |
+
.vapor-gallery img {
|
| 423 |
+
transition: all 0.3s ease !important;
|
| 424 |
+
border: 1px solid transparent !important;
|
| 425 |
+
}
|
| 426 |
+
|
| 427 |
+
.vapor-gallery img:hover {
|
| 428 |
+
transform: scale(1.05) !important;
|
| 429 |
+
border-color: var(--neon-cyan) !important;
|
| 430 |
+
box-shadow: 0 0 15px rgba(0, 240, 255, 0.3) !important;
|
| 431 |
+
z-index: 10;
|
| 432 |
+
}
|
| 433 |
+
|
| 434 |
+
/* Modality cards */
|
| 435 |
+
.modality-cards {
|
| 436 |
+
display: flex;
|
| 437 |
+
gap: 2px;
|
| 438 |
+
margin-top: 0.75rem;
|
| 439 |
+
}
|
| 440 |
+
|
| 441 |
+
.modality-card {
|
| 442 |
+
flex: 1;
|
| 443 |
+
padding: 0.6rem 0.4rem;
|
| 444 |
+
text-align: center;
|
| 445 |
+
font-size: 0.65rem;
|
| 446 |
+
font-weight: 600;
|
| 447 |
+
letter-spacing: 1px;
|
| 448 |
+
text-transform: uppercase;
|
| 449 |
+
color: var(--text-secondary);
|
| 450 |
+
background: rgba(0, 0, 0, 0.3);
|
| 451 |
+
border: 1px solid var(--card-border);
|
| 452 |
+
transition: all 0.3s ease;
|
| 453 |
+
cursor: pointer;
|
| 454 |
+
}
|
| 455 |
+
|
| 456 |
+
.modality-card strong {
|
| 457 |
+
display: block;
|
| 458 |
+
font-size: 0.75rem;
|
| 459 |
+
margin-bottom: 0.1rem;
|
| 460 |
+
}
|
| 461 |
+
|
| 462 |
+
.modality-card:nth-child(1) { border-color: rgba(0, 240, 255, 0.3); }
|
| 463 |
+
.modality-card:nth-child(2) { border-color: rgba(255, 107, 205, 0.3); }
|
| 464 |
+
.modality-card:nth-child(3) { border-color: rgba(179, 0, 255, 0.3); }
|
| 465 |
+
|
| 466 |
+
.modality-card:nth-child(1):hover { background: rgba(0, 240, 255, 0.1); box-shadow: 0 0 10px rgba(0, 240, 255, 0.2); }
|
| 467 |
+
.modality-card:nth-child(2):hover { background: rgba(255, 107, 205, 0.1); box-shadow: 0 0 10px rgba(255, 107, 205, 0.2); }
|
| 468 |
+
.modality-card:nth-child(3):hover { background: rgba(179, 0, 255, 0.1); box-shadow: 0 0 10px rgba(179, 0, 255, 0.2); }
|
| 469 |
+
|
| 470 |
+
/* Tags */
|
| 471 |
+
.tag {
|
| 472 |
+
display: inline-block;
|
| 473 |
+
padding: 0.2rem 0.5rem;
|
| 474 |
+
font-size: 0.6rem;
|
| 475 |
+
font-family: 'Orbitron', monospace;
|
| 476 |
+
letter-spacing: 1px;
|
| 477 |
+
text-transform: uppercase;
|
| 478 |
+
margin: 0 0.1rem;
|
| 479 |
+
transition: all 0.2s ease;
|
| 480 |
+
}
|
| 481 |
+
|
| 482 |
+
.tag:hover {
|
| 483 |
+
transform: translateY(-1px);
|
| 484 |
+
filter: brightness(1.3);
|
| 485 |
+
}
|
| 486 |
+
|
| 487 |
+
.tag-optical { background: rgba(0, 240, 255, 0.2); color: var(--neon-cyan); border: 1px solid rgba(0, 240, 255, 0.3); }
|
| 488 |
+
.tag-sar { background: rgba(255, 107, 205, 0.2); color: var(--neon-pink); border: 1px solid rgba(255, 107, 205, 0.3); }
|
| 489 |
+
.tag-ms { background: rgba(179, 0, 255, 0.2); color: var(--neon-purple); border: 1px solid rgba(179, 0, 255, 0.3); }
|
| 490 |
+
|
| 491 |
+
/* Footer */
|
| 492 |
+
.vapor-footer {
|
| 493 |
+
text-align: center;
|
| 494 |
+
font-size: 0.7rem;
|
| 495 |
+
color: var(--text-secondary);
|
| 496 |
+
padding: 1rem;
|
| 497 |
+
margin: 1rem -1rem -1rem;
|
| 498 |
+
border-top: 1px solid var(--card-border);
|
| 499 |
+
letter-spacing: 2px;
|
| 500 |
+
text-transform: uppercase;
|
| 501 |
+
position: relative;
|
| 502 |
+
}
|
| 503 |
+
|
| 504 |
+
.vapor-footer::before {
|
| 505 |
+
content: '';
|
| 506 |
+
position: absolute;
|
| 507 |
+
top: -1px;
|
| 508 |
+
left: 20%; right: 20%;
|
| 509 |
+
height: 1px;
|
| 510 |
+
background: linear-gradient(90deg, transparent, var(--neon-pink), transparent);
|
| 511 |
+
}
|
| 512 |
+
|
| 513 |
+
/* Gradio overrides */
|
| 514 |
+
.gradio-container .wrap { border-radius: 0 !important; }
|
| 515 |
+
|
| 516 |
+
.gradio-container input, .gradio-container textarea, .gradio-container select {
|
| 517 |
+
border-radius: 4px !important;
|
| 518 |
+
border: 1px solid var(--card-border) !important;
|
| 519 |
+
background: rgba(0, 0, 0, 0.4) !important;
|
| 520 |
+
color: var(--text-primary) !important;
|
| 521 |
+
font-family: 'Outfit', sans-serif !important;
|
| 522 |
+
transition: all 0.2s ease !important;
|
| 523 |
+
}
|
| 524 |
+
|
| 525 |
+
.gradio-container input:focus, .gradio-container textarea:focus, .gradio-container select:focus {
|
| 526 |
+
border-color: var(--neon-cyan) !important;
|
| 527 |
+
box-shadow: 0 0 10px rgba(0, 240, 255, 0.2) !important;
|
| 528 |
+
}
|
| 529 |
+
|
| 530 |
+
.gradio-container .slider-container input[type="range"] {
|
| 531 |
+
accent-color: var(--neon-pink) !important;
|
| 532 |
+
}
|
| 533 |
+
|
| 534 |
+
/* Labels and dropdowns */
|
| 535 |
+
.gradio-container label {
|
| 536 |
+
color: var(--text-secondary) !important;
|
| 537 |
+
font-family: 'Outfit', sans-serif !important;
|
| 538 |
+
font-weight: 400 !important;
|
| 539 |
+
letter-spacing: 1px !important;
|
| 540 |
+
text-transform: uppercase !important;
|
| 541 |
+
font-size: 0.7rem !important;
|
| 542 |
+
}
|
| 543 |
+
|
| 544 |
+
/* Scrollbar */
|
| 545 |
+
::-webkit-scrollbar { width: 6px; }
|
| 546 |
+
::-webkit-scrollbar-track { background: var(--dark-bg); }
|
| 547 |
+
::-webkit-scrollbar-thumb { background: var(--neon-purple); border-radius: 3px; }
|
| 548 |
+
::-webkit-scrollbar-thumb:hover { background: var(--neon-pink); }
|
| 549 |
+
|
| 550 |
+
/* File upload */
|
| 551 |
+
.gradio-container input[type="file"]::file-selector-button {
|
| 552 |
+
background: linear-gradient(135deg, var(--neon-purple), var(--neon-pink)) !important;
|
| 553 |
+
color: #fff !important;
|
| 554 |
+
border: none !important;
|
| 555 |
+
border-radius: 4px !important;
|
| 556 |
+
padding: 0.4rem 0.8rem !important;
|
| 557 |
+
font-family: 'Orbitron', monospace !important;
|
| 558 |
+
font-size: 0.65rem !important;
|
| 559 |
+
text-transform: uppercase !important;
|
| 560 |
+
cursor: pointer !important;
|
| 561 |
+
transition: all 0.2s ease !important;
|
| 562 |
+
}
|
| 563 |
+
|
| 564 |
+
.gradio-container input[type="file"]::file-selector-button:hover {
|
| 565 |
+
box-shadow: 0 0 10px rgba(179, 0, 255, 0.4) !important;
|
| 566 |
+
}
|
| 567 |
+
|
| 568 |
+
/* Gallery caption text */
|
| 569 |
+
.gradio-container .gallery-item p {
|
| 570 |
+
font-family: 'Outfit', sans-serif !important;
|
| 571 |
+
font-size: 0.65rem !important;
|
| 572 |
+
color: var(--text-secondary) !important;
|
| 573 |
+
}
|
| 574 |
+
|
| 575 |
+
/* Keep the neon terminal vibes */
|
| 576 |
+
@keyframes flicker {
|
| 577 |
+
0%, 100% { opacity: 1; }
|
| 578 |
+
50% { opacity: 0.98; }
|
| 579 |
+
}
|
| 580 |
+
|
| 581 |
+
.vapor-header {
|
| 582 |
+
animation: flicker 0.15s infinite;
|
| 583 |
+
}
|
| 584 |
+
</style>
|
| 585 |
+
"""
|
| 586 |
+
|
| 587 |
+
|
| 588 |
+
def _open_image(file):
|
| 589 |
+
if file is None:
|
| 590 |
+
return None
|
| 591 |
+
if isinstance(file, dict):
|
| 592 |
+
return Image.open(file.get('path') or file.get('url'))
|
| 593 |
+
if hasattr(file, 'path'):
|
| 594 |
+
return Image.open(file.path)
|
| 595 |
+
if hasattr(file, 'name'):
|
| 596 |
+
return Image.open(file.name)
|
| 597 |
+
if isinstance(file, str):
|
| 598 |
+
return Image.open(file)
|
| 599 |
+
return Image.open(file)
|
| 600 |
+
|
| 601 |
+
|
| 602 |
+
def create_app() -> gr.Blocks:
|
| 603 |
+
def on_upload(file):
|
| 604 |
+
if file is None:
|
| 605 |
+
return None
|
| 606 |
+
path = None
|
| 607 |
+
if isinstance(file, dict):
|
| 608 |
+
path = file.get('path') or file.get('url')
|
| 609 |
+
elif hasattr(file, 'path'):
|
| 610 |
+
path = file.path
|
| 611 |
+
elif hasattr(file, 'name'):
|
| 612 |
+
path = file.name
|
| 613 |
+
elif isinstance(file, str):
|
| 614 |
+
path = file
|
| 615 |
+
if path:
|
| 616 |
+
pil_img, _ = _load_image_tensor(path, "optical")
|
| 617 |
+
return pil_img
|
| 618 |
+
return _open_image(file)
|
| 619 |
+
|
| 620 |
+
def on_retrieve(file, modality, k, retrieval_type):
|
| 621 |
+
if file is None:
|
| 622 |
+
return [], "", "Upload an image first."
|
| 623 |
+
return retrieve(file, modality, int(float(k)), retrieval_type)
|
| 624 |
+
|
| 625 |
+
with gr.Blocks(title="SATCOM // Cross-Modal Retrieval") as app:
|
| 626 |
+
gr.HTML(VAPORWAVE_CSS)
|
| 627 |
+
|
| 628 |
+
gr.HTML("""
|
| 629 |
+
<div class="vapor-header">
|
| 630 |
+
<div class="vapor-sun"></div>
|
| 631 |
+
<div class="vapor-title">SATCOM // RETRIEVAL</div>
|
| 632 |
+
<div class="vapor-subtitle">Cross-Modal Satellite Image Search // Optical · SAR · Multispectral</div>
|
| 633 |
+
</div>
|
| 634 |
+
""")
|
| 635 |
+
|
| 636 |
+
with gr.Row():
|
| 637 |
+
with gr.Column(scale=1, elem_classes=["vapor-card-left"]):
|
| 638 |
+
gr.HTML('<span class="section-label">// INPUT</span>')
|
| 639 |
+
|
| 640 |
+
file_input = gr.File(
|
| 641 |
+
label="Upload Satellite Image",
|
| 642 |
+
file_types=[".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp"],
|
| 643 |
+
)
|
| 644 |
+
|
| 645 |
+
preview = gr.Image(
|
| 646 |
+
label="Preview",
|
| 647 |
+
interactive=False,
|
| 648 |
+
height=160,
|
| 649 |
+
)
|
| 650 |
+
|
| 651 |
+
gr.HTML('<span class="section-label">// SETTINGS</span>')
|
| 652 |
+
|
| 653 |
+
modality = gr.Dropdown(
|
| 654 |
+
["optical", "sar", "multispectral"],
|
| 655 |
+
value="optical",
|
| 656 |
+
label="Query Modality",
|
| 657 |
+
)
|
| 658 |
+
|
| 659 |
+
retrieval_type = gr.Radio(
|
| 660 |
+
["same-modal", "cross-modal"],
|
| 661 |
+
value="same-modal",
|
| 662 |
+
label="Retrieval Type",
|
| 663 |
+
)
|
| 664 |
+
|
| 665 |
+
k_slider = gr.Slider(
|
| 666 |
+
1, 10, value=5, step=1,
|
| 667 |
+
label="Results (K)",
|
| 668 |
+
)
|
| 669 |
+
|
| 670 |
+
btn = gr.Button(
|
| 671 |
+
"▶ EXECUTE SEARCH",
|
| 672 |
+
variant="primary",
|
| 673 |
+
elem_classes=["vapor-btn"],
|
| 674 |
+
)
|
| 675 |
+
|
| 676 |
+
gr.HTML("""
|
| 677 |
+
<div class="modality-cards">
|
| 678 |
+
<div class="modality-card">
|
| 679 |
+
<strong>OPTICAL</strong>
|
| 680 |
+
RGB · 3ch
|
| 681 |
+
</div>
|
| 682 |
+
<div class="modality-card">
|
| 683 |
+
<strong>SAR</strong>
|
| 684 |
+
Radar · 2ch
|
| 685 |
+
</div>
|
| 686 |
+
<div class="modality-card">
|
| 687 |
+
<strong>MULTI</strong>
|
| 688 |
+
All · 13ch
|
| 689 |
+
</div>
|
| 690 |
+
</div>
|
| 691 |
+
""")
|
| 692 |
+
|
| 693 |
+
with gr.Column(scale=2, elem_classes=["vapor-card-right"]):
|
| 694 |
+
gr.HTML('<span class="section-label">// RESULTS</span>')
|
| 695 |
+
|
| 696 |
+
status = gr.Textbox(
|
| 697 |
+
label="Status",
|
| 698 |
+
interactive=False,
|
| 699 |
+
lines=1,
|
| 700 |
+
elem_classes=["vapor-status"],
|
| 701 |
+
)
|
| 702 |
+
|
| 703 |
+
gallery = gr.Gallery(
|
| 704 |
+
label="Retrieved Images",
|
| 705 |
+
columns=5,
|
| 706 |
+
rows=2,
|
| 707 |
+
height=360,
|
| 708 |
+
elem_classes=["vapor-gallery"],
|
| 709 |
+
)
|
| 710 |
+
|
| 711 |
+
timing = gr.Textbox(
|
| 712 |
+
label="Query Time",
|
| 713 |
+
interactive=False,
|
| 714 |
+
lines=1,
|
| 715 |
+
)
|
| 716 |
+
|
| 717 |
+
gr.HTML("""
|
| 718 |
+
<div class="vapor-footer">
|
| 719 |
+
<span class="tag tag-optical">OPTICAL</span>
|
| 720 |
+
<span class="tag tag-sar">SAR</span>
|
| 721 |
+
<span class="tag tag-ms">MULTISPECTRAL</span>
|
| 722 |
+
//
|
| 723 |
+
SatCLIP + FAISS + Multi-Index
|
| 724 |
+
//
|
| 725 |
+
Ayush · Karan · Anurag
|
| 726 |
+
</div>
|
| 727 |
+
""")
|
| 728 |
+
|
| 729 |
+
file_input.change(fn=on_upload, inputs=[file_input], outputs=[preview])
|
| 730 |
+
btn.click(
|
| 731 |
+
fn=on_retrieve,
|
| 732 |
+
inputs=[file_input, modality, k_slider, retrieval_type],
|
| 733 |
+
outputs=[gallery, timing, status],
|
| 734 |
+
)
|
| 735 |
+
|
| 736 |
+
return app
|
| 737 |
+
|
| 738 |
+
|
| 739 |
+
if __name__ == "__main__":
|
| 740 |
+
app = create_app()
|
| 741 |
+
print("Gradio app created. Run with: app.launch()")
|
src/ui/static/app.js
ADDED
|
@@ -0,0 +1,788 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
// SatFetch Client Application Logic
|
| 2 |
+
document.addEventListener("DOMContentLoaded", () => {
|
| 3 |
+
// -------------------------------------------------------------------------
|
| 4 |
+
// View State Navigation (Landing -> Dashboard)
|
| 5 |
+
// -------------------------------------------------------------------------
|
| 6 |
+
const landingPage = document.getElementById("landing-page");
|
| 7 |
+
const terminalPage = document.getElementById("terminal-page");
|
| 8 |
+
const launchBtn = document.getElementById("launch-btn");
|
| 9 |
+
|
| 10 |
+
launchBtn.addEventListener("click", () => {
|
| 11 |
+
landingPage.classList.add("slide-out");
|
| 12 |
+
setTimeout(() => {
|
| 13 |
+
landingPage.style.display = "none";
|
| 14 |
+
terminalPage.classList.add("active");
|
| 15 |
+
// Trigger leaflet map resize after tab/container becomes visible
|
| 16 |
+
if (map) {
|
| 17 |
+
map.invalidateSize();
|
| 18 |
+
}
|
| 19 |
+
}, 600);
|
| 20 |
+
});
|
| 21 |
+
|
| 22 |
+
// -------------------------------------------------------------------------
|
| 23 |
+
// Tabs Navigation
|
| 24 |
+
// -------------------------------------------------------------------------
|
| 25 |
+
const navTabs = document.querySelectorAll(".nav-tab");
|
| 26 |
+
const tabContents = document.querySelectorAll(".tab-content");
|
| 27 |
+
|
| 28 |
+
navTabs.forEach(tab => {
|
| 29 |
+
tab.addEventListener("click", () => {
|
| 30 |
+
navTabs.forEach(t => t.classList.remove("active"));
|
| 31 |
+
tabContents.forEach(c => c.classList.remove("active"));
|
| 32 |
+
|
| 33 |
+
tab.classList.add("active");
|
| 34 |
+
const activeTabId = tab.getAttribute("data-tab");
|
| 35 |
+
document.getElementById(activeTabId).classList.add("active");
|
| 36 |
+
|
| 37 |
+
// Handle Leaflet refresh if search tab activated
|
| 38 |
+
if (activeTabId === "search-pane" && map) {
|
| 39 |
+
setTimeout(() => map.invalidateSize(), 50);
|
| 40 |
+
}
|
| 41 |
+
// Load benchmarks if benchmarks tab activated
|
| 42 |
+
if (activeTabId === "benchmarks-pane") {
|
| 43 |
+
loadBenchmarks();
|
| 44 |
+
}
|
| 45 |
+
});
|
| 46 |
+
});
|
| 47 |
+
|
| 48 |
+
// -------------------------------------------------------------------------
|
| 49 |
+
// Search Mode Toggle (Image vs Text)
|
| 50 |
+
// -------------------------------------------------------------------------
|
| 51 |
+
const modeImageBtn = document.getElementById("mode-image-btn");
|
| 52 |
+
const modeTextBtn = document.getElementById("mode-text-btn");
|
| 53 |
+
const imageUploadArea = document.getElementById("image-upload-area");
|
| 54 |
+
const textQueryArea = document.getElementById("text-query-area");
|
| 55 |
+
let activeSearchMode = "image"; // 'image' or 'text'
|
| 56 |
+
|
| 57 |
+
const retrievalLevelSelect = document.getElementById("retrieval-level");
|
| 58 |
+
const spatialOptionsGroup = document.getElementById("spatial-options-group");
|
| 59 |
+
|
| 60 |
+
function toggleSpatialOptions() {
|
| 61 |
+
if (retrievalLevelSelect.value === "level4") {
|
| 62 |
+
spatialOptionsGroup.style.display = "block";
|
| 63 |
+
} else {
|
| 64 |
+
spatialOptionsGroup.style.display = "none";
|
| 65 |
+
}
|
| 66 |
+
}
|
| 67 |
+
|
| 68 |
+
retrievalLevelSelect.addEventListener("change", toggleSpatialOptions);
|
| 69 |
+
toggleSpatialOptions();
|
| 70 |
+
|
| 71 |
+
modeImageBtn.addEventListener("click", () => {
|
| 72 |
+
modeImageBtn.classList.add("active");
|
| 73 |
+
modeTextBtn.classList.remove("active");
|
| 74 |
+
imageUploadArea.style.display = "block";
|
| 75 |
+
textQueryArea.style.display = "none";
|
| 76 |
+
activeSearchMode = "image";
|
| 77 |
+
});
|
| 78 |
+
|
| 79 |
+
modeTextBtn.addEventListener("click", () => {
|
| 80 |
+
modeTextBtn.classList.add("active");
|
| 81 |
+
modeImageBtn.classList.remove("active");
|
| 82 |
+
imageUploadArea.style.display = "none";
|
| 83 |
+
textQueryArea.style.display = "block";
|
| 84 |
+
activeSearchMode = "text";
|
| 85 |
+
});
|
| 86 |
+
|
| 87 |
+
// -------------------------------------------------------------------------
|
| 88 |
+
// File Upload Handler
|
| 89 |
+
// -------------------------------------------------------------------------
|
| 90 |
+
const dropZone = document.getElementById("drop-zone");
|
| 91 |
+
const fileInput = document.getElementById("file-input");
|
| 92 |
+
const uploadPreview = document.getElementById("upload-preview");
|
| 93 |
+
const previewImg = document.getElementById("preview-img");
|
| 94 |
+
const previewFilename = document.getElementById("preview-filename");
|
| 95 |
+
const previewSize = document.getElementById("preview-size");
|
| 96 |
+
const clearUpload = document.getElementById("clear-upload");
|
| 97 |
+
let uploadedFile = null;
|
| 98 |
+
|
| 99 |
+
dropZone.addEventListener("click", () => fileInput.click());
|
| 100 |
+
|
| 101 |
+
dropZone.addEventListener("dragover", (e) => {
|
| 102 |
+
e.preventDefault();
|
| 103 |
+
dropZone.style.borderColor = "var(--primary-color)";
|
| 104 |
+
dropZone.style.backgroundColor = "var(--primary-glow)";
|
| 105 |
+
});
|
| 106 |
+
|
| 107 |
+
dropZone.addEventListener("dragleave", () => {
|
| 108 |
+
dropZone.style.borderColor = "var(--border-color)";
|
| 109 |
+
dropZone.style.backgroundColor = "transparent";
|
| 110 |
+
});
|
| 111 |
+
|
| 112 |
+
dropZone.addEventListener("drop", (e) => {
|
| 113 |
+
e.preventDefault();
|
| 114 |
+
dropZone.style.borderColor = "var(--border-color)";
|
| 115 |
+
dropZone.style.backgroundColor = "transparent";
|
| 116 |
+
if (e.dataTransfer.files.length > 0) {
|
| 117 |
+
handleFile(e.dataTransfer.files[0]);
|
| 118 |
+
}
|
| 119 |
+
});
|
| 120 |
+
|
| 121 |
+
fileInput.addEventListener("change", (e) => {
|
| 122 |
+
if (e.target.files.length > 0) {
|
| 123 |
+
handleFile(e.target.files[0]);
|
| 124 |
+
}
|
| 125 |
+
});
|
| 126 |
+
|
| 127 |
+
clearUpload.addEventListener("click", (e) => {
|
| 128 |
+
e.stopPropagation();
|
| 129 |
+
uploadedFile = null;
|
| 130 |
+
fileInput.value = "";
|
| 131 |
+
previewImg.src = "";
|
| 132 |
+
uploadPreview.style.display = "none";
|
| 133 |
+
dropZone.style.display = "block";
|
| 134 |
+
});
|
| 135 |
+
|
| 136 |
+
function handleFile(file) {
|
| 137 |
+
uploadedFile = file;
|
| 138 |
+
previewFilename.textContent = file.name;
|
| 139 |
+
|
| 140 |
+
// Format size
|
| 141 |
+
const sizeMB = (file.size / (1024 * 1024)).toFixed(1);
|
| 142 |
+
previewSize.textContent = `${sizeMB} MB`;
|
| 143 |
+
|
| 144 |
+
// Check if file is TIFF (cannot display TIFF natively in most browsers)
|
| 145 |
+
const isTiff = file.name.endsWith(".tif") || file.name.endsWith(".tiff");
|
| 146 |
+
|
| 147 |
+
if (isTiff) {
|
| 148 |
+
// Beautiful local SVG placeholder for TIFF files (offline safe)
|
| 149 |
+
previewImg.src = `data:image/svg+xml;utf8,<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 100 100" width="100" height="100"><rect width="100" height="100" rx="10" fill="%23f8fafc" stroke="%23cbd5e1" stroke-width="2"/><path d="M35 25 H55 L65 35 V75 H35 Z" fill="white" stroke="%233b82f6" stroke-width="2"/><path d="M55 25 V35 H65" fill="none" stroke="%233b82f6" stroke-width="2"/><text x="50" y="60" font-family="system-ui, sans-serif" font-size="9" font-weight="bold" fill="%231e293b" text-anchor="middle">GeoTIFF</text></svg>`;
|
| 150 |
+
} else {
|
| 151 |
+
const reader = new FileReader();
|
| 152 |
+
reader.onload = (e) => {
|
| 153 |
+
previewImg.src = e.target.result;
|
| 154 |
+
};
|
| 155 |
+
reader.readAsDataURL(file);
|
| 156 |
+
}
|
| 157 |
+
|
| 158 |
+
dropZone.style.display = "none";
|
| 159 |
+
uploadPreview.style.display = "block";
|
| 160 |
+
}
|
| 161 |
+
|
| 162 |
+
// -------------------------------------------------------------------------
|
| 163 |
+
// Leaflet Map Integration
|
| 164 |
+
// -------------------------------------------------------------------------
|
| 165 |
+
let map;
|
| 166 |
+
let centerMarker;
|
| 167 |
+
let searchRadiusCircle;
|
| 168 |
+
let resultsMarkersGroup = L.layerGroup();
|
| 169 |
+
let h3GridGroup = L.layerGroup();
|
| 170 |
+
|
| 171 |
+
const indiaCenter = [20.5937, 78.9629];
|
| 172 |
+
|
| 173 |
+
function initMap() {
|
| 174 |
+
const mapEl = document.getElementById("map");
|
| 175 |
+
if (!mapEl) return;
|
| 176 |
+
|
| 177 |
+
map = L.map("map").setView(indiaCenter, 5);
|
| 178 |
+
|
| 179 |
+
// Light-themed GIS Tile Layer (Esri World Topo Map)
|
| 180 |
+
L.tileLayer('https://server.arcgisonline.com/ArcGIS/rest/services/World_Topo_Map/MapServer/tile/{z}/{y}/{x}', {
|
| 181 |
+
attribution: 'Tiles © Esri — Sources: Esri, DeLorme, NAVTEQ, USGS, NOAA, and the GIS User Community'
|
| 182 |
+
}).addTo(map);
|
| 183 |
+
|
| 184 |
+
// Set pointer marker
|
| 185 |
+
centerMarker = L.marker(indiaCenter, { draggable: true }).addTo(map);
|
| 186 |
+
|
| 187 |
+
// Draw initial radius
|
| 188 |
+
updateMapCircle(indiaCenter, parseInt(document.getElementById("radius-input").value));
|
| 189 |
+
|
| 190 |
+
// Sync coordinate boxes on marker drag
|
| 191 |
+
centerMarker.on("dragend", (e) => {
|
| 192 |
+
const latlng = e.target.getLatLng();
|
| 193 |
+
document.getElementById("lat-input").value = latlng.lat.toFixed(5);
|
| 194 |
+
document.getElementById("lon-input").value = latlng.lng.toFixed(5);
|
| 195 |
+
updateMapCircle([latlng.lat, latlng.lng], parseInt(document.getElementById("radius-input").value));
|
| 196 |
+
});
|
| 197 |
+
|
| 198 |
+
// Click on map to place coordinate
|
| 199 |
+
map.on("click", (e) => {
|
| 200 |
+
const latlng = e.latlng;
|
| 201 |
+
centerMarker.setLatLng(latlng);
|
| 202 |
+
document.getElementById("lat-input").value = latlng.lat.toFixed(5);
|
| 203 |
+
document.getElementById("lon-input").value = latlng.lng.toFixed(5);
|
| 204 |
+
updateMapCircle([latlng.lat, latlng.lng], parseInt(document.getElementById("radius-input").value));
|
| 205 |
+
});
|
| 206 |
+
|
| 207 |
+
resultsMarkersGroup.addTo(map);
|
| 208 |
+
h3GridGroup.addTo(map);
|
| 209 |
+
}
|
| 210 |
+
|
| 211 |
+
function updateMapCircle(coords, radiusKm) {
|
| 212 |
+
if (!map) return;
|
| 213 |
+
if (searchRadiusCircle) {
|
| 214 |
+
map.removeLayer(searchRadiusCircle);
|
| 215 |
+
}
|
| 216 |
+
searchRadiusCircle = L.circle(coords, {
|
| 217 |
+
radius: radiusKm * 1000,
|
| 218 |
+
color: "#2563eb",
|
| 219 |
+
fillColor: "#3b82f6",
|
| 220 |
+
fillOpacity: 0.1,
|
| 221 |
+
weight: 1.5
|
| 222 |
+
}).addTo(map);
|
| 223 |
+
}
|
| 224 |
+
|
| 225 |
+
// Connect slider inputs
|
| 226 |
+
const radiusInput = document.getElementById("radius-input");
|
| 227 |
+
const radiusVal = document.getElementById("radius-val");
|
| 228 |
+
radiusInput.addEventListener("input", (e) => {
|
| 229 |
+
const val = e.target.value;
|
| 230 |
+
radiusVal.textContent = `${val} km`;
|
| 231 |
+
if (map) {
|
| 232 |
+
const lat = parseFloat(document.getElementById("lat-input").value);
|
| 233 |
+
const lon = parseFloat(document.getElementById("lon-input").value);
|
| 234 |
+
updateMapCircle([lat, lon], parseInt(val));
|
| 235 |
+
}
|
| 236 |
+
});
|
| 237 |
+
|
| 238 |
+
const kInput = document.getElementById("k-input");
|
| 239 |
+
const kVal = document.getElementById("k-val");
|
| 240 |
+
kInput.addEventListener("input", (e) => {
|
| 241 |
+
kVal.textContent = e.target.value;
|
| 242 |
+
});
|
| 243 |
+
|
| 244 |
+
initMap();
|
| 245 |
+
|
| 246 |
+
// -------------------------------------------------------------------------
|
| 247 |
+
// Execution and API Requests
|
| 248 |
+
// -------------------------------------------------------------------------
|
| 249 |
+
const searchBtn = document.getElementById("search-btn");
|
| 250 |
+
const resultsEmpty = document.getElementById("results-empty");
|
| 251 |
+
const resultsGrid = document.getElementById("results-grid");
|
| 252 |
+
const resultsTableContainer = document.getElementById("results-table-container");
|
| 253 |
+
const resultsTableBody = document.getElementById("results-table-body");
|
| 254 |
+
const timingFeed = document.getElementById("timing-feed");
|
| 255 |
+
|
| 256 |
+
searchBtn.addEventListener("click", () => {
|
| 257 |
+
if (activeSearchMode === "image" && !uploadedFile) {
|
| 258 |
+
alert("Please upload a query satellite image first.");
|
| 259 |
+
return;
|
| 260 |
+
}
|
| 261 |
+
if (activeSearchMode === "text" && !document.getElementById("text-query-input").value.trim()) {
|
| 262 |
+
alert("Please enter a text search query.");
|
| 263 |
+
return;
|
| 264 |
+
}
|
| 265 |
+
|
| 266 |
+
searchBtn.disabled = true;
|
| 267 |
+
searchBtn.innerHTML = `<i class="fa-solid fa-circle-notch fa-spin"></i> EXECUTING RETRIEVAL...`;
|
| 268 |
+
timingFeed.innerHTML = `<i class="fa-solid fa-satellite fa-spin"></i> Aligning sensor domains...`;
|
| 269 |
+
|
| 270 |
+
const formData = new FormData();
|
| 271 |
+
const level = document.getElementById("retrieval-level").value;
|
| 272 |
+
const k = document.getElementById("k-input").value;
|
| 273 |
+
const modality = document.getElementById("query-modality").value;
|
| 274 |
+
|
| 275 |
+
let url = "/api/search";
|
| 276 |
+
|
| 277 |
+
formData.append("k", k);
|
| 278 |
+
formData.append("level", level);
|
| 279 |
+
formData.append("query_modality", modality);
|
| 280 |
+
|
| 281 |
+
if (level === "level4") {
|
| 282 |
+
const lat = document.getElementById("lat-input").value;
|
| 283 |
+
const lon = document.getElementById("lon-input").value;
|
| 284 |
+
const radius = document.getElementById("radius-input").value;
|
| 285 |
+
formData.append("lat", lat);
|
| 286 |
+
formData.append("lon", lon);
|
| 287 |
+
formData.append("radius_km", radius);
|
| 288 |
+
}
|
| 289 |
+
|
| 290 |
+
if (activeSearchMode === "image") {
|
| 291 |
+
formData.append("file", uploadedFile);
|
| 292 |
+
} else {
|
| 293 |
+
url = "/api/search-text";
|
| 294 |
+
formData.append("text_query", document.getElementById("text-query-input").value);
|
| 295 |
+
}
|
| 296 |
+
|
| 297 |
+
fetch(url, {
|
| 298 |
+
method: "POST",
|
| 299 |
+
body: formData
|
| 300 |
+
})
|
| 301 |
+
.then(response => {
|
| 302 |
+
if (!response.ok) {
|
| 303 |
+
throw new Error("Server error occurred during retrieval.");
|
| 304 |
+
}
|
| 305 |
+
return response.json();
|
| 306 |
+
})
|
| 307 |
+
.then(data => {
|
| 308 |
+
renderResults(data);
|
| 309 |
+
})
|
| 310 |
+
.catch(err => {
|
| 311 |
+
console.error(err);
|
| 312 |
+
alert(`Search failed: ${err.message}`);
|
| 313 |
+
timingFeed.innerHTML = `<i class="fa-solid fa-triangle-exclamation" style="color:var(--error)"></i> Error in search`;
|
| 314 |
+
})
|
| 315 |
+
.finally(() => {
|
| 316 |
+
searchBtn.disabled = false;
|
| 317 |
+
searchBtn.innerHTML = `<i class="fa-solid fa-circle-play"></i> RUN ALIGNMENT & SEARCH`;
|
| 318 |
+
});
|
| 319 |
+
});
|
| 320 |
+
|
| 321 |
+
// -------------------------------------------------------------------------
|
| 322 |
+
// Results Rendering & Leaflet Marker Plotting
|
| 323 |
+
// -------------------------------------------------------------------------
|
| 324 |
+
function renderResults(data) {
|
| 325 |
+
const results = data.results;
|
| 326 |
+
const timeMs = data.query_time_ms.toFixed(0);
|
| 327 |
+
const device = data.device || "cpu";
|
| 328 |
+
|
| 329 |
+
// Update telemetry
|
| 330 |
+
document.getElementById("telemetry-device").textContent = device.toUpperCase();
|
| 331 |
+
timingFeed.innerHTML = `<i class="fa-solid fa-bolt"></i> Completed in ${timeMs}ms`;
|
| 332 |
+
|
| 333 |
+
// Clear existing markers & overlays
|
| 334 |
+
resultsMarkersGroup.clearLayers();
|
| 335 |
+
h3GridGroup.clearLayers();
|
| 336 |
+
|
| 337 |
+
if (results.length === 0) {
|
| 338 |
+
resultsEmpty.style.display = "flex";
|
| 339 |
+
resultsGrid.style.display = "none";
|
| 340 |
+
resultsTableContainer.style.display = "none";
|
| 341 |
+
return;
|
| 342 |
+
}
|
| 343 |
+
|
| 344 |
+
resultsEmpty.style.display = "none";
|
| 345 |
+
resultsGrid.style.display = "grid";
|
| 346 |
+
resultsTableContainer.style.display = "block";
|
| 347 |
+
|
| 348 |
+
resultsGrid.innerHTML = "";
|
| 349 |
+
resultsTableBody.innerHTML = "";
|
| 350 |
+
|
| 351 |
+
const mapBounds = L.latLngBounds();
|
| 352 |
+
|
| 353 |
+
results.forEach((item, index) => {
|
| 354 |
+
const rank = index + 1;
|
| 355 |
+
const score = item.score.toFixed(4);
|
| 356 |
+
const lat = item.lat ? item.lat.toFixed(5) : "N/A";
|
| 357 |
+
const lon = item.lon ? item.lon.toFixed(5) : "N/A";
|
| 358 |
+
const dist = item.distance_km ? `${item.distance_km.toFixed(1)} km` : "N/A";
|
| 359 |
+
const h3Cell = item.h3_cell || "N/A";
|
| 360 |
+
|
| 361 |
+
// Render Card
|
| 362 |
+
const card = document.createElement("div");
|
| 363 |
+
card.className = "result-card card";
|
| 364 |
+
|
| 365 |
+
// True color vs FCC render toggle for MS
|
| 366 |
+
const isMs = item.modality === "multispectral";
|
| 367 |
+
const actionHTML = isMs ? `
|
| 368 |
+
<div class="result-actions">
|
| 369 |
+
<button class="action-btn toggle-fcc-btn" data-path="${item.original_path}" data-active="rgb" title="Toggle NIR False Color Composite (FCC)">
|
| 370 |
+
<i class="fa-solid fa-wand-magic-sparkles"></i> FCC
|
| 371 |
+
</button>
|
| 372 |
+
<button class="action-btn plot-spectral-btn" data-path="${item.original_path}" title="Plot Spectral Signature">
|
| 373 |
+
<i class="fa-solid fa-chart-line"></i>
|
| 374 |
+
</button>
|
| 375 |
+
</div>
|
| 376 |
+
` : '';
|
| 377 |
+
|
| 378 |
+
card.innerHTML = `
|
| 379 |
+
<div class="result-img-wrapper">
|
| 380 |
+
<img src="${item.gallery_path}" alt="${item.class}" id="img-result-${index}">
|
| 381 |
+
<span class="result-score-badge">S: ${score}</span>
|
| 382 |
+
<span class="result-modality-tag">${item.modality}</span>
|
| 383 |
+
${actionHTML}
|
| 384 |
+
</div>
|
| 385 |
+
<div class="result-content">
|
| 386 |
+
<div class="result-meta-title">
|
| 387 |
+
<h4>${item.class}</h4>
|
| 388 |
+
<span class="result-meta-rank">#0${rank}</span>
|
| 389 |
+
</div>
|
| 390 |
+
<div class="result-details">
|
| 391 |
+
<span><strong>Lat:</strong> ${lat}</span>
|
| 392 |
+
<span><strong>Lon:</strong> ${lon}</span>
|
| 393 |
+
<span><strong>Dist:</strong> ${dist}</span>
|
| 394 |
+
<span><strong>H3:</strong> ${h3Cell}</span>
|
| 395 |
+
</div>
|
| 396 |
+
</div>
|
| 397 |
+
`;
|
| 398 |
+
resultsGrid.appendChild(card);
|
| 399 |
+
|
| 400 |
+
// Add FCC toggle listener
|
| 401 |
+
if (isMs) {
|
| 402 |
+
const btn = card.querySelector(".toggle-fcc-btn");
|
| 403 |
+
btn.addEventListener("click", (e) => {
|
| 404 |
+
e.stopPropagation();
|
| 405 |
+
const activeMode = btn.getAttribute("data-active");
|
| 406 |
+
const path = btn.getAttribute("data-path");
|
| 407 |
+
const imgEl = document.getElementById(`img-result-${index}`);
|
| 408 |
+
|
| 409 |
+
if (activeMode === "rgb") {
|
| 410 |
+
btn.setAttribute("data-active", "fcc");
|
| 411 |
+
btn.innerHTML = `<i class="fa-solid fa-image"></i> RGB`;
|
| 412 |
+
imgEl.src = `/api/render-bands?path=${encodeURIComponent(path)}&bands=FCC`;
|
| 413 |
+
} else {
|
| 414 |
+
btn.setAttribute("data-active", "rgb");
|
| 415 |
+
btn.innerHTML = `<i class="fa-solid fa-wand-magic-sparkles"></i> FCC`;
|
| 416 |
+
imgEl.src = `/api/render-bands?path=${encodeURIComponent(path)}&bands=RGB`;
|
| 417 |
+
}
|
| 418 |
+
});
|
| 419 |
+
|
| 420 |
+
// Add spectral plot listener
|
| 421 |
+
const spectralBtn = card.querySelector(".plot-spectral-btn");
|
| 422 |
+
spectralBtn.addEventListener("click", () => {
|
| 423 |
+
openSpectralModal(item.original_path, item.class);
|
| 424 |
+
});
|
| 425 |
+
}
|
| 426 |
+
|
| 427 |
+
// Render Table Row
|
| 428 |
+
const row = document.createElement("tr");
|
| 429 |
+
row.innerHTML = `
|
| 430 |
+
<td><strong>0${rank}</strong></td>
|
| 431 |
+
<td><img src="${item.gallery_path}" class="table-thumbnail"></td>
|
| 432 |
+
<td><strong>${item.class}</strong></td>
|
| 433 |
+
<td><span class="badge">${item.modality.toUpperCase()}</span></td>
|
| 434 |
+
<td><code style="font-family:var(--font-mono)">${h3Cell}</code></td>
|
| 435 |
+
<td><code>${lat}, ${lon}</code></td>
|
| 436 |
+
<td>${dist}</td>
|
| 437 |
+
<td class="table-score">${score}</td>
|
| 438 |
+
`;
|
| 439 |
+
resultsTableBody.appendChild(row);
|
| 440 |
+
|
| 441 |
+
// Plot leaf-let marker
|
| 442 |
+
if (item.lat && item.lon) {
|
| 443 |
+
const pCoords = [item.lat, item.lon];
|
| 444 |
+
const marker = L.marker(pCoords).addTo(resultsMarkersGroup);
|
| 445 |
+
|
| 446 |
+
// Marker popup with thumbnail
|
| 447 |
+
marker.bindPopup(`
|
| 448 |
+
<div style="width:140px; text-align:center;">
|
| 449 |
+
<img src="${item.gallery_path}" style="width:100%; border-radius:4px; margin-bottom:0.375rem;">
|
| 450 |
+
<strong>Rank #${rank} • ${item.class}</strong><br>
|
| 451 |
+
Modality: ${item.modality}<br>
|
| 452 |
+
Score: ${score}<br>
|
| 453 |
+
Dist: ${dist}
|
| 454 |
+
</div>
|
| 455 |
+
`);
|
| 456 |
+
|
| 457 |
+
mapBounds.extend(pCoords);
|
| 458 |
+
|
| 459 |
+
// If H3 cell boundary polygon coordinates are returned, draw the hexagon
|
| 460 |
+
if (item.h3_boundary) {
|
| 461 |
+
L.polygon(item.h3_boundary, {
|
| 462 |
+
color: "#10b981",
|
| 463 |
+
fillColor: "#10b981",
|
| 464 |
+
fillOpacity: 0.15,
|
| 465 |
+
weight: 1
|
| 466 |
+
}).addTo(h3GridGroup);
|
| 467 |
+
}
|
| 468 |
+
}
|
| 469 |
+
});
|
| 470 |
+
|
| 471 |
+
// Fit map bounds to show results
|
| 472 |
+
if (map && results.some(item => item.lat && item.lon)) {
|
| 473 |
+
map.fitBounds(mapBounds, { padding: [50, 50] });
|
| 474 |
+
}
|
| 475 |
+
}
|
| 476 |
+
|
| 477 |
+
// -------------------------------------------------------------------------
|
| 478 |
+
// Spectral Modal & Graphing Logic
|
| 479 |
+
// -------------------------------------------------------------------------
|
| 480 |
+
const modal = document.getElementById("spectral-modal");
|
| 481 |
+
const closeModalBtn = document.getElementById("close-modal-btn");
|
| 482 |
+
let spectralChartInstance = null;
|
| 483 |
+
|
| 484 |
+
closeModalBtn.addEventListener("click", () => {
|
| 485 |
+
modal.classList.remove("active");
|
| 486 |
+
});
|
| 487 |
+
|
| 488 |
+
function openSpectralModal(path, className) {
|
| 489 |
+
modal.classList.add("active");
|
| 490 |
+
document.getElementById("modal-title").textContent = `Sentinel-2 Spectral Signature - ${className}`;
|
| 491 |
+
|
| 492 |
+
fetch(`/api/spectral-signature?path=${encodeURIComponent(path)}`)
|
| 493 |
+
.then(res => res.json())
|
| 494 |
+
.then(data => {
|
| 495 |
+
renderSpectralChart(data.reflectance);
|
| 496 |
+
})
|
| 497 |
+
.catch(err => {
|
| 498 |
+
console.error(err);
|
| 499 |
+
alert("Failed to load spectral bands.");
|
| 500 |
+
});
|
| 501 |
+
}
|
| 502 |
+
|
| 503 |
+
function renderSpectralChart(values) {
|
| 504 |
+
const ctx = document.getElementById("spectral-chart").getContext("2d");
|
| 505 |
+
const bands = ["B01 (Aer)", "B02 (B)", "B03 (G)", "B04 (R)", "B05 (RE1)", "B06 (RE2)", "B07 (RE3)", "B08 (NIR)", "B08A (N2)", "B09 (WV)", "B10 (Cir)", "B11 (SW1)", "B12 (SW2)"];
|
| 506 |
+
|
| 507 |
+
if (spectralChartInstance) {
|
| 508 |
+
spectralChartInstance.destroy();
|
| 509 |
+
}
|
| 510 |
+
|
| 511 |
+
spectralChartInstance = new Chart(ctx, {
|
| 512 |
+
type: "line",
|
| 513 |
+
data: {
|
| 514 |
+
labels: bands,
|
| 515 |
+
datasets: [{
|
| 516 |
+
label: "Reflectance Index",
|
| 517 |
+
data: values,
|
| 518 |
+
borderColor: "#2563eb",
|
| 519 |
+
backgroundColor: "rgba(37, 99, 235, 0.1)",
|
| 520 |
+
borderWidth: 2,
|
| 521 |
+
fill: true,
|
| 522 |
+
tension: 0.3
|
| 523 |
+
}]
|
| 524 |
+
},
|
| 525 |
+
options: {
|
| 526 |
+
responsive: true,
|
| 527 |
+
maintainAspectRatio: false,
|
| 528 |
+
scales: {
|
| 529 |
+
y: {
|
| 530 |
+
beginAtZero: true,
|
| 531 |
+
title: { display: true, text: "Scaled Reflectance Value" }
|
| 532 |
+
}
|
| 533 |
+
}
|
| 534 |
+
}
|
| 535 |
+
});
|
| 536 |
+
}
|
| 537 |
+
|
| 538 |
+
// -------------------------------------------------------------------------
|
| 539 |
+
// System Benchmarks Tab Rendering
|
| 540 |
+
// -------------------------------------------------------------------------
|
| 541 |
+
let benchmarksLoaded = false;
|
| 542 |
+
let recallChartInstance = null;
|
| 543 |
+
let latencyChartInstance = null;
|
| 544 |
+
|
| 545 |
+
function loadBenchmarks() {
|
| 546 |
+
if (benchmarksLoaded) return;
|
| 547 |
+
|
| 548 |
+
fetch("/api/benchmarks")
|
| 549 |
+
.then(res => res.json())
|
| 550 |
+
.then(data => {
|
| 551 |
+
renderBenchmarksGrid(data);
|
| 552 |
+
benchmarksLoaded = true;
|
| 553 |
+
})
|
| 554 |
+
.catch(err => {
|
| 555 |
+
console.error("Failed to load benchmarks", err);
|
| 556 |
+
});
|
| 557 |
+
}
|
| 558 |
+
|
| 559 |
+
function renderBenchmarksGrid(results) {
|
| 560 |
+
// Render comparison table
|
| 561 |
+
const tbody = document.getElementById("benchmarks-table-body");
|
| 562 |
+
tbody.innerHTML = "";
|
| 563 |
+
|
| 564 |
+
results.forEach(r => {
|
| 565 |
+
const tr = document.createElement("tr");
|
| 566 |
+
tr.innerHTML = `
|
| 567 |
+
<td><strong>${r.model}</strong></td>
|
| 568 |
+
<td>${r.same_r1.toFixed(3)}</td>
|
| 569 |
+
<td><strong>${r.same_r5.toFixed(3)}</strong></td>
|
| 570 |
+
<td>${r.same_r10.toFixed(3)}</td>
|
| 571 |
+
<td>${r.cross_r1.toFixed(3)}</td>
|
| 572 |
+
<td><strong>${r.cross_r5.toFixed(3)}</strong></td>
|
| 573 |
+
<td>${r.cross_r10.toFixed(3)}</td>
|
| 574 |
+
<td><code>${r.latency_ms.toFixed(0)} ms</code></td>
|
| 575 |
+
`;
|
| 576 |
+
tbody.appendChild(tr);
|
| 577 |
+
});
|
| 578 |
+
|
| 579 |
+
// Plot recall chart
|
| 580 |
+
const ctxRecall = document.getElementById("recall-chart").getContext("2d");
|
| 581 |
+
const labels = results.map(r => r.model);
|
| 582 |
+
const sameR5 = results.map(r => r.same_r5);
|
| 583 |
+
const crossR5 = results.map(r => r.cross_r5);
|
| 584 |
+
|
| 585 |
+
if (recallChartInstance) recallChartInstance.destroy();
|
| 586 |
+
recallChartInstance = new Chart(ctxRecall, {
|
| 587 |
+
type: "bar",
|
| 588 |
+
data: {
|
| 589 |
+
labels: labels,
|
| 590 |
+
datasets: [
|
| 591 |
+
{
|
| 592 |
+
label: "Same-Modal Recall@5",
|
| 593 |
+
data: sameR5,
|
| 594 |
+
backgroundColor: "rgba(37, 99, 235, 0.7)",
|
| 595 |
+
borderColor: "#2563eb",
|
| 596 |
+
borderWidth: 1
|
| 597 |
+
},
|
| 598 |
+
{
|
| 599 |
+
label: "Cross-Modal Recall@5",
|
| 600 |
+
data: crossR5,
|
| 601 |
+
backgroundColor: "rgba(16, 185, 129, 0.7)",
|
| 602 |
+
borderColor: "#10b981",
|
| 603 |
+
borderWidth: 1
|
| 604 |
+
}
|
| 605 |
+
]
|
| 606 |
+
},
|
| 607 |
+
options: {
|
| 608 |
+
responsive: true,
|
| 609 |
+
maintainAspectRatio: false,
|
| 610 |
+
scales: {
|
| 611 |
+
y: {
|
| 612 |
+
beginAtZero: true,
|
| 613 |
+
max: 0.6,
|
| 614 |
+
title: { display: true, text: "Score" }
|
| 615 |
+
}
|
| 616 |
+
}
|
| 617 |
+
}
|
| 618 |
+
});
|
| 619 |
+
|
| 620 |
+
// Plot latency chart
|
| 621 |
+
const ctxLatency = document.getElementById("latency-chart").getContext("2d");
|
| 622 |
+
const latencies = results.map(r => r.latency_ms);
|
| 623 |
+
|
| 624 |
+
if (latencyChartInstance) latencyChartInstance.destroy();
|
| 625 |
+
latencyChartInstance = new Chart(ctxLatency, {
|
| 626 |
+
type: "bar",
|
| 627 |
+
data: {
|
| 628 |
+
labels: labels,
|
| 629 |
+
datasets: [{
|
| 630 |
+
label: "Average Search Latency (ms)",
|
| 631 |
+
data: latencies,
|
| 632 |
+
backgroundColor: "rgba(245, 158, 11, 0.7)",
|
| 633 |
+
borderColor: "#f59e0b",
|
| 634 |
+
borderWidth: 1
|
| 635 |
+
}]
|
| 636 |
+
},
|
| 637 |
+
options: {
|
| 638 |
+
responsive: true,
|
| 639 |
+
maintainAspectRatio: false,
|
| 640 |
+
scales: {
|
| 641 |
+
y: {
|
| 642 |
+
beginAtZero: true,
|
| 643 |
+
title: { display: true, text: "Latency (ms)" }
|
| 644 |
+
}
|
| 645 |
+
}
|
| 646 |
+
}
|
| 647 |
+
});
|
| 648 |
+
}
|
| 649 |
+
|
| 650 |
+
// -------------------------------------------------------------------------
|
| 651 |
+
// Sandbox Queries Click Handler
|
| 652 |
+
// -------------------------------------------------------------------------
|
| 653 |
+
document.querySelectorAll(".sample-query-tile").forEach(tile => {
|
| 654 |
+
tile.addEventListener("click", async () => {
|
| 655 |
+
const path = tile.getAttribute("data-path");
|
| 656 |
+
const name = tile.getAttribute("data-name");
|
| 657 |
+
const type = tile.getAttribute("data-type");
|
| 658 |
+
|
| 659 |
+
// Visually toggle active loading state
|
| 660 |
+
tile.style.borderColor = "var(--primary-color)";
|
| 661 |
+
|
| 662 |
+
try {
|
| 663 |
+
// Fetch the image as a Blob and create a File object
|
| 664 |
+
const res = await fetch(path);
|
| 665 |
+
const blob = await res.blob();
|
| 666 |
+
const file = new File([blob], name, { type: blob.type || "image/png" });
|
| 667 |
+
|
| 668 |
+
// Set search mode to image query
|
| 669 |
+
document.getElementById("mode-image-btn").click();
|
| 670 |
+
|
| 671 |
+
// Select target query modality dropdown
|
| 672 |
+
document.getElementById("query-modality").value = type;
|
| 673 |
+
|
| 674 |
+
// Load file to preview
|
| 675 |
+
handleFile(file);
|
| 676 |
+
|
| 677 |
+
// Automatically trigger search click
|
| 678 |
+
searchBtn.click();
|
| 679 |
+
} catch (err) {
|
| 680 |
+
console.error("Failed to load sample query file:", err);
|
| 681 |
+
alert("Failed to load sample query file: " + err.message);
|
| 682 |
+
} finally {
|
| 683 |
+
setTimeout(() => {
|
| 684 |
+
tile.style.borderColor = "#ced4da";
|
| 685 |
+
}, 1000);
|
| 686 |
+
}
|
| 687 |
+
});
|
| 688 |
+
});
|
| 689 |
+
|
| 690 |
+
// -------------------------------------------------------------------------
|
| 691 |
+
// Rotating Header Banner
|
| 692 |
+
// -------------------------------------------------------------------------
|
| 693 |
+
const bannerMsgs = document.querySelectorAll(".banner-msg");
|
| 694 |
+
if (bannerMsgs.length > 0) {
|
| 695 |
+
let currentMsgIndex = 0;
|
| 696 |
+
setInterval(() => {
|
| 697 |
+
bannerMsgs[currentMsgIndex].classList.remove("active");
|
| 698 |
+
currentMsgIndex = (currentMsgIndex + 1) % bannerMsgs.length;
|
| 699 |
+
bannerMsgs[currentMsgIndex].classList.add("active");
|
| 700 |
+
}, 5000); // rotates every 5 seconds
|
| 701 |
+
}
|
| 702 |
+
|
| 703 |
+
// -------------------------------------------------------------------------
|
| 704 |
+
// Interactive LinkedIn Profile Links
|
| 705 |
+
// -------------------------------------------------------------------------
|
| 706 |
+
const members = ["ayush", "karan", "anurag"];
|
| 707 |
+
members.forEach(member => {
|
| 708 |
+
const linkEl = document.getElementById(`${member}-linkedin-link`);
|
| 709 |
+
const editEl = document.getElementById(`${member}-edit-linkedin`);
|
| 710 |
+
const inputEl = document.getElementById(`${member}-linkedin-input`);
|
| 711 |
+
|
| 712 |
+
function showLink(url) {
|
| 713 |
+
linkEl.href = url;
|
| 714 |
+
linkEl.style.display = "inline-block";
|
| 715 |
+
editEl.style.display = "inline-block";
|
| 716 |
+
inputEl.style.display = "none";
|
| 717 |
+
}
|
| 718 |
+
|
| 719 |
+
function showInput() {
|
| 720 |
+
linkEl.style.display = "none";
|
| 721 |
+
editEl.style.display = "none";
|
| 722 |
+
inputEl.style.display = "inline-block";
|
| 723 |
+
inputEl.value = localStorage.getItem(`${member}_linkedin_url`) || "";
|
| 724 |
+
}
|
| 725 |
+
|
| 726 |
+
const savedUrl = localStorage.getItem(`${member}_linkedin_url`);
|
| 727 |
+
if (savedUrl) {
|
| 728 |
+
showLink(savedUrl);
|
| 729 |
+
} else {
|
| 730 |
+
showInput();
|
| 731 |
+
}
|
| 732 |
+
|
| 733 |
+
inputEl.addEventListener("change", (e) => {
|
| 734 |
+
let url = e.target.value.trim();
|
| 735 |
+
if (url) {
|
| 736 |
+
if (!url.startsWith("http://") && !url.startsWith("https://")) {
|
| 737 |
+
url = "https://" + url;
|
| 738 |
+
}
|
| 739 |
+
localStorage.setItem(`${member}_linkedin_url`, url);
|
| 740 |
+
showLink(url);
|
| 741 |
+
}
|
| 742 |
+
});
|
| 743 |
+
|
| 744 |
+
editEl.addEventListener("click", (e) => {
|
| 745 |
+
e.preventDefault();
|
| 746 |
+
showInput();
|
| 747 |
+
});
|
| 748 |
+
});
|
| 749 |
+
|
| 750 |
+
// -------------------------------------------------------------------------
|
| 751 |
+
// Benchmarks Simulator Logic
|
| 752 |
+
// -------------------------------------------------------------------------
|
| 753 |
+
const simWeight = document.getElementById("sim-modality-weight");
|
| 754 |
+
const simNoise = document.getElementById("sim-noise-level");
|
| 755 |
+
const simH3Res = document.getElementById("sim-h3-resolution");
|
| 756 |
+
|
| 757 |
+
function updateSimulation() {
|
| 758 |
+
if (!simWeight) return;
|
| 759 |
+
const w = parseFloat(simWeight.value);
|
| 760 |
+
const n = parseFloat(simNoise.value);
|
| 761 |
+
const res = parseInt(simH3Res.value);
|
| 762 |
+
|
| 763 |
+
document.getElementById("sim-weight-val").innerText = w.toFixed(2);
|
| 764 |
+
document.getElementById("sim-noise-val").innerText = n.toFixed(2);
|
| 765 |
+
|
| 766 |
+
// Simulated precision recall curves formulas based on ZS-MC characteristics
|
| 767 |
+
const baseR1 = 0.45 + 0.12 * Math.sin(w * Math.PI) - 0.3 * n;
|
| 768 |
+
const r1 = Math.max(0.1, Math.min(0.99, baseR1));
|
| 769 |
+
const r5 = Math.max(r1 + 0.05, Math.min(0.99, r1 * 1.15 + 0.05));
|
| 770 |
+
const mapVal = Math.max(r5 + 0.05, Math.min(0.99, r5 * 1.12 + 0.04));
|
| 771 |
+
|
| 772 |
+
// Latency depends on H3 resolution level
|
| 773 |
+
const latency = Math.round(28 + (8 - res) * 4 + w * 2 + n * 2);
|
| 774 |
+
|
| 775 |
+
document.getElementById("sim-metric-r1").innerText = (r1 * 100).toFixed(1) + "%";
|
| 776 |
+
document.getElementById("sim-metric-r5").innerText = (r5 * 100).toFixed(1) + "%";
|
| 777 |
+
document.getElementById("sim-metric-map").innerText = (mapVal * 100).toFixed(1) + "%";
|
| 778 |
+
document.getElementById("sim-metric-latency").innerText = latency + " ms";
|
| 779 |
+
}
|
| 780 |
+
|
| 781 |
+
if (simWeight) {
|
| 782 |
+
simWeight.addEventListener("input", updateSimulation);
|
| 783 |
+
simNoise.addEventListener("input", updateSimulation);
|
| 784 |
+
simH3Res.addEventListener("change", updateSimulation);
|
| 785 |
+
updateSimulation();
|
| 786 |
+
}
|
| 787 |
+
|
| 788 |
+
});
|
src/ui/static/app_assets/anurag.jpg
ADDED
|
src/ui/static/app_assets/ayush.jpg
ADDED
|
src/ui/static/app_assets/chart.js
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|
src/ui/static/app_assets/images/marker-icon-2x.png
ADDED
|
|
src/ui/static/app_assets/images/marker-icon.png
ADDED
|
|
src/ui/static/app_assets/images/marker-shadow.png
ADDED
|
src/ui/static/app_assets/karan.jpg
ADDED
|