eo_data_embedding package

eo-data-embedding: multi-modal geospatial embedding search & change detection.

Submodules

eo_data_embedding.baseline module

Supervised CNN baseline for the few-shot comparison.

A ResNet-18 trained from scratch on raw EuroSAT bands — the same 10-band Clay subset and the same split protocol as the linear probe (fixed held-out test set from probe.heldout_split, k-shot training sets from probe.sample_shots). This is the comparison the architecture diagram promises: frozen foundation-model embeddings + linear probe vs a supervised CNN trained directly on the pixels with the same labels.

eo_data_embedding.baseline.build_resnet18(in_chans=10, num_classes=10)

torchvision ResNet-18 with an in_chans-channel stem, no pretrained weights.

ImageNet weights are RGB-only and don’t transfer to 10 multispectral bands, so the baseline trains from scratch — which is exactly the regime the label-efficiency comparison is about.

eo_data_embedding.baseline.cnn_baseline_full(X, y, test_frac=0.2, split_seed=42, **train_kw)

Fully-supervised CNN reference: train on the whole pool, same fixed test set.

Return type:

dict

eo_data_embedding.baseline.cnn_baseline_multi(X, y, shots, seeds=(0, 1, 2, 3, 4), test_frac=0.2, split_seed=42, **train_kw)

k-shot supervised CNN over multiple seeds on the probe’s fixed held-out test set.

Return type:

dict

eo_data_embedding.baseline.train_eval_cnn(X, y, train_idx, test_idx, *, epochs=60, batch_size=64, lr=0.001, weight_decay=0.0001, device='cuda', seed=0)

Train a ResNet-18 on X[train_idx], return macro-F1/accuracy on X[test_idx].

X is a (N, C, H, W) tensor of RAW band values; per-band standardization uses statistics of the training subset only (no test leakage). Augmentation is random flips — EuroSAT patches have no canonical orientation.

Return type:

dict

eo_data_embedding.change module

Embedding-distance change detection (OSCD stretch).

Embed each date of a bitemporal pair (ideally on co-registered tiles — perspective geometry matters here), then threshold the per-tile embedding distance to flag change.

eo_data_embedding.change.binary_change_metrics(y_true, score, threshold)

Change-detection metrics at a FIXED threshold (chosen on train) + threshold-free ROC-AUC.

Returns {f1, precision, recall, iou, kappa, accuracy, roc_auc, threshold}. ROC-AUC is threshold-free; everything else is the operating point at threshold. Kappa and accuracy are the OSCD-literature companions to F1/IoU.

Return type:

dict

eo_data_embedding.change.delta_features(e1, e2, kind='abs')

Difference features for a supervised change probe on frozen embeddings.

e1, e2 are aligned (N, D) embeddings of the two dates (N tiles, or N patches). Returns (N, D) for “abs”/”signed”, (N, 3D) for “concat” ([e1, e2, |e1-e2|]).

Return type:

ndarray

eo_data_embedding.change.embedding_change_score(emb_t1, emb_t2, metric='cosine')

Per-tile change score from two (N, D) embedding arrays.

Return type:

ndarray

eo_data_embedding.change.patch_change_map(p1, p2, grid_hw, metric='cosine')

Per-patch change scores for one tile, reshaped to the spatial grid.

p1, p2 are a single tile’s patch tokens (P, D) from ClayEmbedder.encode(…, return_patches=True). Returns a (gh, gw) array of per-patch distances — a spatial change map at patch resolution (~80 m for Clay’s 8-px patch at 10 m GSD), the granularity zero-shot change methods actually use instead of one global vector per tile.

eo_data_embedding.change.patch_mask_labels(mask_tile, grid_hw, frac=0.05)

Per-patch change labels for one tile: 1 if >frac of a patch cell’s pixels changed.

mask_tile is one tile’s (H, W) 0/1 change mask; grid_hw is the patch grid (gh, gw). Average-pools the mask to the patch grid (fraction changed per patch) and thresholds at frac. Returns a (gh*gw,) int array aligned with patch_change_map(…).reshape(-1).

eo_data_embedding.change.pick_threshold(y_true, score)

Threshold on score that maximises F1 on this split. Pick it on the TRAIN split, then evaluate the held-out split at this fixed threshold — sweeping on the test split itself is an oracle/optimistic operating point. Candidates span the full 0.01–0.99 quantile range of score: a narrow upper-tail grid can pick a threshold above every held-out score, collapsing the test predictions to all-negative (F1 = 0) even when the score is discriminative (ROC-AUC > 0.5).

