feat(models): integrate OLMoEarth embedding model

[go-bananas-wwj]

feat(models): integrate OLMoEarth embedding model

Overview

This PR adds OLMoEarth (Allen AI's Earth-system foundation model) as the 5th embedding model to EarthEmbeddingExplorer, alongside SigLIP, FarSLIP, SatCLIP, and DINOv2.

OLMoEarth is a multimodal, spatio-temporal foundation model trained on Sentinel-2 L2A and 6 derived geospatial maps (OpenStreetMap, WorldCover, SRTM DEM, etc.). It excels at capturing spectral and spatial patterns from 12-band multispectral imagery, making it a strong candidate for pure visual similarity search in remote-sensing applications.

Paper: Herzog et al., 2025 — OlmoEarth: Stable Latent Image Modeling for Multimodal Earth Observation

---

What Changed

Commit	Scope	Description
feat(models)	models/olmoearth_model.py, core/model_manager.py, generate_embeddings.py, requirements.txt	Add OLMoEarth model wrapper with MEAN pooling, L2-normalized cosine search, and MajorTOM band reordering
fix(search_engine, ui)	core/search_engine.py, ui/callbacks.py	Fallback to other models' df_embed when active model lacks parquet_url metadata
fix(majortom)	MajorTOM/embedder/MajorTOM_Embedder.py	Use INTER_CUBIC instead of INTER_NEAREST for 12-band image quality
docs(readme, doc, contributing)	README.md, doc.md, CONTRIBUTING.md	Document OLMoEarth as the 5th model; add dataset link and torch compatibility notes
chore(gitignore)	.gitignore	Exclude generated embedding artifacts (embeddings_*/, *.pkl, *.npy)
style(...)	core/search_engine.py, generate_embeddings.py, ui/callbacks.py	Fix pre-existing ruff warnings and apply formatting

---

Key Technical Details

1. Band Reordering
MajorTOM stores bands as [B01, B02, B03, B04, B05, B06, B07, B08, B8A, B09, B11, B12], but OLMoEarth expects [B02, B03, B04, B08, B05, B06, B07, B8A, B11, B12, B01, B09]. The wrapper handles this transparently in _prepare_input().

2. Pooling Strategy
We use PoolingType.MEAN (matching the official allenai/olmoearth_ml4rs_tutorial) rather than MAX pooling. Raw embeddings have L2 norm ~10.3 and are not normalized during storage/encoding to stay consistent with the official tutorial.

3. Retrieval Normalization
To avoid polar clustering (polar embeddings have systematically higher norms ~10.5 vs ~10.2), we apply F.normalize() only inside search(), converting dot products to cosine similarity in [-1, 1].

4. Metadata Gap Workaround
Our self-generated OLMoEarth embeddings lack parquet_url and parquet_row columns. We added fallback logic so that _fetch_top_k_images(), get_initial_plot(), and download_image_by_location() can borrow metadata from any other loaded model. All 5 models share the same 248,719 samples in identical order (verified by product_id).

---

Environment Compatibility

olmoearth-pretrain-minimal requires torch >= 2.8, < 2.9. For users with older PyTorch versions, we recommend a dedicated conda environment:

conda create -n eee python=3.12
conda activate eee
conda install pytorch==2.8.0 torchvision==0.23.0 pytorch-cuda=12.4 -c pytorch -c nvidia
pip install -r requirements.txt

---

Verification

• ✅ ruff check . — all checks passed
• ✅ ruff format . --check — all formatted
• ✅ python -c "import app" — all 5 models load successfully (248,719 records each)
• ✅ 20 random image queries tested; ~80% top-1 accuracy, no polar clustering

---

Dataset

• ModelScope: WeijieWu/olmoearth_embdding (824 MB, 248,719 rows, 768-dim float32)
• Source images: Major-TOM/Core-S2L2A-249k (same as other 4 models)

Note: configs/config.yaml contains a placeholder ms://Major-TOM/Core-S2RGB-249k-OLMoEarth/... path consistent with the other 4 models. Users should create configs/config_local.yaml with the actual local path or the WeijieWu/olmoearth_embdding ModelScope URL.

---

Checklist

• Model wrapper follows the unified interface (encode_image, search)
• Registered in core/model_manager.py and generate_embeddings.py
• Added to requirements.txt
• Updated README.md, doc.md, and CONTRIBUTING.md
• Ruff linting and formatting passed
• Local import test passed
• Embedding files kept outside the repository (/workspace/EEE/embeddings/)