diff --git a/CHANGELOG.md b/CHANGELOG.md index c5e3697..03e0603 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,11 @@ The format is based on Keep a Changelog, and the project follows Semantic Versio ### Fixed - **Warm-cache model loading no longer touches the Hugging Face Hub.** Every checkpoint download (`hf_hub_download`/`snapshot_download` across all embedders, plus the rshf `from_pretrained` loaders for satmae/scalemae/satmaepp) now resolves against the local HF cache first and only goes online on a cache miss. Previously each fresh process issued a HEAD request to huggingface.co even when weights were fully cached, so Hub outages, 429 rate limits, or blocked networks froze `get_embedding` indefinitely — hit hardest by agent integrations that spawn a new process per call. Consequence of cache-first: cached weights are never re-checked against the Hub; delete the cached file to force a re-download. Shared helpers: `hf_hub_download_cache_first` / `snapshot_download_cache_first` / `resolve_pretrained_source_cache_first` in `embedders/shared.py`. +- **Clay batch prefetched-input path no longer acquires a provider.** `ClayEmbedder.get_embeddings_batch_from_inputs` unconditionally initialized the provider even though prefetched inputs never fetch — invisible in exports (the provider was already live) but it forced Earth Engine auth on machines without GEE when embedding user-provided data. The single-embedding path already followed the lazy-provider convention; the batch path now matches it, with a regression test. + +### Added + +- **Bring-your-own-data API** — compute embeddings from imagery you already have, without any provider fetch or provider auth. Register each piece of imagery once as `UserData(data, collection, spatial=None, bands=None, temporal=None, scale_m=None)` — the pixels plus everything that describes them: collection, one band name per channel (raw provider units, `[C,H,W]` or `[T,C,H,W]`), and where/when they were acquired — then embed with just a model name: `get_embedding_from_data("galileo", data)` / `get_embeddings_batch_from_data(model, datas, batch_size=None)` (batch items carry their own spatial/temporal; same-temporal items dispatch together, results return in input order; `batch_size` caps the per-forward batch for small GPUs, while each model's per-device internal default still applies as a further cap). User data follows the package-wide `input_prep` policy with the same `"tile"` default as the fetch path: arrays larger than a model's input size are cut into model-native tiles at their own resolution and the outputs stitched, so every model sees the full detail regardless of its input size (galileo's 64 px and clay's 256 px get equal treatment); `input_prep="resize"` opts into one-step downsampling, and arrays a single tile covers keep the efficient batched dispatch. Flexible-size models consume user data natively by default via the new `EmbedderBase.resolve_input_image_size(model_config, input_hw=...)` hook — `olmoearth` (FlexiViT) adapts to the input's own size (snapped up to a patch multiple), so any patch runs as one seamless native pass instead of a tile mosaic, with no user action needed; a native pass beyond 512 px warns (attention cost grows quadratically with token count), and an explicit `image_size` in `model_config` restores fixed-size behavior (larger inputs tile at it). Batch items are grouped by native size so same-size items still dispatch together. The declaration is matched against the model's input sensor: superset band sets are sliced and reordered into model band order automatically, while a collection mismatch or missing band refuses the request with a `ModelError` naming what is missing (precomputed models always refuse). `bands` may be omitted only for the canonical case (12-channel S2 L2A → canonical `B1..B12` order); band identity is never guessed otherwise. `spatial` is optional, but models whose forward pass conditions on geometry (new `_requires_georef` embedder flag: clay, prithvi) refuse declarations without it — coordinates are never fabricated. `list_models_for_data` reports, without loading weights, which catalog models a declaration can serve and why the rest cannot. Matching shares the provider band-alias vocabulary (`"RED"` → `"B4"`, `"NIR_NARROW"` → `"B8A"`), collection shorthand aliases (`"s2"`, `"s1"`) resolve to full ids, and S2 declarations whose values look already normalized (max ≤ 1.5) warn about the raw-DN contract. Results carry `meta["user_input"]` provenance (declared bands, bands used, channel indices). See [docs/user_data.md](docs/user_data.md). ## [0.2.1] — 2026-07-27 diff --git a/docs/models.md b/docs/models.md index 9f61052..24a7d3e 100644 --- a/docs/models.md +++ b/docs/models.md @@ -39,25 +39,27 @@ Some detail-page filenames still use older names for compatibility, but the cano ### On-the-fly Foundation Models -| Model ID | Primary Input | Dim | Default Resolution | Temporal style | Notable requirements | Detail | -| ----------------- | -------------------------------- | ---- | ------------------ | ----------------------- | ------------------------------------------------------- | ------------------------------ | -| `prithvi` | S2 6-band | 768 | 30m | multi-frame (auto, ≤4) | required temporal + location side inputs | [detail](models/prithvi.md) | -| `olmoearth` | S2 L2A 12-band / S1 VV/VH | 128–1024 | 10m | multi-frame (auto, ≤12) | FlexiViT; 4 sizes (nano/tiny/base/large) | [detail](models/olmoearth.md) | -| `dofa` | Multispectral + wavelengths | 768 | 10m | single composite | wavelength vector required | [detail](models/dofa.md) | -| `clay` | S2 L2A 10-band | 1024 | 10m | single composite | metadata conditioning (latlon/time/gsd/wavelengths) | [detail](models/clay.md) | -| `terramind` | S2 12-band | 384 | 10m | single composite | ViT-S class; strict z-score normalization | [detail](models/terramind.md) | -| `terrafm` | S2 12-band or S1 VV/VH | 768 | 10m | single composite | dual-modality by channel count | [detail](models/terrafm.md) | -| `thor` | S2 10-band or S1 VV/VH | 768 | 10m | single composite | dual-modality; grouped tokens; native-snap | [detail](models/thor.md) | -| `galileo` | S2 10-band time series | 128 | 10m | multi-frame (auto, ≤12) | nano default; month tokens | [detail](models/galileo.md) | -| `anysat` | S2 10-band time series | 768 | 10m | multi-frame (fixed `T`) | JEPA; `s2_dates` DOY side input | [detail](models/anysat.md) | -| `agrifm` | S2 10-band time series | 1024 | 10m | multi-frame (fixed `T`) | Video Swin; fixed `T` frame stack | [detail](models/agrifm.md) | -| `fomo` | S2 12-band | 768 | 10m | single composite | per-channel spectral modality keys | [detail](models/fomo.md) | -| `wildsat` | S2 RGB | 256 | 10m | single composite | biodiversity training; image_head default | [detail](models/wildsat.md) | -| `satvision` | TOA 14-channel (MODIS) | 4096 | 1000m | single composite | SwinV2 Giant; strict channel calibration | [detail](models/satvision.md) | -| `remoteclip` | S2 RGB (`B4,B3,B2`) | 512 | 10m | single composite | CLIP projection; RGB preprocessing | [detail](models/remoteclip.md) | -| `scalemae` | S2 RGB + scale | 1024 | 10m | single composite | `sensor.scale_m` is a model input | [detail](models/scalemae.md) | -| `satmae` | S2 RGB (`B4,B3,B2`) | 1024 | 10m | single composite | ViT-L; MAE token/grid | [detail](models/satmae.md) | -| `satmaepp` | S2 RGB (`B4,B3,B2`) or S2 10-band | 1024 | 10m | single composite | `modality=rgb` (default) or `s2_10b`; ViT-L; fMoW eval preprocessing; 10-band uses strict band order + grouped-channel tokens | [detail](models/satmaepp.md) | +| Model ID | Primary Input | Dim | Default Resolution | Input size (px) | Temporal style | Notable requirements | Detail | +| ----------------- | -------------------------------- | ---- | ------------------ | --------------- | ----------------------- | ------------------------------------------------------- | ------------------------------ | +| `prithvi` | S2 6-band | 768 | 30m | 224 | multi-frame (auto, ≤4) | required temporal + location side inputs | [detail](models/prithvi.md) | +| `olmoearth` | S2 L2A 12-band / S1 VV/VH | 128–1024 | 10m | 256 (flexible) | multi-frame (auto, ≤12) | FlexiViT; 4 sizes (nano/tiny/base/large) | [detail](models/olmoearth.md) | +| `dofa` | Multispectral + wavelengths | 768 | 10m | 224 | single composite | wavelength vector required | [detail](models/dofa.md) | +| `clay` | S2 L2A 10-band | 1024 | 10m | 256 | single composite | metadata conditioning (latlon/time/gsd/wavelengths) | [detail](models/clay.md) | +| `terramind` | S2 12-band | 384 | 10m | 224 | single composite | ViT-S class; strict z-score normalization | [detail](models/terramind.md) | +| `terrafm` | S2 12-band or S1 VV/VH | 768 | 10m | 224 | single composite | dual-modality by channel count | [detail](models/terrafm.md) | +| `thor` | S2 10-band or S1 VV/VH | 768 | 10m | 288 | single composite | dual-modality; grouped tokens; native-snap | [detail](models/thor.md) | +| `galileo` | S2 10-band time series | 128 | 10m | 64 | multi-frame (auto, ≤12) | nano default; month tokens | [detail](models/galileo.md) | +| `anysat` | S2 10-band time series | 768 | 10m | 24 | multi-frame (fixed `T`) | JEPA; `s2_dates` DOY side input | [detail](models/anysat.md) | +| `agrifm` | S2 10-band time series | 1024 | 10m | 224 | multi-frame (fixed `T`) | Video Swin; fixed `T` frame stack | [detail](models/agrifm.md) | +| `fomo` | S2 12-band | 768 | 10m | 64 | single composite | per-channel spectral modality keys | [detail](models/fomo.md) | +| `wildsat` | S2 RGB | 256 | 10m | 224 | single composite | biodiversity training; image_head default | [detail](models/wildsat.md) | +| `satvision` | TOA 14-channel (MODIS) | 4096 | 1000m | 128 | single composite | SwinV2 Giant; strict channel calibration | [detail](models/satvision.md) | +| `remoteclip` | S2 RGB (`B4,B3,B2`) | 512 | 10m | 224 | single composite | CLIP projection; RGB preprocessing | [detail](models/remoteclip.md) | +| `scalemae` | S2 RGB + scale | 1024 | 10m | 224 | single composite | `sensor.scale_m` is a model input | [detail](models/scalemae.md) | +| `satmae` | S2 RGB (`B4,B3,B2`) | 1024 | 10m | 224 | single composite | ViT-L; MAE token/grid | [detail](models/satmae.md) | +| `satmaepp` | S2 RGB (`B4,B3,B2`) or S2 10-band | 1024 | 10m | 224 (rgb) / 96 (s2_10b) | single composite | `modality=rgb` (default) or `s2_10b`; ViT-L; fMoW eval preprocessing; 10-band uses strict band order + grouped-channel tokens | [detail](models/satmaepp.md) | + +**Input size (px)** is the fixed spatial size each model's encoder consumes: under the default `input_prep="tile"`, inputs larger than it are cut into tiles of this size at native resolution and the outputs stitched (both for provider fetches and user-provided data — see [User Data API](user_data.md)); under `input_prep="resize"` they are downsampled to it in one step. Together with Default Resolution it gives the native footprint of one forward pass, e.g. galileo 64 px × 10 m ≈ 640 m. `olmoearth` (FlexiViT) accepts any size divisible by its patch size — 256 is its training tile size; user-provided data is consumed natively at its own size by default (warning above 512 px), and an explicit `model_config` `image_size` restores fixed-size behavior. `anysat` and `prithvi` sizes are env-tunable (`RS_EMBED_ANYSAT_IMG`, `RS_EMBED_PRITHVI_IMG`). --- diff --git a/docs/user_data.md b/docs/user_data.md new file mode 100644 index 0000000..ba3c59a --- /dev/null +++ b/docs/user_data.md @@ -0,0 +1,173 @@ +# API: User-Provided Data + +You already have imagery — patches from your own dataset, exported GeoTIFFs, a training cube on disk. This page covers the bring-your-own-data API: computing embeddings from that imagery directly, with no provider fetch and no provider auth. + +Related pages: [API: Embedding](api_embedding.md), [API: Specs and Data Structures](api_specs.md), [Spatial ROI Handling](spatial_roi.md). + +!!! abstract "The one idea" + You **register** each piece of imagery once as a `UserData` — the pixels plus everything that describes them: which collection they came from, one band name per channel, and where/when they were acquired. From then on, embedding takes only a model name. The declaration, not the array shape, is the contract: a model whose bands your declaration covers gets its channels sliced out automatically; a model it cannot satisfy is **refused with the exact reason**, never served silently wrong data. + +--- + +## The two-step flow + +```mermaid +flowchart LR + REG["1. Register\nUserData(pixels, collection,\nbands, spatial, temporal)"] --> MATCH["2. Match\ndeclaration vs the model's\ninput sensor"] + MATCH -->|covers all bands| SLICE["slice + reorder\nchannels to model order"] --> EMB["Embedding"] + MATCH -->|"missing band /\nwrong collection"| REFUSE["ModelError\nnaming what is missing"] +``` + +Every on-the-fly model declares an input sensor (a collection plus an ordered band list), and your declaration is matched against it per request. + +**Superset data is accepted.** If your declaration covers all bands the model needs, the needed channels are sliced out and reordered automatically. One 12-band Sentinel-2 L2A cube serves models that need 3, 6, or 10 of those bands — `galileo` takes its 10, an RGB model takes `B4/B3/B2`, all from the same declaration. Band aliases resolve the same way provider fetches resolve them (`"RED"` → `"B4"`, `"NIR_NARROW"` → `"B8A"`, …). + +**Insufficient data is refused.** A collection mismatch (e.g. S2 data offered to a MODIS model) or a missing band raises `ModelError` naming exactly what is missing. Precomputed models (`tessera`, `gse`, `copernicus`) are always refused — they have no imagery input. Call [`list_models_for_data`](#list_models_for_data) to see the verdict for every catalog model up front. + +!!! warning "Values must be raw provider units" + Pass exactly what a provider fetch would return — for Sentinel-2 L2A that is surface-reflectance DN in `0..10000`, **not** reflectance rescaled to `0..1`. Per-model normalization stays inside each embedder, so you never need to know a model's normalization; but data that looks already normalized (max ≤ 1.5 on an S2 declaration) triggers a `UserWarning`, because it would otherwise produce silently wrong embeddings. + +--- + +## UserData + +```python +from rs_embed import UserData + +UserData( + data: np.ndarray, # [C,H,W] or [T,C,H,W], raw provider values + collection: str, # e.g. "COPERNICUS/S2_SR_HARMONIZED" or alias "s2" + spatial: SpatialSpec | None = None, # where the imagery is (PointBuffer / BBox) + bands: tuple[str, ...] | None = None, # one band name per channel; None = canonical order + temporal: TemporalSpec | None = None, # when the imagery was acquired + scale_m: int | None = None, # optional nominal pixel size, provenance only +) +``` + +**`collection` names the sensor product — and thereby the units.** Short aliases resolve to full ids: `"s2"` / `"sentinel-2"` / `"s2-l2a"` → `COPERNICUS/S2_SR_HARMONIZED`; `"s1"` / `"sentinel-1"` → `COPERNICUS/S1_GRD`. Full collection ids pass through unchanged. + +**`bands` may be omitted only for the canonical case.** An S2 L2A declaration with exactly 12 channels defaults to the canonical order `B1, B2, B3, B4, B5, B6, B7, B8, B8A, B9, B11, B12`. Any other channel count, order, or collection must name its bands — band identity is never guessed from channel count. + +**`spatial` and `temporal` travel with the data**, because they describe *this imagery* (where it is, when it was acquired), not the API call. Both are optional, but supply them whenever you have them: models that condition on time read `temporal`, and models whose forward pass conditions on geometry (lat/lon or GSD encodings: `clay`, `prithvi`) **refuse** declarations without `spatial` — coordinates are never fabricated. Everything else accepts ungeoreferenced data and merely loses location provenance in the metadata. + +!!! note "Multi-frame arrays" + A `[T,C,H,W]` array is only meaningful for time-series models (`galileo`, `prithvi`, `olmoearth`, `anysat`, `agrifm`); single-frame models reject it. + +--- + +## Functions + +### get_embedding_from_data + +A complete example, starting from a file on disk. Say you have a Sentinel-2 L2A patch saved as a GeoTIFF — 12 bands in the canonical order, raw surface-reflectance DN: + +```python +import rasterio # example only; not an rs-embed dependency +from rasterio.warp import transform_bounds + +from rs_embed import UserData, get_embedding_from_data +from rs_embed.core.specs import BBox, TemporalSpec + +# 1. Load the pixels and the footprint from the file. +with rasterio.open("maize_field_2022.tif") as src: + pixels = src.read() # numpy array, shape [12, H, W] + left, bottom, right, top = transform_bounds( # footprint -> lon/lat degrees + src.crs, "EPSG:4326", *src.bounds + ) + +# 2. Register the imagery: the pixels plus everything that describes them. +data = UserData( + data=pixels, + collection="s2", + spatial=BBox(minlon=left, minlat=bottom, maxlon=right, maxlat=top), + temporal=TemporalSpec.range("2022-06-01", "2022-09-01"), +) + +# 3. Embed — just name the model. +emb = get_embedding_from_data("galileo", data) + +print(emb.data.shape) # pooled feature vector, shape [D] +print(emb.meta["user_input"]) # which bands/channels were actually used +``` + +If your data is already a numpy array (e.g. one sample from a training dataset), skip step 1 — anything `[C,H,W]` in raw provider units works as `data=`. + +Returns one `Embedding`. `meta["user_input"]` records the declaration and the channel selection actually fed to the model (`declared_bands`, `bands_used`, `channel_indices`), and `meta["input_prep"]` records how the array was sized (see [Input size handling](#input-size-handling)). + +### get_embeddings_batch_from_data + +Continuing the example above — a whole directory of patches into one feature matrix: + +```python +from pathlib import Path + +import numpy as np + +from rs_embed import get_embeddings_batch_from_data + +datas = [] +for path in sorted(Path("patches/").glob("*.tif")): + with rasterio.open(path) as src: + left, bottom, right, top = transform_bounds(src.crs, "EPSG:4326", *src.bounds) + datas.append( + UserData( + data=src.read(), + collection="s2", + spatial=BBox(minlon=left, minlat=bottom, maxlon=right, maxlat=top), + temporal=TemporalSpec.range("2022-06-01", "2022-09-01"), + ) + ) + +embs = get_embeddings_batch_from_data("galileo", datas, batch_size=16) + +X = np.stack([e.data for e in embs]) # [N, D] — ready for sklearn, clustering, ... +``` + +Each item is matched independently and carries its own `spatial` / `temporal`, so one batch can mix locations, dates, and even band orders. Items sharing a temporal are dispatched together (models with true batching benefit); results always come back in input order. + +!!! tip "Fitting your GPU with `batch_size`" + `batch_size` caps how many items reach one model forward batch — set a small value for a small GPU. Models keep their own per-device internal default (e.g. `clay` 32 on CUDA / 4 on CPU) as a further cap, so `batch_size` lowers but does not raise a model's forward batch; to raise it, use the model's `RS_EMBED__BATCH_SIZE` environment variable. + +### list_models_for_data + +```python +from rs_embed import list_models_for_data + +report = list_models_for_data(data) +[r["model"] for r in report if r["compatible"]] +``` + +Runs the same matching against every catalog model without loading weights. Each entry has `model`, `compatible`, `bands_used`, and `reason` (why the model is incompatible). + +--- + +## Input size handling + +User data follows the package-wide `input_prep` policy, defaulting to **`"tile"`** — the same fairness semantics as the provider-fetch path (see [Spatial ROI Handling](spatial_roi.md) and the *Input size* column in [Models Overview](models.md)). + +**Larger than the model's input size — tiled, not squashed.** The array is cut into model-native tiles at its own resolution, each tile embedded, and the outputs stitched. Every model sees the full detail regardless of its input size: a 256×256 patch reaches `clay` as one 256-px pass and `galileo` as a 4×4 grid of 64-px tiles, instead of `galileo` silently losing 15/16 of the pixels to a resize. Pass `input_prep="resize"` for one-step downsampling instead (faster, lossy); `meta["input_prep"]` records which path ran. + +**At or below the model's input size — straight through.** There is nothing to tile: the array goes to the embedder, which resizes up if needed (`prithvi` can pad instead via `RS_EMBED_PRITHVI_PREP=pad`). Very small arrays run fine but carry only the information your pixels have — expect weak embeddings below roughly half the model's input size. + +**Non-square — tiling handles it cleanly.** Edge tiles are padded and the outputs cropped back. Under `"resize"`, a plain resize distorts the aspect ratio — another reason to keep the tile default for non-square patches. + +!!! note "Flexible-size models run your data natively" + `olmoearth` (FlexiViT) accepts any patch-divisible input size, so your data is consumed **natively by default** — a 512×512 patch runs as one seamless pass instead of a 2×2 tile mosaic, with no configuration needed (the side length is snapped up to a patch multiple). Passing an explicit `image_size` in model kwargs restores fixed-size behavior (larger inputs tile at it). + +!!! warning "Very large native passes" + A flexible-size native pass beyond **512 px** emits a warning: attention cost grows quadratically with token count, so very large patches can be slow or exhaust GPU memory. Tile them via `input_prep=InputPrepSpec(mode="tile", tile_size=...)` or downsample via `"resize"`. + +--- + +## Refusal semantics + +| Situation | Result | +| --------- | ------ | +| Declaration covers all model bands | Accepted; channels sliced/reordered | +| Missing band(s) | `ModelError` listing the missing band names | +| Collection mismatch | `ModelError` (raw units differ across collections) | +| Precomputed model | `ModelError` (no imagery input) | +| Channel count ≠ declared bands | `SpecError` from `UserData.validate()` | +| `bands=None` outside the canonical case | `SpecError` (declare bands explicitly) | +| Missing `spatial` on a georef-conditioned model (`clay`, `prithvi`) | `ModelError` (coordinates are never fabricated) | +| S2 values look normalized (max ≤ 1.5) | `UserWarning`, request still runs | diff --git a/mkdocs.yml b/mkdocs.yml index 70c8efb..9e1b39e 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -106,6 +106,7 @@ nav: - Overview: api.md - Specs & Data Structures: api_specs.md - Embedding API: api_embedding.md + - User Data API: user_data.md - Export API: api_export.md - Load API: api_load.md - Inspect API: api_inspect.md diff --git a/src/rs_embed/__init__.py b/src/rs_embed/__init__.py index 87b742c..76aa0de 100644 --- a/src/rs_embed/__init__.py +++ b/src/rs_embed/__init__.py @@ -12,11 +12,14 @@ describe_model, export_batch, get_embedding, + get_embedding_from_data, get_embeddings_batch, + get_embeddings_batch_from_data, inspect_gee_patch, inspect_model_input, inspect_provider_patch, list_models, + list_models_for_data, reset_runtime, ) from .core._warnings import disable_pretty_warnings, enable_pretty_warnings @@ -35,6 +38,7 @@ ExportModelRequest, ExportTarget, ModelConfig, + UserData, ) from .load import ExportResult, ModelResult, load_export from .model import Model @@ -68,6 +72,11 @@ "list_models", "describe_model", "reset_runtime", + # User-provided data API + "UserData", + "get_embedding_from_data", + "get_embeddings_batch_from_data", + "list_models_for_data", # Export API "export_batch", # Load API diff --git a/src/rs_embed/api.py b/src/rs_embed/api.py index 4881db3..0cea3ea 100644 --- a/src/rs_embed/api.py +++ b/src/rs_embed/api.py @@ -46,6 +46,7 @@ ExportConfig, ExportModelRequest, ExportTarget, + UserData, ) from .core.validation import ( assert_supported as _assert_supported, @@ -59,6 +60,9 @@ from .embedders.catalog import MODEL_ALIASES, MODEL_SPECS from .embedders.meta import temporal_to_range as _temporal_to_range from .providers.fetch import fetch_sensor_patch_chw as _fetch_sensor_patch_chw +from .providers.prefetch_plan import ( + select_prefetched_channels as _select_prefetched_channels, +) from .providers.resolution import create_provider_for_backend from .tools.export_requests import ( maybe_return_completed_combined_resume as _maybe_return_completed_combined_resume, @@ -105,7 +109,22 @@ from .tools.runtime import ( run_embedding_request as _run_embedding_request_shared, ) +from .tools.runtime import ( + run_user_input_request as _run_user_input_request, +) from .tools.tiling import _resolve_input_prep_spec as _resolve_input_prep_spec +from .tools.user_data import ( + match_user_data_to_sensor as _match_user_data_to_sensor, +) +from .tools.user_data import ( + normalize_collection_id as _normalize_collection_id, +) +from .tools.user_data import ( + resolve_declared_bands as _resolve_declared_bands, +) +from .tools.user_data import ( + warn_on_suspicious_value_range as _warn_on_suspicious_value_range, +) # ----------------------------------------------------------------------------- # Internal helpers @@ -398,6 +417,314 @@ def get_embeddings_batch( ) +# ----------------------------------------------------------------------------- +# Public: embeddings from user-provided data +# ----------------------------------------------------------------------------- + + +def _resolve_user_data_sensor(model_n: str, *, modality: str | None) -> SensorSpec: + """Resolve the sensor a user-data request must satisfy, or refuse.""" + embedder_cls = _get_embedder_cls(model_n) + if getattr(embedder_cls, "_is_precomputed", False): + raise ModelError( + f"Model '{model_n}' serves precomputed embeddings; it does not " + "embed user-provided imagery. Use an on-the-fly model instead." + ) + sensor_eff = _resolve_sensor_for_model( + model_n, + sensor=None, + fetch=None, + modality=modality, + default_when_missing=True, + ) + if sensor_eff is None: + raise ModelError( + f"Model '{model_n}' declares no input sensor" + + (f" for modality='{modality}'" if modality else "") + + "; there is nothing to match user-provided data against." + ) + return sensor_eff + + +def _require_georef_for_model(model_n: str, data: UserData) -> None: + """Refuse georef-conditioned models when a declaration has no spatial.""" + if data.spatial is None and getattr(_get_embedder_cls(model_n), "_requires_georef", False): + raise ModelError( + f"Model '{model_n}' conditions on request geometry (lat/lon or " + "GSD encodings); it cannot embed data without UserData.spatial, " + "and coordinates are never fabricated. Provide spatial or choose " + "a model without this requirement (see list_models_for_data)." + ) + + +def _prepare_user_data_inputs( + model_n: str, + datas: list[UserData], + sensor_eff: SensorSpec, + output: OutputSpec, +) -> tuple[list[np.ndarray], list[dict[str, Any]]]: + """Validate, match, and slice each user-data item to model band order.""" + arrays: list[np.ndarray] = [] + metas: list[dict[str, Any]] = [] + for data in datas: + data.validate() + _require_georef_for_model(model_n, data) + if data.spatial is not None: + _validate_specs(spatial=data.spatial, temporal=data.temporal, output=output) + elif data.temporal is not None: + data.temporal.validate() + idx = _match_user_data_to_sensor(data, sensor_eff, model_name=model_n) + _warn_on_suspicious_value_range(data) + # select_prefetched_channels casts to float32 and returns the input + # object unchanged for an identity selection, so no redundant copy. + arrays.append(_select_prefetched_channels(data.data, idx)) + metas.append( + { + "source": "user_data", + "collection": _normalize_collection_id(data.collection), + "declared_bands": list(_resolve_declared_bands(data)), + "bands_used": list(sensor_eff.bands), + "channel_indices": list(idx), + "declared_scale_m": data.scale_m, + } + ) + return arrays, metas + + +def get_embedding_from_data( + model: str, + data: UserData, + *, + modality: str | None = None, + output: OutputSpec = OutputSpec.pooled(), + device: str = "auto", + input_prep: InputPrepSpec | str | None = None, + **model_kwargs: Any, +) -> Embedding: + """Compute an embedding from user-provided imagery (no provider fetch). + + The bring-your-own-data counterpart of :func:`get_embedding`: instead of + fetching imagery from a provider, the caller registers a + :class:`~rs_embed.UserData` describing the imagery completely — pixels, + collection + band names (raw provider units), and where/when it was + acquired — and then only names the model. The declaration is matched + against the model's input sensor: superset band sets are sliced and + reordered automatically, while a collection mismatch or missing band + refuses the request with a :class:`ModelError` naming what is missing. + Call :func:`list_models_for_data` to see which models a declaration can + serve. + + Parameters + ---------- + model : str + Model identifier or alias. Precomputed models are refused (they have + no imagery input). + data : UserData + The registered imagery. ``data.data`` must be ``[C,H,W]`` (or + ``[T,C,H,W]`` for time-series models) with raw provider values; + ``data.spatial`` / ``data.temporal`` carry where and when it was + acquired. + modality : str or None + Optional modality selector for models with multiple input branches; + the declaration is matched against that modality's sensor profile. + output : OutputSpec + Output representation policy. + device : str + Target inference device. + input_prep : InputPrepSpec or str or None + How arrays larger than the model's input size are handled. ``None`` + (the default) uses the package default ``"tile"``: the array is cut + into model-native tiles at its own resolution and the outputs + stitched — every model sees the full detail, regardless of its input + size (galileo's 64 px and clay's 256 px get equal treatment). Pass + ``"resize"`` to instead downsample the whole array to the model input + size in one step. Arrays a single tile already covers are unaffected + either way. + **model_kwargs + Model-specific settings, as in :func:`get_embedding`. + + Returns + ------- + Embedding + Normalized embedding; ``meta['user_input']`` records the declaration + and the channel selection that was fed to the model, and + ``meta['input_prep']`` records how the array was prepared. + + Raises + ------ + ModelError + If the model cannot take user data (precomputed / no input sensor) or + the declaration does not satisfy the model's sensor. + SpecError + If the declaration fails validation. + + Examples + -------- + >>> pixels = np.load("s2_patch.npy") # [12, H, W] raw S2 L2A DN (0..10000) + >>> data = UserData( + ... data=pixels, + ... collection="s2", # 12 channels in canonical order -> bands may be omitted + ... spatial=PointBuffer(lon=-88.2, lat=40.1, buffer_m=640), + ... temporal=TemporalSpec.year(2022), + ... ) + >>> emb = get_embedding_from_data("galileo", data) + """ + return get_embeddings_batch_from_data( + model, + [data], + modality=modality, + output=output, + device=device, + input_prep=input_prep, + **model_kwargs, + )[0] + + +def get_embeddings_batch_from_data( + model: str, + datas: list[UserData], + *, + modality: str | None = None, + output: OutputSpec = OutputSpec.pooled(), + device: str = "auto", + batch_size: int | None = None, + input_prep: InputPrepSpec | str | None = None, + **model_kwargs: Any, +) -> list[Embedding]: + """Compute embeddings for multiple user-provided inputs. + + Batch counterpart of :func:`get_embedding_from_data`; each item is + matched independently, so items may declare different band orders and + carry their own ``spatial`` / ``temporal``. Items sharing a temporal are + dispatched together (models with true batching benefit); results always + come back in input order. + + Parameters + ---------- + model : str + Model identifier or alias. + datas : list[UserData] + Registered imagery, one per item. + modality : str or None + Optional modality selector. + output : OutputSpec + Output representation policy. + device : str + Target inference device. + batch_size : int or None + Upper bound on how many items reach one model forward batch — use a + small value to fit a small GPU. Models keep their own per-device + internal default as a further cap, so this lowers but does not raise + a model's forward batch (raise via the model's + ``RS_EMBED__BATCH_SIZE`` environment variable). Items large + enough to be tiled are processed one at a time (their tiles are + batched by the model internally). + input_prep : InputPrepSpec or str or None + How arrays larger than the model's input size are handled, as in + :func:`get_embedding_from_data`: ``None`` uses the package default + ``"tile"`` (native-resolution tiles + stitched outputs), ``"resize"`` + downsamples instead. + **model_kwargs + Model-specific settings, as in :func:`get_embedding`. + + Returns + ------- + list[Embedding] + Embeddings in the same order as *datas*. + + Raises + ------ + ModelError + If *datas* is empty, ``batch_size`` is invalid, the model cannot + take user data, or any declaration does not satisfy the model's + sensor. + SpecError + If any declaration fails validation. + """ + model_config = model_kwargs or None + model_n = _normalize_model_name(model) + if not isinstance(datas, list) or len(datas) == 0: + raise ModelError("datas must be a non-empty list[UserData].") + sensor_eff = _resolve_user_data_sensor(model_n, modality=modality) + arrays, metas = _prepare_user_data_inputs(model_n, datas, sensor_eff, output) + + # The embedder batch path takes one temporal per call, so dispatch one + # sub-request per distinct temporal and scatter results back to input + # order. TemporalSpec is frozen/hashable; dict preserves insertion order. + groups: dict[TemporalSpec | None, list[int]] = {} + for i, data in enumerate(datas): + groups.setdefault(data.temporal, []).append(i) + + results: list[Embedding | None] = [None] * len(datas) + for temporal, indices in groups.items(): + embs = _run_user_input_request( + model_n=model_n, + spatials=[datas[i].spatial for i in indices], + input_arrays=[arrays[i] for i in indices], + temporal=temporal, + sensor=sensor_eff, + model_config=model_config, + output=output, + device=device, + input_metas=[metas[i] for i in indices], + batch_size=batch_size, + input_prep=input_prep, + ) + for i, emb in zip(indices, embs, strict=True): + results[i] = emb + out = [emb for emb in results if emb is not None] + if len(out) != len(datas): + raise ModelError("Internal error: temporal-grouped dispatch lost items.") + return out + + +def list_models_for_data(data: UserData) -> list[dict[str, Any]]: + """Report which catalog models a user-data declaration can serve. + + Runs the same matching as :func:`get_embedding_from_data` against every + model in the catalog without loading any weights. + + Parameters + ---------- + data : UserData + The imagery declaration to check. ``data.data`` may be a small dummy + array; only its shape is validated here. + + Returns + ------- + list[dict[str, Any]] + One entry per catalog model, sorted by name, with keys ``model`` + (str), ``compatible`` (bool), ``bands_used`` (the model's band order, + or ``None``), and ``reason`` (why the model is incompatible, or + ``None``). + + Examples + -------- + >>> report = list_models_for_data(my_s2_declaration) + >>> [r["model"] for r in report if r["compatible"]] + """ + data.validate() + report: list[dict[str, Any]] = [] + for model_id in sorted(MODEL_SPECS.keys()): + entry: dict[str, Any] = { + "model": model_id, + "compatible": False, + "bands_used": None, + "reason": None, + } + try: + sensor = _resolve_user_data_sensor(model_id, modality=None) + _require_georef_for_model(model_id, data) + _match_user_data_to_sensor(data, sensor, model_name=model_id) + except ModelError as exc: + entry["reason"] = str(exc) + else: + entry["compatible"] = True + entry["bands_used"] = list(sensor.bands) + report.append(entry) + return report + + # ----------------------------------------------------------------------------- # Public: batch export (core) # ----------------------------------------------------------------------------- diff --git a/src/rs_embed/core/types.py b/src/rs_embed/core/types.py index d5e1457..17ad113 100644 --- a/src/rs_embed/core/types.py +++ b/src/rs_embed/core/types.py @@ -12,7 +12,8 @@ import numpy as np -from .specs import FetchSpec, InputPrepSpec, SensorSpec +from .errors import SpecError +from .specs import FetchSpec, InputPrepSpec, SensorSpec, SpatialSpec, TemporalSpec # ── Enums ────────────────────────────────────────────────────────── @@ -143,6 +144,95 @@ class FetchResult: meta: dict[str, Any] +@dataclass(frozen=True) +class UserData: + """User-provided imagery with a declaration of what the pixels are. + + The full description of one piece of user imagery: the pixels, what + sensor they came from, and where/when they were acquired. Register once, + then any bring-your-own-data entrypoint + (:func:`rs_embed.get_embedding_from_data` and friends) needs only the + model name. The declaration is matched against the model's input sensor + and the request is refused when the data cannot satisfy it — so the + declaration, not the array shape, is the contract. + + Values must be raw provider units for *collection* (e.g. Sentinel-2 L2A + surface-reflectance DN in ``0..10000``), exactly what a provider fetch + would return; per-model normalization stays the embedder's job. + + Attributes + ---------- + data : np.ndarray + Pixel array, ``[C,H,W]`` or ``[T,C,H,W]`` with channels in *bands* + order. Multi-frame arrays are only meaningful for time-series models; + single-frame models reject them. + collection : str + Provider collection the pixels came from, e.g. + ``"COPERNICUS/S2_SR_HARMONIZED"`` or a short alias like ``"s2"``. + spatial : SpatialSpec or None + Where the imagery is located. Optional, but models whose forward pass + conditions on geometry (e.g. lat/lon or GSD encodings — clay, + prithvi) refuse declarations without it; coordinates are never + fabricated. Supply it whenever you have it. + bands : tuple[str, ...] or None + Band name per channel, e.g. ``("B2", "B3", "B4", ...)``. Aliases like + ``"RED"`` resolve the same way provider fetches resolve them. + ``None`` means "the collection's canonical full band order"; this is + only accepted for collections with a documented canonical order (e.g. + the 12-band Sentinel-2 L2A set) and when the channel count matches — + otherwise the declaration must name its bands. + temporal : TemporalSpec or None + When the imagery was acquired, for models that condition on time. + scale_m : int or None + Optional nominal pixel size in meters, recorded as provenance. + """ + + data: np.ndarray + collection: str + spatial: SpatialSpec | None = None + bands: tuple[str, ...] | None = None + temporal: TemporalSpec | None = None + scale_m: int | None = None + + def validate(self) -> None: + """Validate the declaration's internal consistency. + + Band identity against a defaulted (``bands=None``) declaration is + resolved later by the matching layer; this only checks what the + object can know about itself. + + Raises + ------ + SpecError + If the collection/bands declaration is empty or malformed, or the + array is not CHW/TCHW with one channel per declared band. + """ + if not str(self.collection or "").strip(): + raise SpecError("UserData.collection must be a non-empty collection id or alias.") + if self.bands is not None and ( + len(self.bands) == 0 or any(not str(b or "").strip() for b in self.bands) + ): + raise SpecError( + "UserData.bands must be a non-empty tuple of band names, or " + "None for the collection's canonical band order." + ) + arr = np.asarray(self.data) + if arr.ndim not in (3, 4): + raise SpecError( + "UserData.data must be [C,H,W] or [T,C,H,W], " + f"got shape={tuple(int(v) for v in arr.shape)}." + ) + if self.bands is not None: + channels = int(arr.shape[-3]) + if channels != len(self.bands): + raise SpecError( + f"UserData.data has {channels} channels but declares " + f"{len(self.bands)} bands; one band name per channel is required." + ) + if self.scale_m is not None and int(self.scale_m) <= 0: + raise SpecError("UserData.scale_m must be positive when provided.") + + # ── Typed results ────────────────────────────────────────────────── diff --git a/src/rs_embed/embedders/base.py b/src/rs_embed/embedders/base.py index bcd5c39..69d8b06 100644 --- a/src/rs_embed/embedders/base.py +++ b/src/rs_embed/embedders/base.py @@ -53,6 +53,10 @@ class EmbedderBase: # The embedder performs its own spatial tiling based on request size; # API-side ``input_prep`` has no effect and a non-resize request warns. _manages_own_input_prep: bool = False + # The forward pass conditions on request geometry (e.g. lat/lon or GSD + # encodings derived from ``spatial``); requests without a spatial (such as + # ungeoreferenced user-provided data) must be refused, never fabricated. + _requires_georef: bool = False def __init__(self) -> None: self._providers: dict[str, ProviderBase] = {} @@ -197,6 +201,25 @@ def fetch_input( return FetchResult(data=raw, meta=roi_fetch_meta(geo_roi) or {}) + def resolve_input_image_size( + self, + model_config: dict[str, Any] | None, + *, + input_hw: tuple[int, int] | None = None, + ) -> int | None: + """The input size the encoder will consume for this request, or ``None``. + + Generic layers (the tiler's tile-size resolution and the user-data + routing) consult this before falling back to the static + ``describe().defaults.image_size``. Flexible-size embedders + (FlexiViT-style models) override it: an explicitly configured + ``model_config`` size wins, and otherwise — when *input_hw* is given — + they adapt to the input itself, so any input runs as one native pass + instead of being tiled at the training tile size. The default returns + ``None`` (fixed-size model; use the static default). + """ + return None + def tiled_dispatch_model_config( self, model_config: dict[str, Any] | None, diff --git a/src/rs_embed/embedders/onthefly_clay.py b/src/rs_embed/embedders/onthefly_clay.py index 9c8ebe9..8dff0bb 100644 --- a/src/rs_embed/embedders/onthefly_clay.py +++ b/src/rs_embed/embedders/onthefly_clay.py @@ -551,6 +551,9 @@ class ClayEmbedder(EmbedderBase): # Clay needs a square token grid → base.fetch_input enlarges a rectangular # ROI to a square of real imagery; the output is cropped back to the ROI. _requires_square_input = True + # The encoder conditions on lat/lon (and GSD) metadata derived from the + # request geometry, so a request without a spatial must be refused. + _requires_georef = True DEFAULT_FETCH_WORKERS = 8 DEFAULT_BATCH_CPU = 4 DEFAULT_BATCH_CUDA = 32 @@ -921,7 +924,8 @@ def get_embeddings_batch_from_inputs( backend=backend, device=device, ) - self._get_provider(backend) + # Inputs are prefetched — same lazy-provider convention as the single + # path: never acquire a provider (and its live auth) without a fetch. t = temporal_to_range(temporal) ss = sensor or self._default_sensor() diff --git a/src/rs_embed/embedders/onthefly_olmoearth.py b/src/rs_embed/embedders/onthefly_olmoearth.py index 40d7304..9c00b04 100644 --- a/src/rs_embed/embedders/onthefly_olmoearth.py +++ b/src/rs_embed/embedders/onthefly_olmoearth.py @@ -779,6 +779,41 @@ class OlmoEarthEmbedder(EmbedderBase): model_config_batch_inputs=True, ) + def resolve_input_image_size( + self, + model_config: dict[str, Any] | None, + *, + input_hw: tuple[int, int] | None = None, + ) -> int: + """FlexiViT accepts any patch-divisible size; adapt to the input. + + An explicitly configured size (``model_config['image_size']`` or the + env override) wins and behaves like a fixed-size model (larger inputs + tile at it). Otherwise, when the input shape is known, the encoder + consumes it natively — the longer side snapped up to a patch + multiple — so user-provided data of any size runs as one seamless + pass instead of being tiled at the 256-px training tile size. + """ + explicit = model_config_value(model_config, "image_size") is not None or bool( + os.environ.get("RS_EMBED_OLMOEARTH_IMAGE_SIZE", "").strip() + ) + if explicit or input_hw is None: + return _resolve_geometry(model_config)[0] + patch_size = _resolve_patch_size(model_config) + side = max(int(input_hw[0]), int(input_hw[1])) + return max(patch_size, -(-side // patch_size) * patch_size) + + def tiled_dispatch_model_config( + self, + model_config: dict[str, Any] | None, + *, + tile_size: int, + ) -> dict[str, Any] | None: + """Consume dispatched inputs at exactly *tile_size* (no internal resize).""" + out = dict(model_config or {}) + out["image_size"] = int(tile_size) + return out + def describe(self) -> dict[str, Any]: return { "type": "onthefly", diff --git a/src/rs_embed/embedders/onthefly_prithvi.py b/src/rs_embed/embedders/onthefly_prithvi.py index ca39d01..7cd3ca4 100644 --- a/src/rs_embed/embedders/onthefly_prithvi.py +++ b/src/rs_embed/embedders/onthefly_prithvi.py @@ -911,6 +911,9 @@ class PrithviEOV2S2_6B_Embedder(EmbedderBase): # ROI to a square of real imagery (base.fetch_input for the single-frame # path; the multi-frame fetch_input below for multi) and crop back to the ROI. _requires_square_input = True + # The TL checkpoints condition on location coords derived from the request + # geometry, so a request without a spatial must be refused. + _requires_georef = True DEFAULT_MODEL_KEY = "prithvi_eo_v2_100_tl" DEFAULT_IMAGE_SIZE = 224 DEFAULT_IMAGE_SCALE_M = 30 # notebook used 30m diff --git a/src/rs_embed/pipelines/inference.py b/src/rs_embed/pipelines/inference.py index 13a562f..56ac271 100644 --- a/src/rs_embed/pipelines/inference.py +++ b/src/rs_embed/pipelines/inference.py @@ -438,7 +438,7 @@ def _run_batch_tiled( spec = input_prep_resolved if input_prep_resolved is not None else self.input_prep_resolved out: dict[int, TaskResult] = {} try: - params = _resolve_tile_params(embedder, spec) + params = _resolve_tile_params(embedder, spec, model_config) tile_size = params.tile_size if tile_size <= 0: return out, False diff --git a/src/rs_embed/providers/gee_utils.py b/src/rs_embed/providers/gee_utils.py index e14228a..35cb90c 100644 --- a/src/rs_embed/providers/gee_utils.py +++ b/src/rs_embed/providers/gee_utils.py @@ -269,7 +269,12 @@ def _gee_error_message(exc: Exception) -> str: # ── Band alias resolution ───────────────────────────────────────────────────── -def _resolve_band_aliases(collection: str, bands: tuple[str, ...]) -> tuple[str, ...]: +def resolve_band_aliases(collection: str, bands: tuple[str, ...]) -> tuple[str, ...]: + """Resolve human-friendly band aliases to collection-native band names. + + Pure lookup over static alias tables (no provider access); unknown names + and unknown collections pass through unchanged. + """ if not bands: return bands c = (collection or "").upper() @@ -288,6 +293,10 @@ def _resolve_band_aliases(collection: str, bands: tuple[str, ...]) -> tuple[str, return tuple(amap.get((b or "").upper(), b) for b in bands) +# Backwards-compatible alias kept for existing imports/tests. +_resolve_band_aliases = resolve_band_aliases + + # ── Cloud-cover filtering ───────────────────────────────────────────────────── # Scene-level cloud-cover property per collection family. GEE property filters diff --git a/src/rs_embed/tools/runtime.py b/src/rs_embed/tools/runtime.py index 03a0ac4..58b0032 100644 --- a/src/rs_embed/tools/runtime.py +++ b/src/rs_embed/tools/runtime.py @@ -41,6 +41,11 @@ _T = TypeVar("_T") +# A flexible-size model consuming a user input natively beyond this side +# length gets a UserWarning: ViT attention cost grows quadratically with +# token count, so very large native passes can be slow or exhaust GPU memory. +_FLEX_NATIVE_WARN_PX = 512 + def resolve_device_auto_torch(device: str) -> str: if device != "auto": @@ -802,3 +807,221 @@ def run_embedding_request( ctx=ctx, output=output, ) + + +def run_user_input_request( + *, + model_n: str, + spatials: list[SpatialSpec], + input_arrays: list[np.ndarray], + temporal: TemporalSpec | None, + sensor: SensorSpec, + model_config: dict[str, Any] | None, + output: OutputSpec, + device: str, + input_metas: list[dict[str, Any]] | None = None, + batch_size: int | None = None, + input_prep: Any | None = None, +) -> list[Embedding]: + """Run an embedding request over user-provided inputs (no provider fetch). + + *input_arrays* must already be matched and sliced to the model's band + order for *sensor* (see ``tools.user_data.match_user_data_to_sensor``); + this function only dispatches them through the embedder's prefetched-input + path. ``backend="auto"`` is passed through so the cached embedder instance + is shared with the default fetch path, but with an input array present no + embedder resolves a provider, so no provider auth is required. + + Input prep follows the package default: under ``"tile"`` (the default), + arrays larger than the model's tile size go through the shared tiler + (native-resolution tiles + stitched grids — the same fairness semantics as + the fetch path), while arrays that cannot tile keep the efficient batched + dispatch. ``"resize"`` sends every array straight to the embedder, which + downsamples to its fixed input size. + + Parameters + ---------- + model_n : str + Canonical model name. + spatials : list[SpatialSpec] + Spatial metadata per input (location context for models that condition + on geometry, and for embedding meta). + input_arrays : list[np.ndarray] + Model-band-order arrays aligned with *spatials*. + temporal : TemporalSpec or None + Temporal metadata for models that condition on time. + sensor : SensorSpec + The model's resolved input sensor (band order authority). + model_config : dict or None + Optional model-specific settings. + output : OutputSpec + Requested output layout. + device : str + Target inference device. + input_metas : list of dict or None + Optional per-item provenance recorded as ``meta['user_input']``. + batch_size : int or None + Upper bound on how many items reach one embedder batch call: the + dispatch is chunked to this size. Embedders with a smaller internal + per-device forward batch still split further, so this caps GPU memory + but cannot raise a model's own default. ``None`` dispatches all items + in one call (the embedder's internal chunking applies). Items routed + through the tiler are processed one at a time regardless (their tiles + are batched by the embedder internally). + input_prep : InputPrepSpec or str or None + API-side input preprocessing policy; ``None`` uses the package + default ``"tile"``. See the function description. + + Returns + ------- + list[Embedding] + Embeddings aligned with *spatials*. + + Raises + ------ + ModelError + If lengths mismatch, ``batch_size`` is invalid, or the embedder + cannot take prefetched inputs. + """ + if len(spatials) != len(input_arrays): + raise ModelError( + f"spatials/input arrays length mismatch: {len(spatials)} != {len(input_arrays)}" + ) + if batch_size is not None and int(batch_size) < 1: + raise ModelError(f"batch_size must be >= 1, got {batch_size}.") + device_n = normalize_device_name(device) + embedder, lock = get_embedder_bundle_cached(model_n, "auto", device_n) + if not embedder_accepts_input_chw(type(embedder)): + raise ModelError( + f"Model '{model_n}' does not accept prefetched inputs (input_chw); " + "it cannot embed user-provided data." + ) + assert_supported(embedder, backend="auto", output=output, temporal=temporal) + + if model_config is not None: + require_model_config_support( + embedder=embedder, + model_config=model_config, + method_name="get_embeddings_batch_from_inputs", + ) + + from .tiling import ( + _call_embedder_get_embedding_with_input_prep, + _resolve_tile_params, + _stamp_input_prep_meta, + ) + + ( + input_prep_eff, + input_prep_resolved, + requested_mode, + _model_policy, + ) = resolve_model_aware_input_prep(model_n=model_n, input_prep=input_prep, output=output) + + # Route per item. Fixed-size models: arrays larger than the model's tile + # size go through the shared tiler (native resolution preserved via + # tiles), everything else keeps the efficient batched dispatch. + # Flexible-size models (resolve_input_image_size adapts to input_hw) + # consume any array as one seamless native pass instead — no tiling, with + # a warning above _FLEX_NATIVE_WARN_PX since attention cost grows + # quadratically with token count. An explicitly configured size restores + # fixed-size behavior. + tiled_set: set[int] = set() + native_sizes: dict[int, int | None] = {} + if str(getattr(input_prep_resolved, "mode", "tile")) != "resize": + params = _resolve_tile_params(embedder, input_prep_resolved, model_config) + for i, x in enumerate(input_arrays): + h, w = int(x.shape[-2]), int(x.shape[-1]) + native: int | None = None + try: + v = embedder.resolve_input_image_size(model_config, input_hw=(h, w)) + native = int(v) if v is not None and int(v) > 0 else None + except ModelError: + raise + except Exception as _e: + native = None + threshold = native if native is not None else params.tile_size + if threshold > 0 and (h > threshold or w > threshold): + tiled_set.add(i) + else: + native_sizes[i] = native + else: + native_sizes = dict.fromkeys(range(len(input_arrays))) + + max_native = max((v for v in native_sizes.values() if v is not None), default=0) + if max_native > _FLEX_NATIVE_WARN_PX: + warnings.warn( + f"Model '{model_n}' will consume a {max_native}px input as one " + f"native pass (> {_FLEX_NATIVE_WARN_PX}px). Attention cost grows " + "quadratically with token count; expect high GPU memory/time. " + "Pass input_prep=InputPrepSpec(mode='tile', tile_size=...) to " + "tile instead, or 'resize' to downsample.", + UserWarning, + stacklevel=3, + ) + + results: list[Embedding | None] = [None] * len(input_arrays) + + # Direct items grouped by native size: one embedder batch call per group, + # with the size injected via tiled_dispatch_model_config so a flexible + # model consumes the group's inputs at exactly that size. + direct_groups: dict[int | None, list[int]] = {} + for i in sorted(native_sizes): + direct_groups.setdefault(native_sizes[i], []).append(i) + for native, indices in direct_groups.items(): + group_config = model_config + if native is not None: + group_config = embedder.tiled_dispatch_model_config(model_config, tile_size=native) + if group_config is not None and group_config is not model_config: + require_model_config_support( + embedder=embedder, + model_config=group_config, + method_name="get_embeddings_batch_from_inputs", + ) + step = int(batch_size) if batch_size is not None else max(1, len(indices)) + for s0 in range(0, len(indices), step): + chunk = indices[s0 : s0 + step] + kwargs: dict[str, Any] = { + "spatials": [spatials[i] for i in chunk], + "input_chws": [input_arrays[i] for i in chunk], + "temporal": temporal, + "sensor": sensor, + "output": output, + "backend": "auto", + "device": device_n, + } + if group_config is not None: + kwargs["model_config"] = group_config + with lock: + out = embedder.get_embeddings_batch_from_inputs(**kwargs) + for i, emb in zip(chunk, out, strict=True): + emb = normalize_embedding_output(emb=emb, output=output) + results[i] = _stamp_input_prep_meta( + emb, + requested_mode=requested_mode, + resolved_mode="native" if native is not None else "resize", + ) + + for i in sorted(tiled_set): + with lock: + results[i] = _call_embedder_get_embedding_with_input_prep( + embedder=embedder, + spatial=spatials[i], + temporal=temporal, + sensor=sensor, + output=output, + backend="auto", + device=device_n, + input_chw=input_arrays[i], + input_prep=input_prep_eff, + model_config=model_config, + ) + + embs = [emb for emb in results if emb is not None] + if len(embs) != len(input_arrays): + raise ModelError("Internal error: user-input dispatch lost items.") + if input_metas is not None: + for emb, item_meta in zip(embs, input_metas, strict=False): + if item_meta and isinstance(getattr(emb, "meta", None), dict): + emb.meta.setdefault("user_input", item_meta) + return embs diff --git a/src/rs_embed/tools/tiling.py b/src/rs_embed/tools/tiling.py index 17419a8..69788e4 100644 --- a/src/rs_embed/tools/tiling.py +++ b/src/rs_embed/tools/tiling.py @@ -230,13 +230,16 @@ class _TileParams: def _resolve_tile_params( embedder: Any, input_prep: _ResolvedInputPrepSpec, + model_config: dict[str, Any] | None = None, ) -> _TileParams: """Resolve tile size/stride and the padding policy for an embedder + spec. - ``tile_size`` is the explicit ``input_prep.tile_size`` or, failing that, the - model's advertised ``describe().defaults.image_size``; ``0`` signals that no - tile size could be determined and the caller should fall back to a plain - call. + ``tile_size`` is the explicit ``input_prep.tile_size``, else the size the + embedder resolves for this request via ``resolve_input_image_size( + model_config)`` (how flexible-size models raise their tiling threshold), + else the model's advertised ``describe().defaults.image_size``; ``0`` + signals that no tile size could be determined and the caller should fall + back to a plain call. Edge tiles are padded to square whenever ``pad_edges`` is on (the default), for every model. The tiling layer must never hand a model a rectangular @@ -249,7 +252,18 @@ def _resolve_tile_params( escape hatch for models that natively handle rectangular inputs with a proportional output grid. """ - model_img = _embedder_default_image_size(embedder) + model_img: int | None = None + hook = getattr(embedder, "resolve_input_image_size", None) + if callable(hook): + try: + v = hook(model_config) + model_img = int(v) if v is not None and int(v) > 0 else None + except ModelError: + raise + except Exception as _e: + model_img = None + if model_img is None: + model_img = _embedder_default_image_size(embedder) tile_size = int(input_prep.tile_size or model_img or 0) stride = int(input_prep.tile_stride or tile_size) return _TileParams( @@ -806,7 +820,7 @@ def _call_embedder_get_embedding_tiled( fetch_meta: dict[str, Any] | None = None, ) -> Embedding: x = np.asarray(input_chw, dtype=np.float32) - params = _resolve_tile_params(embedder, input_prep) + params = _resolve_tile_params(embedder, input_prep, model_config) tile_size = params.tile_size model_fixed_size = params.model_fixed_size effective_pad_edges = params.effective_pad_edges diff --git a/src/rs_embed/tools/user_data.py b/src/rs_embed/tools/user_data.py new file mode 100644 index 0000000..5062a06 --- /dev/null +++ b/src/rs_embed/tools/user_data.py @@ -0,0 +1,202 @@ +"""Match user-provided imagery declarations against model input sensors. + +The bring-your-own-data entrypoints let callers compute embeddings from +arrays they already have instead of provider-fetched imagery. The caller +declares what the array is (:class:`~rs_embed.core.types.UserData`: +collection + band names, raw provider units); this module decides whether +that declaration satisfies a model's resolved :class:`SensorSpec` and, when +it does, which user channels to feed the model in which order. + +Policy: the model's required bands must be a subset of the declared bands +(superset data is sliced and reordered automatically); a collection mismatch +or a missing band refuses the request with a :class:`ModelError` naming what +is missing. Band vocabulary is shared with provider fetches via +:func:`~rs_embed.providers.gee_utils.resolve_band_aliases`. +""" + +from __future__ import annotations + +import warnings + +import numpy as np + +from ..core.errors import ModelError, SpecError +from ..core.specs import SensorSpec +from ..core.types import UserData +from ..providers.gee_utils import resolve_band_aliases + +# Short user-facing aliases for provider collection ids. Full ids always pass +# through unchanged, so this stays a convenience layer, not a registry. +_COLLECTION_ALIASES: dict[str, str] = { + "s2": "COPERNICUS/S2_SR_HARMONIZED", + "s2_sr": "COPERNICUS/S2_SR_HARMONIZED", + "s2_l2a": "COPERNICUS/S2_SR_HARMONIZED", + "sentinel2": "COPERNICUS/S2_SR_HARMONIZED", + "sentinel_2": "COPERNICUS/S2_SR_HARMONIZED", + "s1": "COPERNICUS/S1_GRD", + "s1_grd": "COPERNICUS/S1_GRD", + "sentinel1": "COPERNICUS/S1_GRD", + "sentinel_1": "COPERNICUS/S1_GRD", +} + +# Collections whose raw units are surface-reflectance DN in 0..10000; used +# only for the best-effort "looks already normalized" warning below. +_DN_0_10000_COLLECTION_MARKERS: tuple[str, ...] = ("COPERNICUS/S2",) + +# Canonical full band order per collection, used when a declaration omits +# ``bands``. Only collections with one unambiguous canonical order belong +# here; a declaration whose channel order differs must name its bands. +_DEFAULT_BANDS_BY_COLLECTION: dict[str, tuple[str, ...]] = { + "COPERNICUS/S2_SR_HARMONIZED": ( + "B1", + "B2", + "B3", + "B4", + "B5", + "B6", + "B7", + "B8", + "B8A", + "B9", + "B11", + "B12", + ), +} + + +def normalize_collection_id(collection: str) -> str: + """Resolve a user-facing collection alias to a full collection id.""" + raw = str(collection or "").strip() + key = raw.lower().replace("-", "_").replace(" ", "_") + return _COLLECTION_ALIASES.get(key, raw) + + +def canonical_band_names(collection_id: str, bands: tuple[str, ...]) -> tuple[str, ...]: + """Alias-resolve band names and fold case for comparison.""" + resolved = resolve_band_aliases(collection_id, tuple(str(b) for b in bands)) + return tuple(b.upper() for b in resolved) + + +def resolve_declared_bands(data: UserData) -> tuple[str, ...]: + """Return the band names a declaration covers, defaulting when omitted. + + A ``bands=None`` declaration means "the collection's canonical full band + order"; that default only exists for collections listed in + ``_DEFAULT_BANDS_BY_COLLECTION`` and only when the array's channel count + matches exactly — anything else must name its bands, because guessing + band identity from channel count is precisely the silent-wrongness this + layer exists to refuse. + + Raises + ------ + SpecError + If ``bands`` is omitted and the collection has no canonical order, + or the channel count does not match that order. + """ + if data.bands is not None: + return tuple(data.bands) + collection_id = normalize_collection_id(data.collection) + default = _DEFAULT_BANDS_BY_COLLECTION.get(collection_id.upper()) + if default is None: + raise SpecError( + f"UserData.bands was omitted, but collection '{data.collection}' " + "has no canonical band order to default to; declare one band " + "name per channel." + ) + channels = int(np.asarray(data.data).shape[-3]) + if channels != len(default): + raise SpecError( + f"UserData.bands was omitted; the canonical order for " + f"'{collection_id}' has {len(default)} bands {list(default)}, but " + f"the array has {channels} channels. Declare bands explicitly." + ) + return default + + +def match_user_data_to_sensor( + data: UserData, + sensor: SensorSpec, + *, + model_name: str, +) -> tuple[int, ...]: + """Match a user-data declaration against a model's input sensor. + + Parameters + ---------- + data : UserData + User declaration (collection, bands, array). Must already be + validated via :meth:`UserData.validate`. + sensor : SensorSpec + The model's resolved input sensor to satisfy. + model_name : str + Model name used in refusal messages. + + Returns + ------- + tuple[int, ...] + Channel indices into the user array's band axis, in the model's band + order, suitable for + :func:`~rs_embed.providers.prefetch_plan.select_prefetched_channels`. + + Raises + ------ + ModelError + If the declared collection does not match the model's, the + declaration repeats a band name, or a required band is missing. + SpecError + If ``bands`` was omitted and no canonical default applies (see + :func:`resolve_declared_bands`). + """ + user_collection = normalize_collection_id(data.collection) + model_collection = normalize_collection_id(sensor.collection) + if user_collection.upper() != model_collection.upper(): + raise ModelError( + f"Model '{model_name}' expects imagery from collection " + f"'{sensor.collection}', but the provided data is declared as " + f"'{data.collection}'. Raw units differ across collections, so " + "this data cannot serve the model." + ) + + user_bands = canonical_band_names(user_collection, resolve_declared_bands(data)) + if len(set(user_bands)) != len(user_bands): + dupes = sorted({b for b in user_bands if user_bands.count(b) > 1}) + raise ModelError( + f"UserData.bands declares duplicate band name(s) {dupes}; " + "channel selection would be ambiguous." + ) + model_bands = canonical_band_names(model_collection, tuple(sensor.bands)) + + missing = [b for b in model_bands if b not in user_bands] + if missing: + raise ModelError( + f"Model '{model_name}' needs bands {list(model_bands)} from " + f"'{sensor.collection}', but the provided data lacks {missing} " + f"(declared bands: {list(user_bands)})." + ) + return tuple(user_bands.index(b) for b in model_bands) + + +def warn_on_suspicious_value_range(data: UserData) -> None: + """Warn when values look already normalized for a raw-DN collection. + + The user-data contract expects raw provider units; reflectance already + scaled to ``0..1`` fed into a DN-normalizing embedder produces silently + wrong embeddings, which this best-effort check surfaces early. + """ + collection_id = normalize_collection_id(data.collection).upper() + if not any(marker in collection_id for marker in _DN_0_10000_COLLECTION_MARKERS): + return + arr = np.asarray(data.data) + finite = arr[np.isfinite(arr)] + if finite.size == 0: + return + max_value = float(finite.max()) + if 0.0 < max_value <= 1.5: + warnings.warn( + f"UserData declared as '{data.collection}' has max value " + f"{max_value:.3g}; raw surface-reflectance DN (0..10000) is " + "expected. If your data is scaled reflectance, multiply by 10000 " + "before embedding.", + UserWarning, + stacklevel=3, + ) diff --git a/tests/test_api_from_data.py b/tests/test_api_from_data.py new file mode 100644 index 0000000..5f30f4e --- /dev/null +++ b/tests/test_api_from_data.py @@ -0,0 +1,423 @@ +"""Tests for the bring-your-own-data API (get_embedding_from_data & friends). + +These use mock embedders registered in the test so they don't require GEE, +torch, or any real model weights. list_models_for_data is additionally +exercised against the real catalog, which only reads static describe() +metadata. +""" + +import numpy as np +import pytest + +from rs_embed import ( + UserData, + get_embedding_from_data, + get_embeddings_batch_from_data, + list_models_for_data, +) +from rs_embed.core import registry +from rs_embed.core.embedding import Embedding +from rs_embed.core.errors import ModelError, SpecError +from rs_embed.core.specs import OutputSpec, PointBuffer, TemporalSpec +from rs_embed.embedders.base import EmbedderBase + +S2 = "COPERNICUS/S2_SR_HARMONIZED" +_POINT = PointBuffer(lon=-88.2, lat=40.1, buffer_m=320) + + +class _MockFromDataEmbedder(EmbedderBase): + """Captures the inputs it receives; no I/O.""" + + model_name = "mock_from_data" + last_input = None + last_sensor = None + last_model_config = None + seen_temporals: list = [] + + def describe(self): + return { + "type": "mock", + "backend": ["gee", "auto"], + "output": ["pooled", "grid"], + "inputs": {"collection": S2, "bands": ["B4", "B3", "B2"]}, + } + + def get_embedding( + self, + *, + spatial, + temporal, + sensor, + output, + backend, + device="auto", + input_chw=None, + model_config=None, + ): + type(self).last_input = input_chw + type(self).last_sensor = sensor + type(self).last_model_config = model_config + type(self).seen_temporals.append(temporal) + return Embedding( + data=np.arange(4, dtype=np.float32), + meta={"model": self.model_name, "output": output.mode, "temporal": temporal}, + ) + + +class _MockPrecomputedFromDataEmbedder(EmbedderBase): + model_name = "mock_from_data_precomputed" + _is_precomputed = True + + def describe(self): + return {"type": "precomputed", "backend": ["local"], "output": ["pooled"]} + + +class _MockGeorefFromDataEmbedder(_MockFromDataEmbedder): + model_name = "mock_from_data_georef" + _requires_georef = True + + +class _MockChunkTrackingEmbedder(_MockFromDataEmbedder): + model_name = "mock_from_data_chunks" + seen_chunk_sizes: list = [] + + def get_embeddings_batch_from_inputs(self, *, spatials, input_chws, **kwargs): + type(self).seen_chunk_sizes.append(len(input_chws)) + return super().get_embeddings_batch_from_inputs( + spatials=spatials, input_chws=input_chws, **kwargs + ) + + +class _MockTiledFromDataEmbedder(_MockFromDataEmbedder): + """Advertises a tile size (defaults.image_size) so the tiler engages.""" + + model_name = "mock_from_data_tiled" + seen_input_shapes: list = [] + from_inputs_calls: int = 0 + + def describe(self): + desc = super().describe() + desc["defaults"] = {"image_size": 4} + return desc + + def get_embedding(self, **kwargs): + x = kwargs.get("input_chw") + if x is not None: + type(self).seen_input_shapes.append(tuple(np.asarray(x).shape)) + return super().get_embedding(**kwargs) + + def get_embeddings_batch_from_inputs(self, *, spatials, input_chws, **kwargs): + type(self).from_inputs_calls += 1 + return super().get_embeddings_batch_from_inputs( + spatials=spatials, input_chws=input_chws, **kwargs + ) + + +class _MockFlexTiledEmbedder(_MockTiledFromDataEmbedder): + """Flexible-size model: adapts to the input (patch multiple of 4) unless + an explicit model_config image_size restores fixed-size behavior.""" + + model_name = "mock_from_data_flex" + + def resolve_input_image_size(self, model_config, *, input_hw=None): + explicit = (model_config or {}).get("image_size") + if explicit is not None: + return int(explicit) + if input_hw is None: + return 4 + side = max(int(input_hw[0]), int(input_hw[1])) + return max(4, -(-side // 4) * 4) + + +@pytest.fixture(autouse=True) +def register_mocks(): + registry.register("mock_from_data")(_MockFromDataEmbedder) + registry.register("mock_from_data_precomputed")(_MockPrecomputedFromDataEmbedder) + registry.register("mock_from_data_georef")(_MockGeorefFromDataEmbedder) + registry.register("mock_from_data_chunks")(_MockChunkTrackingEmbedder) + registry.register("mock_from_data_tiled")(_MockTiledFromDataEmbedder) + registry.register("mock_from_data_flex")(_MockFlexTiledEmbedder) + _MockChunkTrackingEmbedder.seen_chunk_sizes = [] + _MockTiledFromDataEmbedder.seen_input_shapes = [] + _MockTiledFromDataEmbedder.from_inputs_calls = 0 + _MockFlexTiledEmbedder.seen_input_shapes = [] + _MockFlexTiledEmbedder.from_inputs_calls = 0 + _MockFromDataEmbedder.last_input = None + _MockFromDataEmbedder.last_sensor = None + _MockFromDataEmbedder.last_model_config = None + _MockFromDataEmbedder.seen_temporals = [] + yield + + +def _twelve_band_userdata( + *, fill_by_channel=True, hw=(8, 8), temporal=None, declare_bands=True, spatial=_POINT +): + bands = ("B1", "B2", "B3", "B4", "B5", "B6", "B7", "B8", "B8A", "B9", "B11", "B12") + if fill_by_channel: + arr = np.stack([np.full(hw, i * 1000.0, dtype=np.float32) for i in range(len(bands))]) + else: + arr = np.full((len(bands), *hw), 5000.0, dtype=np.float32) + return UserData( + data=arr, + collection="s2", + spatial=spatial, + bands=bands if declare_bands else None, + temporal=temporal, + ) + + +# ── single ───────────────────────────────────────────────────────── + + +def test_superset_data_is_sliced_to_model_band_order(): + emb = get_embedding_from_data("mock_from_data", _twelve_band_userdata()) + x = _MockFromDataEmbedder.last_input + assert x is not None and x.shape == (3, 8, 8) + # model order (B4, B3, B2) -> channels 3, 2, 1 of the declared 12-band cube + assert [float(x[i, 0, 0]) for i in range(3)] == [3000.0, 2000.0, 1000.0] + assert emb.data.shape == (4,) + + +def test_omitted_bands_default_to_canonical_s2_order(): + emb = get_embedding_from_data("mock_from_data", _twelve_band_userdata(declare_bands=False)) + x = _MockFromDataEmbedder.last_input + assert [float(x[i, 0, 0]) for i in range(3)] == [3000.0, 2000.0, 1000.0] + assert emb.meta["user_input"]["declared_bands"][:3] == ["B1", "B2", "B3"] + + +def test_temporal_travels_with_the_data(): + t = TemporalSpec.year(2022) + emb = get_embedding_from_data("mock_from_data", _twelve_band_userdata(temporal=t)) + assert emb.meta["temporal"] == t + + +def test_user_input_meta_records_declaration_and_selection(): + emb = get_embedding_from_data("mock_from_data", _twelve_band_userdata()) + ui = emb.meta["user_input"] + assert ui["source"] == "user_data" + assert ui["collection"] == S2 + assert ui["bands_used"] == ["B4", "B3", "B2"] + assert ui["channel_indices"] == [3, 2, 1] + + +def test_missing_band_refuses_with_band_name(): + data = UserData( + data=np.full((2, 8, 8), 5000.0, dtype=np.float32), + collection="s2", + spatial=_POINT, + bands=("B2", "B3"), + ) + with pytest.raises(ModelError, match=r"lacks \['B4'\]"): + get_embedding_from_data("mock_from_data", data) + + +def test_collection_mismatch_refuses(): + data = UserData( + data=np.full((3, 8, 8), 0.5, dtype=np.float32), + collection="s1", + spatial=_POINT, + bands=("VV", "VH", "ANGLE"), + ) + with pytest.raises(ModelError, match="collection"): + get_embedding_from_data("mock_from_data", data) + + +def test_precomputed_model_refuses(): + with pytest.raises(ModelError, match="precomputed"): + get_embedding_from_data("mock_from_data_precomputed", _twelve_band_userdata()) + + +def test_missing_spatial_is_accepted_by_non_georef_models(): + emb = get_embedding_from_data("mock_from_data", _twelve_band_userdata(spatial=None)) + assert emb.data.shape == (4,) + assert _MockFromDataEmbedder.last_input is not None + + +def test_missing_spatial_refuses_georef_conditioned_models(): + with pytest.raises(ModelError, match="conditions on request geometry"): + get_embedding_from_data("mock_from_data_georef", _twelve_band_userdata(spatial=None)) + + +def test_georef_model_accepts_data_with_spatial(): + emb = get_embedding_from_data("mock_from_data_georef", _twelve_band_userdata()) + assert emb.data.shape == (4,) + + +def test_invalid_userdata_raises_specerror(): + bad = UserData( + data=np.zeros((4, 8, 8), dtype=np.float32), + collection="s2", + spatial=_POINT, + bands=("B2", "B3"), + ) + with pytest.raises(SpecError, match="channels"): + get_embedding_from_data("mock_from_data", bad) + + +def test_model_kwargs_are_forwarded_as_model_config(): + get_embedding_from_data("mock_from_data", _twelve_band_userdata(), variant="large") + assert _MockFromDataEmbedder.last_model_config == {"variant": "large"} + + +def test_output_mode_is_passed_through(): + emb = get_embedding_from_data( + "mock_from_data", + _twelve_band_userdata(), + output=OutputSpec.pooled(), + ) + assert emb.meta["output"] == "pooled" + + +# ── batch ────────────────────────────────────────────────────────── + + +def test_batch_returns_one_embedding_per_item(): + datas = [_twelve_band_userdata(), _twelve_band_userdata()] + embs = get_embeddings_batch_from_data("mock_from_data", datas) + assert len(embs) == 2 + assert all(e.meta["user_input"]["channel_indices"] == [3, 2, 1] for e in embs) + + +def test_batch_groups_by_temporal_and_preserves_order(): + t22, t23 = TemporalSpec.year(2022), TemporalSpec.year(2023) + datas = [ + _twelve_band_userdata(temporal=t22), + _twelve_band_userdata(temporal=t23), + _twelve_band_userdata(temporal=t22), + ] + embs = get_embeddings_batch_from_data("mock_from_data", datas) + assert [e.meta["temporal"] for e in embs] == [t22, t23, t22] + # two distinct temporals -> two dispatches covering all three items + assert sorted(_MockFromDataEmbedder.seen_temporals, key=str) == sorted([t22, t22, t23], key=str) + + +def test_batch_empty_refuses(): + with pytest.raises(ModelError, match="non-empty"): + get_embeddings_batch_from_data("mock_from_data", []) + + +def test_batch_size_chunks_the_dispatch(): + datas = [_twelve_band_userdata() for _ in range(5)] + embs = get_embeddings_batch_from_data("mock_from_data_chunks", datas, batch_size=2) + assert len(embs) == 5 + assert _MockChunkTrackingEmbedder.seen_chunk_sizes == [2, 2, 1] + + +def test_batch_size_default_dispatches_once(): + datas = [_twelve_band_userdata() for _ in range(3)] + get_embeddings_batch_from_data("mock_from_data_chunks", datas) + assert _MockChunkTrackingEmbedder.seen_chunk_sizes == [3] + + +def test_invalid_batch_size_refuses(): + with pytest.raises(ModelError, match="batch_size"): + get_embeddings_batch_from_data("mock_from_data", [_twelve_band_userdata()], batch_size=0) + + +# ── input_prep: tile default vs resize ───────────────────────────── + + +def test_large_input_is_tiled_by_default(): + # 8x8 input, model tile size 4 -> 2x2 grid of native-resolution 4x4 tiles. + emb = get_embedding_from_data("mock_from_data_tiled", _twelve_band_userdata(hw=(8, 8))) + assert _MockTiledFromDataEmbedder.seen_input_shapes == [(3, 4, 4)] * 4 + prep = emb.meta["input_prep"] + assert prep["resolved_mode"] == "tile" + assert emb.meta["user_input"]["bands_used"] == ["B4", "B3", "B2"] + + +def test_input_prep_resize_opts_out_of_tiling(): + emb = get_embedding_from_data( + "mock_from_data_tiled", _twelve_band_userdata(hw=(8, 8)), input_prep="resize" + ) + # One call with the full-size array; the embedder handles the downsample. + assert _MockTiledFromDataEmbedder.seen_input_shapes == [(3, 8, 8)] + assert emb.meta["input_prep"]["resolved_mode"] == "resize" + + +def test_tile_sized_inputs_keep_batched_dispatch_under_tile_default(): + # Items a single tile already covers cannot tile -> one batched dispatch. + datas = [_twelve_band_userdata(hw=(4, 4)), _twelve_band_userdata(hw=(4, 4))] + embs = get_embeddings_batch_from_data("mock_from_data_tiled", datas) + assert _MockTiledFromDataEmbedder.from_inputs_calls == 1 + assert _MockTiledFromDataEmbedder.seen_input_shapes == [(3, 4, 4)] * 2 + assert all(e.meta["input_prep"]["resolved_mode"] == "resize" for e in embs) + + +def test_flexible_model_consumes_large_input_natively_by_default(): + emb = get_embedding_from_data("mock_from_data_flex", _twelve_band_userdata(hw=(8, 8))) + assert _MockFlexTiledEmbedder.seen_input_shapes == [(3, 8, 8)] + assert emb.meta["input_prep"]["resolved_mode"] == "native" + + +def test_flexible_native_pass_warns_above_threshold(monkeypatch): + import rs_embed.tools.runtime as rt + + monkeypatch.setattr(rt, "_FLEX_NATIVE_WARN_PX", 6) + with pytest.warns(UserWarning, match="native pass"): + get_embedding_from_data("mock_from_data_flex", _twelve_band_userdata(hw=(8, 8))) + + +def test_flexible_explicit_image_size_restores_tiling(): + emb = get_embedding_from_data( + "mock_from_data_flex", _twelve_band_userdata(hw=(8, 8)), image_size=4 + ) + assert _MockFlexTiledEmbedder.seen_input_shapes == [(3, 4, 4)] * 4 + assert emb.meta["input_prep"]["resolved_mode"] == "tile" + + +def test_flexible_items_group_by_native_size(): + datas = [ + _twelve_band_userdata(hw=(8, 8)), + _twelve_band_userdata(hw=(12, 12)), + _twelve_band_userdata(hw=(8, 8)), + ] + embs = get_embeddings_batch_from_data("mock_from_data_flex", datas) + # two distinct native sizes -> two batch dispatches, all full-size inputs + assert _MockFlexTiledEmbedder.from_inputs_calls == 2 + assert sorted(_MockFlexTiledEmbedder.seen_input_shapes) == [(3, 8, 8), (3, 8, 8), (3, 12, 12)] + assert len(embs) == 3 + + +# ── list_models_for_data over the real catalog ───────────────────── + + +def test_list_models_for_data_over_catalog_with_12_band_s2(): + report = list_models_for_data(_twelve_band_userdata(fill_by_channel=False)) + by_model = {r["model"]: r for r in report} + + # Precomputed models are incompatible with an explicit reason. + for name in ("tessera", "gse", "copernicus"): + assert not by_model[name]["compatible"] + assert "precomputed" in by_model[name]["reason"] + + # MODIS-based SatVision cannot take S2 data. + assert not by_model["satvision"]["compatible"] + + # A 12-band S2 L2A cube serves the S2-default models. + for name in ("galileo", "prithvi", "clay", "scalemae", "satmae"): + assert by_model[name]["compatible"], by_model[name]["reason"] + assert by_model[name]["bands_used"] + + +def test_list_models_for_data_without_spatial_flags_georef_models(): + report = list_models_for_data(_twelve_band_userdata(fill_by_channel=False, spatial=None)) + by_model = {r["model"]: r for r in report} + for name in ("clay", "prithvi"): + assert not by_model[name]["compatible"] + assert "geometry" in by_model[name]["reason"] + assert by_model["galileo"]["compatible"] + + +def test_list_models_for_data_rgb_only_declaration(): + data = UserData( + data=np.full((3, 8, 8), 5000.0, dtype=np.float32), + collection="s2", + spatial=_POINT, + bands=("B4", "B3", "B2"), + ) + report = list_models_for_data(data) + by_model = {r["model"]: r for r in report} + assert by_model["scalemae"]["compatible"] + assert not by_model["galileo"]["compatible"] + assert "lacks" in by_model["galileo"]["reason"] diff --git a/tests/test_capabilities_contract.py b/tests/test_capabilities_contract.py index 8b28f4a..916f7fe 100644 --- a/tests/test_capabilities_contract.py +++ b/tests/test_capabilities_contract.py @@ -103,8 +103,14 @@ def test_manages_own_input_prep_flag_matches_expected_models(): assert flagged == {"gse"} -def test_tiled_dispatch_hook_only_overridden_by_thor(): - """Behavior parity with the removed tiling-level 'thor' hardcode.""" +def test_tiled_dispatch_hook_only_overridden_by_declared_models(): + """The tiled-dispatch config hook is a deliberate, pinned registry. + + thor: behavior parity with the removed tiling-level 'thor' hardcode. + olmoearth: FlexiViT consumes dispatched inputs at exactly ``tile_size`` + (no internal resize) — also the channel through which the flexible + native-size user-data path sets the encoder's per-request size. + """ from rs_embed.embedders.base import EmbedderBase overriding = { @@ -113,7 +119,7 @@ def test_tiled_dispatch_hook_only_overridden_by_thor(): if get_embedder_cls(model_id).tiled_dispatch_model_config is not EmbedderBase.tiled_dispatch_model_config } - assert overriding == {"thor"} + assert overriding == {"thor", "olmoearth"} thor_cls = get_embedder_cls("thor") cfg = thor_cls().tiled_dispatch_model_config({"variant": "base"}, tile_size=288) @@ -123,6 +129,13 @@ def test_tiled_dispatch_hook_only_overridden_by_thor(): "_input_prep_tile_size": 288, } + oe_cls = get_embedder_cls("olmoearth") + assert oe_cls().tiled_dispatch_model_config({"variant": "base"}, tile_size=128) == { + "variant": "base", + "image_size": 128, + } + assert oe_cls().tiled_dispatch_model_config(None, tile_size=256) == {"image_size": 256} + def test_declaration_overrides_signature_for_routing(): """Routing must trust the declaration, not the signature, when declared.""" diff --git a/tests/test_clay_embedder.py b/tests/test_clay_embedder.py index 906214c..1c6554c 100644 --- a/tests/test_clay_embedder.py +++ b/tests/test_clay_embedder.py @@ -304,11 +304,36 @@ def _spy_load(*, model_size, device): assert seen["model_size"] == "base" +def test_clay_batch_from_inputs_needs_no_provider(monkeypatch): + """get_embeddings_batch_from_inputs must honor the same lazy-provider + convention as the single path: prefetched inputs never touch the provider + (a user-data batch on a machine without GEE must work).""" + import rs_embed.embedders.onthefly_clay as clay + + emb = ClayEmbedder() + + def _boom(_backend): + raise AssertionError("prefetched batch inputs must not touch the provider") + + monkeypatch.setattr(emb, "_get_provider", _boom) + monkeypatch.setattr(clay, "_load_clay_model", _fake_load) + monkeypatch.setattr(clay, "_clay_forward_tokens_and_cls_batch", _fake_forward) + + out = emb.get_embeddings_batch_from_inputs( + spatials=[PointBuffer(lon=0.0, lat=0.0, buffer_m=256)], + input_chws=[np.full((10, 8, 8), 5000.0, dtype=np.float32)], + temporal=TemporalSpec.year(2021), + output=OutputSpec.pooled(), + backend="auto", + ) + assert len(out) == 1 + assert out[0].data.shape == (2,) + + def test_clay_batch_from_inputs_crops_per_item_roi(monkeypatch): import rs_embed.embedders.onthefly_clay as clay emb = ClayEmbedder() - monkeypatch.setattr(emb, "_get_provider", lambda _backend: object()) monkeypatch.setattr(clay, "_load_clay_model", _fake_load) monkeypatch.setattr(clay, "_clay_forward_tokens_and_cls_batch", _fake_forward) diff --git a/tests/test_olmoearth.py b/tests/test_olmoearth.py index bd696f1..bac808b 100644 --- a/tests/test_olmoearth.py +++ b/tests/test_olmoearth.py @@ -1421,3 +1421,42 @@ def test_get_embeddings_batch_from_inputs_multi_variable_frames(monkeypatch): assert out[1].meta["n_frames"] == 1 for e in out: assert e.data.shape == (128,) + + +# ── resolve_input_image_size (flexible-size tiling threshold) ────── + + +def test_resolve_input_image_size_defaults_to_training_tile(): + emb = oe.OlmoEarthEmbedder() + assert emb.resolve_input_image_size(None) == oe._DEFAULT_IMAGE_SIZE + + +def test_resolve_input_image_size_honors_model_config(): + emb = oe.OlmoEarthEmbedder() + assert emb.resolve_input_image_size({"image_size": 512, "patch_size": 8}) == 512 + + +def test_resolve_input_image_size_rejects_non_divisible(): + emb = oe.OlmoEarthEmbedder() + with pytest.raises(ModelError, match="divisible"): + emb.resolve_input_image_size({"image_size": 130, "patch_size": 4}) + + +def test_resolve_input_image_size_adapts_to_input(): + emb = oe.OlmoEarthEmbedder() + # already patch-divisible -> consumed natively as-is + assert emb.resolve_input_image_size(None, input_hw=(300, 300)) == 300 + # snapped up to the next patch multiple (default patch 4) + assert emb.resolve_input_image_size(None, input_hw=(301, 299)) == 304 + # explicit config wins over the input + assert ( + emb.resolve_input_image_size({"image_size": 512, "patch_size": 8}, input_hw=(300, 300)) + == 512 + ) + + +def test_tiled_dispatch_model_config_injects_image_size(): + emb = oe.OlmoEarthEmbedder() + assert emb.tiled_dispatch_model_config(None, tile_size=256) == {"image_size": 256} + out = emb.tiled_dispatch_model_config({"variant": "base"}, tile_size=128) + assert out == {"variant": "base", "image_size": 128} diff --git a/tests/test_user_data_matching.py b/tests/test_user_data_matching.py new file mode 100644 index 0000000..5a762bc --- /dev/null +++ b/tests/test_user_data_matching.py @@ -0,0 +1,240 @@ +"""Tests for the user-data declaration + sensor matching layer. + +Pure unit tests over ``core.types.UserData`` and ``tools.user_data``; no +provider access, no model weights. +""" + +import numpy as np +import pytest + +from rs_embed.core.errors import ModelError, SpecError +from rs_embed.core.specs import PointBuffer, SensorSpec +from rs_embed.core.types import UserData +from rs_embed.providers.prefetch_plan import select_prefetched_channels +from rs_embed.tools.user_data import ( + canonical_band_names, + match_user_data_to_sensor, + normalize_collection_id, + resolve_declared_bands, + warn_on_suspicious_value_range, +) + +S2 = "COPERNICUS/S2_SR_HARMONIZED" +_POINT = PointBuffer(lon=-88.2, lat=40.1, buffer_m=320) + + +def _user_data(bands, *, collection=S2, shape_hw=(4, 4), tchw=False, fill=None): + c = len(bands) + if fill is None: + # channel i is constant i, so tests can assert selection order by value + base = np.stack([np.full(shape_hw, i, dtype=np.float32) for i in range(c)]) + else: + base = np.full((c, *shape_hw), fill, dtype=np.float32) + data = np.repeat(base[None, ...], 2, axis=0) if tchw else base + return UserData(data=data, collection=collection, spatial=_POINT, bands=tuple(bands)) + + +def _sensor(bands, *, collection=S2): + return SensorSpec(collection=collection, bands=tuple(bands)) + + +# ── UserData.validate ────────────────────────────────────────────── + + +def test_validate_accepts_chw_and_tchw(): + _user_data(["B2", "B3"]).validate() + _user_data(["B2", "B3"], tchw=True).validate() + + +def test_validate_rejects_channel_band_mismatch(): + bad = UserData( + data=np.zeros((3, 4, 4), dtype=np.float32), + collection=S2, + spatial=_POINT, + bands=("B2", "B3"), + ) + with pytest.raises(SpecError, match="3 channels"): + bad.validate() + + +def test_validate_accepts_missing_spatial(): + # spatial is optional at the declaration level; georef-conditioned models + # refuse at request time instead (never fabricate coordinates). + UserData( + data=np.zeros((2, 4, 4), dtype=np.float32), + collection=S2, + bands=("B2", "B3"), + ).validate() + + +def test_validate_rejects_bad_ndim_and_empty_fields(): + with pytest.raises(SpecError, match=r"\[C,H,W\]"): + UserData(data=np.zeros((4, 4)), collection=S2, spatial=_POINT, bands=("B2",)).validate() + with pytest.raises(SpecError, match="collection"): + UserData( + data=np.zeros((1, 4, 4)), collection=" ", spatial=_POINT, bands=("B2",) + ).validate() + with pytest.raises(SpecError, match="bands"): + UserData(data=np.zeros((1, 4, 4)), collection=S2, spatial=_POINT, bands=()).validate() + with pytest.raises(SpecError, match="scale_m"): + UserData( + data=np.zeros((1, 4, 4)), collection=S2, spatial=_POINT, bands=("B2",), scale_m=0 + ).validate() + + +# ── collection + band normalization ─────────────────────────────── + + +def test_collection_aliases_resolve_and_full_ids_pass_through(): + assert normalize_collection_id("s2") == S2 + assert normalize_collection_id("Sentinel-2") == S2 + assert normalize_collection_id("s2-l2a") == S2 + assert normalize_collection_id("s1") == "COPERNICUS/S1_GRD" + assert normalize_collection_id(S2) == S2 + assert normalize_collection_id("SOME/OTHER/COLLECTION") == "SOME/OTHER/COLLECTION" + + +def test_canonical_band_names_resolve_aliases_and_case(): + assert canonical_band_names(S2, ("RED", "green", "b2")) == ("B4", "B3", "B2") + assert canonical_band_names(S2, ("SWIR_1", "NIR_NARROW")) == ("B11", "B8A") + + +# ── default band order (bands=None) ──────────────────────────────── + +_S2_CANONICAL_12 = ("B1", "B2", "B3", "B4", "B5", "B6", "B7", "B8", "B8A", "B9", "B11", "B12") + + +def test_omitted_bands_default_to_canonical_s2_order(): + data = UserData( + data=np.zeros((12, 4, 4), dtype=np.float32), + collection="s2", + spatial=_POINT, + ) + data.validate() + assert resolve_declared_bands(data) == _S2_CANONICAL_12 + idx = match_user_data_to_sensor(data, _sensor(["B4", "B3", "B2"]), model_name="m") + assert idx == (3, 2, 1) + + +def test_omitted_bands_with_wrong_channel_count_refuse(): + data = UserData( + data=np.zeros((10, 4, 4), dtype=np.float32), + collection="s2", + spatial=_POINT, + ) + with pytest.raises(SpecError, match="Declare bands explicitly"): + resolve_declared_bands(data) + + +def test_omitted_bands_without_canonical_order_refuse(): + data = UserData( + data=np.zeros((2, 4, 4), dtype=np.float32), + collection="s1", + spatial=_POINT, + ) + with pytest.raises(SpecError, match="no canonical band order"): + resolve_declared_bands(data) + + +# ── matching ─────────────────────────────────────────────────────── + + +def test_superset_data_is_sliced_into_model_band_order(): + data = _user_data(["B1", "B2", "B3", "B4", "B8"]) + idx = match_user_data_to_sensor(data, _sensor(["B4", "B3", "B2"]), model_name="m") + assert idx == (3, 2, 1) + sliced = select_prefetched_channels(data.data, idx) + assert sliced.shape == (3, 4, 4) + assert [float(sliced[i, 0, 0]) for i in range(3)] == [3.0, 2.0, 1.0] + + +def test_exact_match_is_identity(): + data = _user_data(["B2", "B3", "B4"]) + idx = match_user_data_to_sensor(data, _sensor(["B2", "B3", "B4"]), model_name="m") + assert idx == (0, 1, 2) + + +def test_tchw_slicing_selects_channel_axis(): + data = _user_data(["B1", "B2", "B3", "B4"], tchw=True) + idx = match_user_data_to_sensor(data, _sensor(["B4", "B2"]), model_name="m") + sliced = select_prefetched_channels(data.data, idx) + assert sliced.shape == (2, 2, 4, 4) + assert [float(sliced[0, i, 0, 0]) for i in range(2)] == [3.0, 1.0] + + +def test_alias_bands_match_model_alias_bands(): + # Prithvi-style: model declares HLS-style names (NIR_NARROW -> B8A), user + # declares S2 names; both sides canonicalize to the same vocabulary. + data = _user_data(["B2", "B3", "B4", "B8A", "B11", "B12"]) + sensor = _sensor(["BLUE", "GREEN", "RED", "NIR_NARROW", "SWIR_1", "SWIR_2"]) + idx = match_user_data_to_sensor(data, sensor, model_name="m") + assert idx == (0, 1, 2, 3, 4, 5) + + +def test_alias_bands_refuse_when_canonical_band_missing(): + # NIR_NARROW canonicalizes to B8A; broad-NIR B8 does not satisfy it. + data = _user_data(["B2", "B3", "B4", "B8", "B11", "B12"]) + sensor = _sensor(["BLUE", "GREEN", "RED", "NIR_NARROW", "SWIR_1", "SWIR_2"]) + with pytest.raises(ModelError, match=r"lacks \['B8A'\]"): + match_user_data_to_sensor(data, sensor, model_name="m") + + +def test_missing_band_refuses_with_names(): + data = _user_data(["B2", "B3", "B4"]) + with pytest.raises(ModelError, match=r"lacks \['B8'\]"): + match_user_data_to_sensor(data, _sensor(["B2", "B3", "B4", "B8"]), model_name="m") + + +def test_collection_mismatch_refuses(): + data = _user_data(["VV", "VH"], collection="s1") + with pytest.raises(ModelError, match="collection"): + match_user_data_to_sensor(data, _sensor(["B2", "B3"]), model_name="m") + + +def test_duplicate_user_bands_refuse(): + data = UserData( + data=np.zeros((3, 4, 4), dtype=np.float32), + collection=S2, + spatial=_POINT, + bands=("B2", "RED", "B4"), # RED aliases to B4 -> duplicate + ) + with pytest.raises(ModelError, match="duplicate"): + match_user_data_to_sensor(data, _sensor(["B2", "B4"]), model_name="m") + + +# ── value-range warning ──────────────────────────────────────────── + + +def test_normalized_looking_s2_values_warn(): + data = _user_data(["B2", "B3"], fill=0.3) + with pytest.warns(UserWarning, match="0..10000"): + warn_on_suspicious_value_range(data) + + +def test_raw_dn_s2_values_do_not_warn(): + data = _user_data(["B2", "B3"], fill=4321.0) + with warnings_disabled_check(): + warn_on_suspicious_value_range(data) + + +def test_non_s2_collections_never_warn(): + data = _user_data(["VV", "VH"], collection="s1", fill=0.02) + with warnings_disabled_check(): + warn_on_suspicious_value_range(data) + + +class warnings_disabled_check: + """Context asserting no UserWarning was emitted inside the block.""" + + def __enter__(self): + import warnings as _warnings + + self._catcher = _warnings.catch_warnings(record=True) + self._records = self._catcher.__enter__() + _warnings.simplefilter("always") + return self + + def __exit__(self, exc_type, exc, tb): + self._catcher.__exit__(exc_type, exc, tb) + assert not [w for w in self._records if issubclass(w.category, UserWarning)] + return False