Return type:

float

eo_data_embedding.change.tile_image(img, size=256)

Pad a (C,H,W) image to multiples of size and return tiles (N,C,size,size).

Uses reflect padding when possible; reflect requires each pad amount < its dimension, so scenes too small for that (e.g. OSCD test scenes shorter than one tile) fall back to replicate (edge) padding, which has no such limit.

eo_data_embedding.change.tile_mask_labels(mask, size=256, frac=0.05)

Per-tile change label: 1 if >frac of the tile’s pixels are changed. Returns (N,) int.

eo_data_embedding.clay_metadata module

Verified Clay v1.5 band metadata (from Clay-foundation/model configs/metadata.yaml).

These are the exact values the model was trained with — wavelengths, per-band mean/std for normalization, band order, and GSD. Sources are quoted in research/04-clay-integration.md.

Clay normalizes pixels as (x - mean) / std, with means/stds reshaped to [1, C, 1, 1]. waves is the per-band central wavelength the model’s patch embedder conditions on.

eo_data_embedding.cli module

eo-data-embedding command-line interface.

eo-data-embedding demo is the plug-and-play CPU demo and works from any install. The phase subcommands (extract/search/probe/…) are thin pass-throughs to the scripts in scripts/, so they only resolve in a source checkout (git clone + pip install -e .).

eo_data_embedding.cli.main(argv=None)

Return type:

int

eo_data_embedding.config module

Lightweight YAML config loader.

Phase scripts pull their argparse defaults from configs/default.yaml via load_config() + cfg_get(), so the config file is the single place to change paths / batch sizes / shots. CLI flags still override. Missing file → {} (scripts fall back to their hardcoded defaults).

eo_data_embedding.config.cfg_get(cfg, dotted, default=None)

Nested lookup, e.g. cfg_get(cfg, "embed.store_path", "artifacts/embeddings.parquet").

Return type:

Any

eo_data_embedding.config.load_config(path='configs/default.yaml')

Load the project config as a dict; returns {} if the file is absent.

Return type:

dict

eo_data_embedding.data module

Dataset loading via TorchGeo.

Phase 0 can run fully synthetic (no download). Phase 1+ pulls real EO patches.

eo_data_embedding.data.bigearthnet_subset(root='data/', n=2000, seed=42)

Aligned multi-modal subset of BigEarthNet-MM, bands reordered for Clay.

Returns a dict:

s2 -> (n, 10, H, W) raw Sentinel-2 reflectance in Clay band order s1 -> (n, 2, H, W) raw Sentinel-1 backscatter (VV, VH) labels -> (n,) int, primary class (argmax of the 19-class multi-hot — a documented

simplification so the few-shot probe is single-label)

ids -> (n,) patch indices (shared across modalities → enables cross-modal retrieval)

Pixels are returned RAW; ClayEmbedder applies Clay’s per-band normalization. NOTE: BigEarthNet is multi-label; reducing to the primary class is a deliberate Phase-1 simplification for the probe demo. Verify TorchGeo’s band order at runtime (see clay_metadata).

eo_data_embedding.data.eurosat_sample(root='data/')

One EuroSAT (Sentinel-2) sample as (image_tensor, label). Downloads ~90MB once.

eo_data_embedding.data.eurosat_subset(root='data/', n=2000, seed=42)

EuroSAT subset (Sentinel-2 optical), bands reordered for Clay. ~2 GB download, single-label.

Returns dict: s2 -> (n, 10, H, W) raw, labels -> (n,) int (10-class), ids -> (n,). No SAR (EuroSAT is optical-only) — fast real-data path for the few-shot probe + retrieval.

eo_data_embedding.data.oscd_pairs(root='data/', split='train', download=False)

OSCD bitemporal change-detection pairs, bands reordered for Clay.

Returns a list of dicts: {id, t1 (10,H,W), t2 (10,H,W), mask (H,W) 0/1}. Raw pixels. OSCD images are full Sentinel-2 scenes of varying size (tiled to 256 in phase5).

root can point to an external/NAS mount holding the (already extracted) OSCD dataset; keep download=False to read it in place without writing to local disk.

