Clay GeoAI Embeddings - Cosine Similarity Validation

End-to-end pipeline that uses the Clay Foundation Model (v1.5) to generate geospatial embeddings from real Sentinel-2 satellite imagery and validates them using cosine similarity.

Overview

Clay is an open-source foundation model for Earth observation built on a Vision Transformer (ViT) architecture, trained via self-supervised learning using a Masked Autoencoder (MAE) approach. It takes multi-band satellite imagery along with spatiotemporal metadata (location, time) and produces dense vector embeddings that encode the semantic content of the landscape.

This project fetches real satellite data for a small region (Bangalore, India), generates Clay embeddings for image patches, and validates that the embeddings capture meaningful spatial structure via cosine similarity analysis.

Project Structure

embeddings/
├── run_clay_embeddings.py     # Full pipeline: fetch data -> embeddings -> similarity
├── validate_pipeline.py       # Lightweight validation (no checkpoint needed)
├── configs/
│   └── metadata.yaml          # Clay sensor metadata (wavelengths, normalization stats)
├── v1.5/
│   └── clay-v1.5.ckpt         # Clay v1.5 model checkpoint (~5GB)
├── requirements.txt           # pip dependencies
├── environment.yml            # conda/mamba environment spec
├── .venv/                     # Python virtual environment
└── clay_similarity_results.png # Output visualization (after running)

Setup

1. Create environment and install dependencies

Using pip + venv:

python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

Using conda/mamba:

conda env create --file environment.yml
conda activate clay-geoai

2. Download Clay v1.5 checkpoint

The model weights (~5GB) are hosted on HuggingFace:

python -c "from huggingface_hub import hf_hub_download; hf_hub_download('made-with-clay/Clay', 'v1.5/clay-v1.5.ckpt', local_dir='.')"

Or manually:

mkdir -p v1.5
wget -P v1.5/ https://huggingface.co/made-with-clay/Clay/resolve/main/v1.5/clay-v1.5.ckpt

3. Sensor metadata

The configs/metadata.yaml file (sourced from the Clay repo) contains per-sensor band order, wavelengths, normalization statistics (mean/std), and GSD for all supported sensors (Sentinel-2, Landsat, NAIP, MODIS, etc.).

Usage

Validate the pipeline first (no checkpoint needed)

python validate_pipeline.py

This tests:

• STAC data fetch from Microsoft Planetary Computer
• Patch extraction from the satellite image
• Cosine similarity computation (with simulated embeddings)
• Clay model package import

Run the full pipeline

python run_clay_embeddings.py

Set a custom checkpoint path if needed:

CLAY_CHECKPOINT=path/to/clay-v1.5.ckpt python run_clay_embeddings.py

Pipeline Steps

Step 1: Fetch Sentinel-2 Imagery

• Queries the Microsoft Planetary Computer STAC API
• Searches for Sentinel-2 L2A scenes over Bangalore, India (bbox: [77.55, 12.93, 77.65, 13.03])
• Filters for low cloud cover (<10%) in the date range Jan-Mar 2024
• Downloads 10 spectral bands (Blue, Green, Red, RedEdge 1-3, NIR, NIR08, SWIR16, SWIR22) at 10m resolution
• Reprojects to UTM zone 43N (EPSG:32643)

Step 2: Patch Extraction

• Splits the full image into non-overlapping 256x256 pixel patches
• Filters out patches with NaN values or zero variance
• Caps at 16 patches for memory efficiency

Step 3: Normalization

• Normalizes each patch using Clay's per-band mean and standard deviation from configs/metadata.yaml
• Formula: normalized = (pixel_value - mean) / std

Step 4: Generate Embeddings with Clay v1.5

• Loads the Clay v1.5 checkpoint (ViT-Large, ~1.5B parameters)
• Sets encoder to inference mode: mask_ratio=0, shuffle=False
• Constructs the datacube dict for each patch:
  • pixels: Normalized image tensor [1, 10, 256, 256]
  • time: Sinusoidal encoding of acquisition week/hour [1, 4]
  • latlon: Sinusoidal encoding of center lat/lon [1, 4]
  • gsd: Ground sample distance (10m for Sentinel-2)
  • waves: Band wavelengths in micrometers
• Extracts the CLS token (index 0) from the encoder output as the patch embedding
• Each patch produces a 1024-dimensional embedding vector
• Processes one patch at a time to stay within GPU/MPS memory limits

Step 5: Cosine Similarity Validation

Computes the full pairwise cosine similarity matrix and runs three validation checks:

1. Spatial coherence: Adjacent patches (Manhattan distance = 1) should have higher average similarity than distant patches (distance >= 4). This validates that the embeddings capture spatial continuity.

2. Self-similarity: Diagonal values should be exactly 1.0 (a patch is identical to itself).

3. Extremes analysis: Identifies the most and least similar patch pairs with their grid coordinates.

Step 6: Visualization

Saves clay_similarity_results.png with three panels:

• Cosine similarity heatmap: Full N x N pairwise similarity matrix
• RGB patch strip: Visual reference of the first 8 patches
• PCA 2D projection: Embedding space reduced to 2D to visualize clustering

Key Technical Details

Component	Detail
Model	Clay v1.5 (ViT-Large, MAE-based)
Embedding dim	1024 (CLS token from encoder)
Input size	256 x 256 pixels, 10 bands
Patch size (ViT)	8 x 8 pixels
Sensor	Sentinel-2 L2A
Data source	Microsoft Planetary Computer (free, no auth required)
Region	Bangalore, India (urban + vegetation mix)
Similarity metric	Cosine similarity via scikit-learn

Dependencies

• PyTorch >= 2.0 (with MPS/CUDA support)
• claymodel - Clay Foundation Model package
• pystac-client + planetary-computer - STAC catalog access
• stackstac + rioxarray - Cloud-optimized GeoTIFF loading
• scikit-learn - Cosine similarity, PCA
• matplotlib - Visualization

References

• Clay Foundation Model Documentation
• Clay GitHub Repository
• Clay Model Weights on HuggingFace
• Microsoft Planetary Computer
• GeoAI with Python