eo_data_embedding.data.ssl4eo_crossmodal(n=1000, split='val', device_batch=8)

Stream N aligned Sentinel-1 (SAR) + Sentinel-2 (optical) tiles from SSL4EO-S12 v1.1.

Uses the official webdataset streaming loader — only the first ~N samples are pulled, NOT the whole dataset. Returns dict: s2 (N,10,H,W), s1 (N,2,H,W), ids — paired by location for cross-modal retrieval. Takes time index 0 of the 4 timestamps; bands reordered for Clay.

Requires (install on the GPU host, see research/05-crossmodal.md):

pip install webdataset pip install “git+https://github.com/DLR-MF-DAS/SSL4EO-S12-v1.1.git” # provides ssl4eos12_dataset

VERIFY-AT-RUNTIME: batch key names (“S2L2A”/”S1GRD”) and the 12-band S2 order (assumed [B01,B02,B03,B04,B05,B06,B07,B08,B8A,B09,B11,B12] like BigEarthNet).

eo_data_embedding.data.synthetic_batch(batch=2, chans=3, size=224)

Deterministic synthetic image batch for the no-download sanity path.

Return type:

Tensor

eo_data_embedding.demo module

CPU-only “try it” demo: similarity search + a live probe prediction over EuroSAT.

Plug-and-play (eo-data-embedding demo): downloads EuroSAT, fetches a small prebuilt bundle (precomputed frozen-Clay embeddings + the trained few-shot probe), then serves a Gradio UI. No GPU, no Clay at runtime (embeddings are precomputed), no training (the probe is loaded from the bundle).

Each trial picks a random EuroSAT tile from the probe’s held-out test split (so the prediction is on data the probe never saw), runs the probe to predict its land-use class, and shows the nearest neighbours from a FAISS index.

eo_data_embedding.demo.ensure_eurosat(root='data/')

Return the EuroSAT (Sentinel-2, all bands) train split, downloading it once if absent.

eo_data_embedding.demo.fetch_bundle(dest='demo', url='https://github.com/AstroCan17/eo-data-embedding/releases/latest/download/demo-bundle.zip')

Download + extract the demo bundle (embeddings.parquet + probe.npz) into dest once.

Return type:

Path

eo_data_embedding.demo.main(argv=None)

Return type:

int

eo_data_embedding.demo.serve(bundle='demo', data_root='data/', port=7860)

Launch the Gradio demo: random held-out tile -> live probe prediction + nearest neighbours.

eo_data_embedding.embed module

Embedding backbones.

Phase 0 uses a lightweight timm ViT so the pipeline is verifiable on CPU/Colab. Phase 1 swaps in a real geospatial foundation model (Clay / Prithvi) — same interface: an encoder that maps an image batch (B, C, H, W) -> embeddings (B, D).

class eo_data_embedding.embed.ClayEmbedder(checkpoint=None, modality='s2', device='cuda', image_size=None, metadata_path=None)

Bases: object

Clay v1.5 geospatial foundation model wrapped to the canonical encode(x) -> (B, 1024).

Handles the full multi-band / SAR path: builds Clay’s datacube (pixels + waves + gsd + time + latlon), normalizes per the verified band stats, runs the FROZEN encoder, and returns the class-token embedding. Input x is RAW (un-normalized) reflectance/backscatter of shape (B, C, H, W) with C bands in Clay’s expected order for modality (see clay_metadata).

Install (on the GPU host): pip install claymodel and download the checkpoint from HuggingFace made-with-clay/Clay (clay-v1.5.ckpt). See research/04-clay-integration.md.

VERIFY-AT-RUNTIME (Clay’s API has drifted across versions):

• encoder call path: model.model.encoder(datacube) vs model.encoder(…)

• datacube time/latlon shapes ([B,2] per current main; some versions use [B,4])

Both are isolated below and easy to flip.

Methods

__call__(x[, return_patches])	Raw (B, C, H, W) bands -> embeddings (float, CPU).
encode(x[, return_patches])	Raw (B, C, H, W) bands -> embeddings (float, CPU).

encode(x, return_patches=False)

Raw (B, C, H, W) bands -> embeddings (float, CPU).

Default: the (B, 1024) class-token vector every phase depends on. return_patches=True: the per-patch tokens (B, P, 1024) plus the patch grid (gh, gw) — for spatial change maps that need patch-level resolution rather than one vector per tile.

class eo_data_embedding.embed.ViTEmbedder(backbone='vit_small_patch16_224', in_chans=3, pretrained=True, device='cpu')

Bases: Module

Frozen ViT encoder returning a single embedding vector per image.

Uses timm forward_features + global pool. This is the sanity/baseline backbone; for EO foundation models see load_clay / load_prithvi (Phase 1).

Methods

add_module(name, module)	Add a child module to the current module.
apply(fn)	Apply fn recursively to every submodule (as returned by .children()) as well as self.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Return an iterator over module buffers.
children()	Return an iterator over immediate children modules.
compile(*args, **kwargs)	Compile this Module's forward using torch.compile().
cpu()	Move all model parameters and buffers to the CPU.
cuda([device])	Move all model parameters and buffers to the GPU.
double()	Casts all floating point parameters and buffers to double datatype.
encode(x)	Canonical interface every phase depends on: images -> (B, D) embeddings.
eval()	Set the module in evaluation mode.
extra_repr()	Return the extra representation of the module.
float()	Casts all floating point parameters and buffers to float datatype.
forward(x)	Define the computation performed at every call.
get_buffer(target)	Return the buffer given by target if it exists, otherwise throw an error.
get_extra_state()	Return any extra state to include in the module's state_dict.
get_parameter(target)	Return the parameter given by target if it exists, otherwise throw an error.
get_submodule(target)	Return the submodule given by target if it exists, otherwise throw an error.
half()	Casts all floating point parameters and buffers to half datatype.
ipu([device])	Move all model parameters and buffers to the IPU.
load_state_dict(state_dict[, strict, assign])	Copy parameters and buffers from state_dict into this module and its descendants.
modules([remove_duplicate])	Return an iterator over all modules in the network.
mtia([device])	Move all model parameters and buffers to the MTIA.
named_buffers([prefix, recurse, ...])	Return an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Return an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Return an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Return an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
parameters([recurse])	Return an iterator over module parameters.
register_backward_hook(hook)	Register a backward hook on the module.
register_buffer(name, tensor[, persistent])	Add a buffer to the module.
register_forward_hook(hook, *[, prepend, ...])	Register a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Register a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Register a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Register a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Register a post-hook to be run after module's load_state_dict() is called.
register_load_state_dict_pre_hook(hook)	Register a pre-hook to be run before module's load_state_dict() is called.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Add a parameter to the module.
register_state_dict_post_hook(hook)	Register a post-hook for the state_dict() method.
register_state_dict_pre_hook(hook)	Register a pre-hook for the state_dict() method.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
set_extra_state(state)	Set extra state contained in the loaded state_dict.
set_submodule(target, module[, strict])	Set the submodule given by target if it exists, otherwise throw an error.
share_memory()	See torch.Tensor.share_memory_().
state_dict(*args[, destination, prefix, ...])	Return a dictionary containing references to the whole state of the module.
to(*args, **kwargs)	Move and/or cast the parameters and buffers.
to_empty(*, device[, recurse])	Move the parameters and buffers to the specified device without copying storage.
train([mode])	Set the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
xpu([device])	Move all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Reset gradients of all model parameters.

__call__

encode(x)

Canonical interface every phase depends on: images -> (B, D) embeddings.

Return type:

Tensor

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses. :rtype: Tensor

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

eo_data_embedding.embed.load_embedder(name='timm-vit', **kw)

Factory returning anything with encode(x) -> (B, D).

eo_data_embedding.log module

Minimal logging setup shared by the phase scripts.

Use log = get_logger("extract") then log.info("...") instead of print. Level comes from the GEO_LOG_LEVEL env var (default INFO). The logger name appears as [extract] in each line, replacing the old print(f"[extract] ...") tag.

eo_data_embedding.log.get_logger(name)

Return type:

Logger

eo_data_embedding.probe module

Few-shot linear probe on frozen embeddings.

Demonstrates the foundation-model value prop: train a linear classifier on top of frozen embeddings with very few labels per class and compare against a fully-supervised baseline. The headline result is the few-shot vs full-label metric table.

Protocol: heldout_split carves out one stratified test set that stays FIXED across shot levels and seeds; sample_shots draws each k-shot training set from the remaining pool. linear_probe_multi repeats that over several seeds and reports mean±std — single-seed few-shot numbers (especially 5-shot) vary too much to quote alone.

class eo_data_embedding.probe.LinearProbe(coef, intercept, classes)

Bases: object

A saved linear probe applied with a pure-numpy forward pass.

Holds a fitted classifier’s coef_ (n_classes, D), intercept_ (n_classes,) and classes_. Storing the learned parameters (not a pickled estimator) keeps the demo bundle independent of the scikit-learn version it was trained with.

Methods

decision
predict
predict_proba

decision(X)

Return type:

ndarray

predict(X)

Return type:

ndarray

predict_proba(X)

Return type:

ndarray

eo_data_embedding.probe.few_shot_split(labels, shots, seed=42)

Indices for shots labelled examples per class, rest as the test pool.

Legacy single-draw protocol (test pool changes with shots) — kept for the Phase-0 smoke gate; the reported results use linear_probe_multi.

eo_data_embedding.probe.full_probe(X, y, test_frac=0.2, split_seed=42)

Fully-supervised reference: train on the entire pool, evaluate on the same fixed test set.

Return type:

dict

eo_data_embedding.probe.heldout_split(labels, test_frac=0.2, seed=42)

Stratified (train-pool, test) index split; the test set is the fixed evaluation set.

eo_data_embedding.probe.linear_probe(X, y, shots, seed=42)

Single-seed probe with shots labels/class (legacy protocol; see linear_probe_multi).

Return type:

dict

eo_data_embedding.probe.linear_probe_multi(X, y, shots, seeds=(0, 1, 2, 3, 4), test_frac=0.2, split_seed=42)

k-shot probe over multiple seeds on one fixed held-out test set; mean±std metrics.

The test set depends only on (test_frac, split_seed), so every shot level and every seed is evaluated on identical data — numbers are comparable across rows.

Return type:

dict

eo_data_embedding.probe.load_probe(path)

Load an npz probe written by save_probe into a numpy-only LinearProbe.

Return type:

LinearProbe

eo_data_embedding.probe.sample_shots(labels, pool_idx, shots, seed)

Draw shots training indices per class from the train pool only.

Return type:

ndarray

eo_data_embedding.probe.save_probe(clf, path)

Persist a fitted linear probe as a version-independent npz (coef / intercept / classes).

Return type:

str

eo_data_embedding.probe.train_probe(X, y, test_frac=0.2, split_seed=42)

Fit the demo classifier on the train pool (the fixed held-out test set is excluded).

Returns (clf, test_idx) — a fitted scikit-learn LogisticRegression and the held-out indices, so callers can both persist the probe and report honest accuracy on data it never saw.

eo_data_embedding.search module

FAISS similarity search over the embedding store.

eo_data_embedding.search.build_index(vectors, normalize=True)

Build a FAISS index. Cosine similarity via inner product on L2-normalized vectors.

eo_data_embedding.search.retrieval_metrics(neigh_labels, query_labels, class_total=None)

Label-based retrieval quality over each query’s ranked top-k neighbours (self removed).

neigh_labels is (Q, k): the class label of every query’s k nearest neighbours, nearest first. query_labels is (Q,). class_total is the per-class corpus size (array indexed by label) used as recall’s denominator; if omitted it is derived from query_labels (valid when the queries are the whole corpus, as in phase 2). Queries whose class is a singleton (no other relevant item) are excluded from recall/mAP.

Return type:

dict

Returns {precision, recall, map, k}:

• precision@k — fraction of retrieved neighbours sharing the query’s class.

• recall@k — retrieved relevant / total relevant (corpus same-class count minus self).

• mAP@k — mean over queries of AP@k = Σ_i P@i·rel_i / min(R, k), R = total relevant.

eo_data_embedding.search.search(index, queries, top_k=12, normalize=True)

Return (distances, indices) for each query row.

eo_data_embedding.store module

Parquet-backed embedding store.

The whole point of decoupling: extract embeddings ONCE (heavy GPU pass), persist them, then run search / probe / change-detection cheaply many times.

eo_data_embedding.store.load_embeddings(path)

Return type:

DataFrame

eo_data_embedding.store.save_embeddings(path, ids, vectors, modality, labels=None)

Persist embeddings as parquet. vectors is (N, D).

Return type:

Path

eo_data_embedding.store.stack_vectors(df)

(N, D) float32 matrix from the vector column.

Return type:

ndarray