From 4ec862f3b7f6ffcf9d3639bf1b8db86325ba66b6 Mon Sep 17 00:00:00 2001 From: Dinghye <1269226384@qq.com> Date: Thu, 20 Aug 2026 17:37:41 -0500 Subject: [PATCH 01/12] feat: user-data declararion+ sensor matching layer --- src/rs_embed/core/types.py | 64 ++++++++++ src/rs_embed/providers/gee_utils.py | 11 +- src/rs_embed/tools/user_data.py | 143 ++++++++++++++++++++++ tests/test_user_data_matching.py | 181 ++++++++++++++++++++++++++++ 4 files changed, 398 insertions(+), 1 deletion(-) create mode 100644 src/rs_embed/tools/user_data.py create mode 100644 tests/test_user_data_matching.py diff --git a/src/rs_embed/core/types.py b/src/rs_embed/core/types.py index d5e1457..bbebab6 100644 --- a/src/rs_embed/core/types.py +++ b/src/rs_embed/core/types.py @@ -12,6 +12,7 @@ import numpy as np +from .errors import SpecError from .specs import FetchSpec, InputPrepSpec, SensorSpec # ── Enums ────────────────────────────────────────────────────────── @@ -143,6 +144,69 @@ class FetchResult: meta: dict[str, Any] +@dataclass(frozen=True) +class UserData: + """User-provided imagery with a declaration of what the pixels are. + + The bring-your-own-data entrypoints (:func:`rs_embed.get_embedding_from_data` + and friends) match this declaration against a model's input sensor and + refuse the request 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"``. + bands : tuple[str, ...] + Band name per channel, e.g. ``("B2", "B3", "B4", ...)``. Aliases like + ``"RED"`` resolve the same way provider fetches resolve them. + scale_m : int or None + Optional nominal pixel size in meters, recorded as provenance. + """ + + data: np.ndarray + collection: str + bands: tuple[str, ...] + scale_m: int | None = None + + def validate(self) -> None: + """Validate the declaration's internal consistency. + + 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 not self.bands 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.") + 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)}." + ) + 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/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/user_data.py b/src/rs_embed/tools/user_data.py new file mode 100644 index 0000000..b0f2ffc --- /dev/null +++ b/src/rs_embed/tools/user_data.py @@ -0,0 +1,143 @@ +"""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 +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",) + + +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 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. + """ + 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, tuple(data.bands)) + 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_user_data_matching.py b/tests/test_user_data_matching.py new file mode 100644 index 0000000..b909d51 --- /dev/null +++ b/tests/test_user_data_matching.py @@ -0,0 +1,181 @@ +"""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 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, + warn_on_suspicious_value_range, +) + +S2 = "COPERNICUS/S2_SR_HARMONIZED" + + +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, 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, bands=("B2", "B3")) + with pytest.raises(SpecError, match="3 channels"): + bad.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, bands=("B2",)).validate() + with pytest.raises(SpecError, match="collection"): + UserData(data=np.zeros((1, 4, 4)), collection=" ", bands=("B2",)).validate() + with pytest.raises(SpecError, match="bands"): + UserData(data=np.zeros((1, 4, 4)), collection=S2, bands=()).validate() + with pytest.raises(SpecError, match="scale_m"): + UserData(data=np.zeros((1, 4, 4)), collection=S2, 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") + + +# ── 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, + 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 From 6cd2dc03df00b16e53b6bf33a8e0775e8b7a42c2 Mon Sep 17 00:00:00 2001 From: Dinghye <1269226384@qq.com> Date: Thu, 20 Aug 2026 17:41:14 -0500 Subject: [PATCH 02/12] feat: get_embeddings_from_data/get_embeddings_batch_from_data --- src/rs_embed/__init__.py | 9 ++ src/rs_embed/api.py | 270 ++++++++++++++++++++++++++++++++++ src/rs_embed/tools/runtime.py | 92 ++++++++++++ tests/test_api_from_data.py | 240 ++++++++++++++++++++++++++++++ 4 files changed, 611 insertions(+) create mode 100644 tests/test_api_from_data.py 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..7f002f0 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,19 @@ 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 ( + warn_on_suspicious_value_range as _warn_on_suspicious_value_range, +) # ----------------------------------------------------------------------------- # Internal helpers @@ -398,6 +414,260 @@ 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 _prepare_user_data_inputs( + model_n: str, + datas: list[UserData], + sensor_eff: SensorSpec, +) -> 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() + idx = _match_user_data_to_sensor(data, sensor_eff, model_name=model_n) + _warn_on_suspicious_value_range(data) + arrays.append(_select_prefetched_channels(np.asarray(data.data, dtype=np.float32), idx)) + metas.append( + { + "source": "user_data", + "collection": _normalize_collection_id(data.collection), + "declared_bands": list(data.bands), + "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, + spatial: SpatialSpec, + temporal: TemporalSpec | None = None, + modality: str | None = None, + output: OutputSpec = OutputSpec.pooled(), + device: str = "auto", + **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 supplies a + :class:`~rs_embed.UserData` declaring what the pixels are (collection + + band names, raw provider units). 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 imagery and its declaration. ``data.data`` must be ``[C,H,W]`` + (or ``[T,C,H,W]`` for time-series models) with raw provider values in + the declared band order. + spatial : SpatialSpec + Where the imagery is located. Required: several models condition on + geometry (e.g. lat/lon or GSD embeddings), and embedding metadata + records it. + temporal : TemporalSpec or None + When the imagery was acquired, for models that condition on time. + 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. + **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. + + 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 *data*, *spatial*, or *temporal* fail validation. + + Examples + -------- + >>> emb = get_embedding_from_data( + ... "galileo", + ... data=UserData(data=chw, collection="s2", bands=("B2", "B3", "B4", ...)), + ... spatial=PointBuffer(lon=-88.2, lat=40.1, buffer_m=640), + ... temporal=TemporalSpec.year(2022), + ... ) + """ + return get_embeddings_batch_from_data( + model, + datas=[data], + spatials=[spatial], + temporal=temporal, + modality=modality, + output=output, + device=device, + **model_kwargs, + )[0] + + +def get_embeddings_batch_from_data( + model: str, + *, + datas: list[UserData], + spatials: list[SpatialSpec], + temporal: TemporalSpec | None = None, + modality: str | None = None, + output: OutputSpec = OutputSpec.pooled(), + device: str = "auto", + **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 as long + as every declaration satisfies the model. + + Parameters + ---------- + model : str + Model identifier or alias. + datas : list[UserData] + User imagery declarations, one per item. + spatials : list[SpatialSpec] + Locations aligned with *datas*. + temporal : TemporalSpec or None + Shared acquisition-time filter. + modality : str or None + Optional modality selector. + output : OutputSpec + Output representation policy. + device : str + Target inference device. + **model_kwargs + Model-specific settings, as in :func:`get_embedding`. + + Returns + ------- + list[Embedding] + Embeddings in the same order as *datas*. + + Raises + ------ + ModelError + If lengths mismatch, the model cannot take user data, or any + declaration does not satisfy the model's sensor. + SpecError + If any spec 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].") + if len(datas) != len(spatials): + raise ModelError(f"datas/spatials length mismatch: {len(datas)} != {len(spatials)}") + _validate_spatial_list(spatials=spatials, temporal=temporal, output=output) + sensor_eff = _resolve_user_data_sensor(model_n, modality=modality) + arrays, metas = _prepare_user_data_inputs(model_n, datas, sensor_eff) + return _run_user_input_request( + model_n=model_n, + spatials=spatials, + input_arrays=arrays, + temporal=temporal, + sensor=sensor_eff, + model_config=model_config, + output=output, + device=device, + input_metas=metas, + ) + + +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) + _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/tools/runtime.py b/src/rs_embed/tools/runtime.py index 03a0ac4..8902039 100644 --- a/src/rs_embed/tools/runtime.py +++ b/src/rs_embed/tools/runtime.py @@ -802,3 +802,95 @@ 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, +) -> 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. + + 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']``. + + Returns + ------- + list[Embedding] + Embeddings aligned with *spatials*. + + Raises + ------ + ModelError + If lengths mismatch 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)}" + ) + 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) + + kwargs: dict[str, Any] = { + "spatials": spatials, + "input_chws": input_arrays, + "temporal": temporal, + "sensor": sensor, + "output": output, + "backend": "auto", + "device": device_n, + } + if model_config is not None: + require_model_config_support( + embedder=embedder, + model_config=model_config, + method_name="get_embeddings_batch_from_inputs", + ) + kwargs["model_config"] = model_config + with lock: + embs = embedder.get_embeddings_batch_from_inputs(**kwargs) + embs = [normalize_embedding_output(emb=emb, output=output) for emb in embs] + 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/tests/test_api_from_data.py b/tests/test_api_from_data.py new file mode 100644 index 0000000..0b92516 --- /dev/null +++ b/tests/test_api_from_data.py @@ -0,0 +1,240 @@ +"""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 input_chw it receives; no I/O.""" + + model_name = "mock_from_data" + last_input = None + last_sensor = None + last_model_config = None + + 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 + return Embedding( + data=np.arange(4, dtype=np.float32), + meta={"model": self.model_name, "output": output.mode}, + ) + + +class _MockPrecomputedFromDataEmbedder(EmbedderBase): + model_name = "mock_from_data_precomputed" + _is_precomputed = True + + def describe(self): + return {"type": "precomputed", "backend": ["local"], "output": ["pooled"]} + + +@pytest.fixture(autouse=True) +def register_mocks(): + registry.register("mock_from_data")(_MockFromDataEmbedder) + registry.register("mock_from_data_precomputed")(_MockPrecomputedFromDataEmbedder) + _MockFromDataEmbedder.last_input = None + _MockFromDataEmbedder.last_sensor = None + _MockFromDataEmbedder.last_model_config = None + yield + + +def _twelve_band_userdata(*, fill_by_channel=True, hw=(8, 8)): + 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", bands=bands) + + +# ── single ───────────────────────────────────────────────────────── + + +def test_superset_data_is_sliced_to_model_band_order(): + emb = get_embedding_from_data( + "mock_from_data", + data=_twelve_band_userdata(), + spatial=_POINT, + temporal=TemporalSpec.year(2022), + ) + 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_user_input_meta_records_declaration_and_selection(): + emb = get_embedding_from_data( + "mock_from_data", + data=_twelve_band_userdata(), + spatial=_POINT, + ) + 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", + bands=("B2", "B3"), + ) + with pytest.raises(ModelError, match=r"lacks \['B4'\]"): + get_embedding_from_data("mock_from_data", data=data, spatial=_POINT) + + +def test_collection_mismatch_refuses(): + data = UserData( + data=np.full((3, 8, 8), 0.5, dtype=np.float32), + collection="s1", + bands=("VV", "VH", "ANGLE"), + ) + with pytest.raises(ModelError, match="collection"): + get_embedding_from_data("mock_from_data", data=data, spatial=_POINT) + + +def test_precomputed_model_refuses(): + with pytest.raises(ModelError, match="precomputed"): + get_embedding_from_data( + "mock_from_data_precomputed", + data=_twelve_band_userdata(), + spatial=_POINT, + ) + + +def test_invalid_userdata_raises_specerror(): + bad = UserData( + data=np.zeros((4, 8, 8), dtype=np.float32), + collection="s2", + bands=("B2", "B3"), + ) + with pytest.raises(SpecError, match="channels"): + get_embedding_from_data("mock_from_data", data=bad, spatial=_POINT) + + +def test_model_kwargs_are_forwarded_as_model_config(): + get_embedding_from_data( + "mock_from_data", + data=_twelve_band_userdata(), + spatial=_POINT, + variant="large", + ) + assert _MockFromDataEmbedder.last_model_config == {"variant": "large"} + + +def test_grid_output_mode_is_passed_through(): + emb = get_embedding_from_data( + "mock_from_data", + data=_twelve_band_userdata(), + spatial=_POINT, + 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=datas, + spatials=[_POINT, _POINT], + ) + assert len(embs) == 2 + assert all(e.meta["user_input"]["channel_indices"] == [3, 2, 1] for e in embs) + + +def test_batch_length_mismatch_refuses(): + with pytest.raises(ModelError, match="length mismatch"): + get_embeddings_batch_from_data( + "mock_from_data", + datas=[_twelve_band_userdata()], + spatials=[_POINT, _POINT], + ) + + +def test_batch_empty_refuses(): + with pytest.raises(ModelError, match="non-empty"): + get_embeddings_batch_from_data("mock_from_data", datas=[], spatials=[]) + + +# ── 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_rgb_only_declaration(): + data = UserData( + data=np.full((3, 8, 8), 5000.0, dtype=np.float32), + collection="s2", + 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"] From be4ace47597628e09392e117ab9435a5e05c26bc Mon Sep 17 00:00:00 2001 From: Dinghye <1269226384@qq.com> Date: Thu, 20 Aug 2026 17:41:43 -0500 Subject: [PATCH 03/12] docs:user api page+changelog --- CHANGELOG.md | 4 +++ docs/user_data.md | 85 +++++++++++++++++++++++++++++++++++++++++++++++ mkdocs.yml | 1 + 3 files changed, 90 insertions(+) create mode 100644 docs/user_data.md diff --git a/CHANGELOG.md b/CHANGELOG.md index a967749..43cd625 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,10 @@ The format is based on Keep a Changelog, and the project follows Semantic Versio ## [Unreleased] +### Added + +- **Bring-your-own-data API** — compute embeddings from imagery you already have, without any provider fetch or provider auth. `UserData(data, collection, bands, scale_m=None)` declares what an array is (collection + one band name per channel, raw provider units, `[C,H,W]` or `[T,C,H,W]`); `get_embedding_from_data` / `get_embeddings_batch_from_data` match the declaration 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). `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 An export-correctness fix for time-series models plus a new input-inspection API. Multi-model exports that combined a temporal model (olmoearth/agrifm/anysat/galileo/prithvi) with other models on a shared collection should be re-exported (see Fixed). diff --git a/docs/user_data.md b/docs/user_data.md new file mode 100644 index 0000000..6ebbdec --- /dev/null +++ b/docs/user_data.md @@ -0,0 +1,85 @@ +# API: User-Provided Data + +This page covers the bring-your-own-data API: computing embeddings from imagery you already have, instead of provider-fetched imagery. + +Related pages: [API: Embedding](api_embedding.md), [API: Specs and Data Structures](api_specs.md). + +--- + +## Concept + +Every on-the-fly model declares an input sensor (a collection plus an ordered band list). The bring-your-own-data path asks you to declare what your array is — the collection it came from and one band name per channel — and matches that declaration against the model's sensor: + +- **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 can serve models that need 3, 6, or 10 of those bands. +- **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. + +Values must be **raw provider units** for the declared collection (e.g. Sentinel-2 L2A surface-reflectance DN in `0..10000`), exactly what a provider fetch would return. Per-model normalization stays inside each embedder, so you never need to know a model's normalization. Data that looks already normalized (max ≤ 1.5 on an S2 declaration) triggers a warning. + +`spatial` is still required: several models condition on geometry (lat/lon, GSD), and embedding metadata records provenance. Pass the location your imagery covers. + +--- + +## 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" + bands: tuple[str, ...], # one band name per channel, e.g. ("B2", "B3", ...) + scale_m: int | None = None, # optional nominal pixel size, recorded as provenance +) +``` + +Collection aliases: `"s2"` / `"sentinel-2"` / `"s2-l2a"` → `COPERNICUS/S2_SR_HARMONIZED`; `"s1"` / `"sentinel-1"` → `COPERNICUS/S1_GRD`. Full collection ids pass through unchanged. Band aliases resolve the same way provider fetches resolve them (`"RED"` → `"B4"`, `"NIR_NARROW"` → `"B8A"`, …). + +Multi-frame `[T,C,H,W]` arrays are only meaningful for time-series models (galileo, prithvi, olmoearth, anysat, agrifm); single-frame models reject them. + +--- + +## Functions + +### get_embedding_from_data + +```python +from rs_embed import UserData, get_embedding_from_data +from rs_embed.core.specs import PointBuffer, TemporalSpec + +emb = get_embedding_from_data( + "galileo", + data=UserData(data=cube, collection="s2", bands=bands), + spatial=PointBuffer(lon=-88.2, lat=40.1, buffer_m=640), + temporal=TemporalSpec.year(2022), +) +``` + +Returns one `Embedding`; `meta["user_input"]` records the declaration and the channel selection actually fed to the model (`bands_used`, `channel_indices`). + +### get_embeddings_batch_from_data + +Batch counterpart: `datas: list[UserData]` aligned with `spatials: list[SpatialSpec]`. Each item is matched independently, so items may declare different band orders. + +### list_models_for_data + +```python +from rs_embed import list_models_for_data + +report = list_models_for_data(my_declaration) +[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). + +--- + +## 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()` | +| 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 From 1c704301aa35f92ce823afee2885c4f0ebf2c30e Mon Sep 17 00:00:00 2001 From: Dinghye <1269226384@qq.com> Date: Thu, 20 Aug 2026 18:03:10 -0500 Subject: [PATCH 04/12] feat: _requires_georef embedder flag on geometry-conditioned models --- src/rs_embed/embedders/base.py | 4 ++++ src/rs_embed/embedders/onthefly_clay.py | 3 +++ src/rs_embed/embedders/onthefly_prithvi.py | 3 +++ 3 files changed, 10 insertions(+) diff --git a/src/rs_embed/embedders/base.py b/src/rs_embed/embedders/base.py index bcd5c39..d1aebdc 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] = {} diff --git a/src/rs_embed/embedders/onthefly_clay.py b/src/rs_embed/embedders/onthefly_clay.py index 5aaeb82..f1b40ac 100644 --- a/src/rs_embed/embedders/onthefly_clay.py +++ b/src/rs_embed/embedders/onthefly_clay.py @@ -549,6 +549,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 diff --git a/src/rs_embed/embedders/onthefly_prithvi.py b/src/rs_embed/embedders/onthefly_prithvi.py index e89d341..9b2b353 100644 --- a/src/rs_embed/embedders/onthefly_prithvi.py +++ b/src/rs_embed/embedders/onthefly_prithvi.py @@ -912,6 +912,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 From ffec755b8d36044831501dfe92a3aab29b55e06d Mon Sep 17 00:00:00 2001 From: Dinghye <1269226384@qq.com> Date: Thu, 20 Aug 2026 18:03:58 -0500 Subject: [PATCH 05/12] refactor: UserData carries its own context; spatial opt via georef flag --- CHANGELOG.md | 2 +- docs/user_data.md | 43 +++++++--- src/rs_embed/api.py | 136 ++++++++++++++++++------------- src/rs_embed/core/types.py | 56 +++++++++---- src/rs_embed/tools/user_data.py | 63 +++++++++++++- tests/test_api_from_data.py | 130 +++++++++++++++++++---------- tests/test_user_data_matching.py | 73 +++++++++++++++-- 7 files changed, 365 insertions(+), 138 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 43cd625..69db03d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,7 +10,7 @@ The format is based on Keep a Changelog, and the project follows Semantic Versio ### Added -- **Bring-your-own-data API** — compute embeddings from imagery you already have, without any provider fetch or provider auth. `UserData(data, collection, bands, scale_m=None)` declares what an array is (collection + one band name per channel, raw provider units, `[C,H,W]` or `[T,C,H,W]`); `get_embedding_from_data` / `get_embeddings_batch_from_data` match the declaration 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). `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). +- **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 items carry their own spatial/temporal; same-temporal items dispatch together, results return in input order). 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/user_data.md b/docs/user_data.md index 6ebbdec..c6db39b 100644 --- a/docs/user_data.md +++ b/docs/user_data.md @@ -8,15 +8,18 @@ Related pages: [API: Embedding](api_embedding.md), [API: Specs and Data Structur ## Concept -Every on-the-fly model declares an input sensor (a collection plus an ordered band list). The bring-your-own-data path asks you to declare what your array is — the collection it came from and one band name per channel — and matches that declaration against the model's sensor: +The flow has two steps: + +1. **Register** each piece of imagery 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. +2. **Embed** by naming a model: `get_embedding_from_data("galileo", data)`. Nothing else is needed — the declaration already carries the full context. + +Every on-the-fly model declares an input sensor (a collection plus an ordered band list). Your declaration is matched against it: - **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 can serve models that need 3, 6, or 10 of those bands. - **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. Values must be **raw provider units** for the declared collection (e.g. Sentinel-2 L2A surface-reflectance DN in `0..10000`), exactly what a provider fetch would return. Per-model normalization stays inside each embedder, so you never need to know a model's normalization. Data that looks already normalized (max ≤ 1.5 on an S2 declaration) triggers a warning. -`spatial` is still required: several models condition on geometry (lat/lon, GSD), and embedding metadata records provenance. Pass the location your imagery covers. - --- ## UserData @@ -25,13 +28,19 @@ Values must be **raw provider units** for the declared collection (e.g. Sentinel 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" - bands: tuple[str, ...], # one band name per channel, e.g. ("B2", "B3", ...) - scale_m: int | None = None, # optional nominal pixel size, recorded as provenance + 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 ) ``` +- **`spatial` is optional but supply it whenever you have it** — models whose forward pass conditions on geometry (lat/lon or GSD encodings: `clay`, `prithvi`) refuse declarations without it, because coordinates are never fabricated. All other models accept ungeoreferenced data; they just lose location provenance in the metadata. `list_models_for_data` on a spatial-less declaration reports which models refuse for this reason. +- **`temporal` travels with the data** — it is the acquisition time of *this* imagery, so it lives here rather than on the API call. Models that condition on time read it; omitting it falls back to the package default window. +- **`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. + Collection aliases: `"s2"` / `"sentinel-2"` / `"s2-l2a"` → `COPERNICUS/S2_SR_HARMONIZED`; `"s1"` / `"sentinel-1"` → `COPERNICUS/S1_GRD`. Full collection ids pass through unchanged. Band aliases resolve the same way provider fetches resolve them (`"RED"` → `"B4"`, `"NIR_NARROW"` → `"B8A"`, …). Multi-frame `[T,C,H,W]` arrays are only meaningful for time-series models (galileo, prithvi, olmoearth, anysat, agrifm); single-frame models reject them. @@ -43,29 +52,35 @@ Multi-frame `[T,C,H,W]` arrays are only meaningful for time-series models (galil ### get_embedding_from_data ```python +import numpy as np from rs_embed import UserData, get_embedding_from_data from rs_embed.core.specs import PointBuffer, TemporalSpec -emb = get_embedding_from_data( - "galileo", - data=UserData(data=cube, collection="s2", bands=bands), +data = UserData( + data=cube, # [12, H, W] raw S2 L2A DN + collection="s2", spatial=PointBuffer(lon=-88.2, lat=40.1, buffer_m=640), temporal=TemporalSpec.year(2022), ) +emb = get_embedding_from_data("galileo", data) ``` -Returns one `Embedding`; `meta["user_input"]` records the declaration and the channel selection actually fed to the model (`bands_used`, `channel_indices`). +Returns one `Embedding`; `meta["user_input"]` records the declaration and the channel selection actually fed to the model (`declared_bands`, `bands_used`, `channel_indices`). ### get_embeddings_batch_from_data -Batch counterpart: `datas: list[UserData]` aligned with `spatials: list[SpatialSpec]`. Each item is matched independently, so items may declare different band orders. +```python +embs = get_embeddings_batch_from_data("galileo", datas) # datas: list[UserData] +``` + +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. ### list_models_for_data ```python from rs_embed import list_models_for_data -report = list_models_for_data(my_declaration) +report = list_models_for_data(data) [r["model"] for r in report if r["compatible"]] ``` @@ -82,4 +97,6 @@ Runs the same matching against every catalog model without loading weights. Each | 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/src/rs_embed/api.py b/src/rs_embed/api.py index 7f002f0..796a3ff 100644 --- a/src/rs_embed/api.py +++ b/src/rs_embed/api.py @@ -119,6 +119,9 @@ 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, ) @@ -443,24 +446,43 @@ def _resolve_user_data_sensor(model_n: str, *, modality: str | None) -> SensorSp 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) - arrays.append(_select_prefetched_channels(np.asarray(data.data, dtype=np.float32), idx)) + # 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(data.bands), + "declared_bands": list(_resolve_declared_bands(data)), "bands_used": list(sensor_eff.bands), "channel_indices": list(idx), "declared_scale_m": data.scale_m, @@ -471,10 +493,8 @@ def _prepare_user_data_inputs( def get_embedding_from_data( model: str, - *, data: UserData, - spatial: SpatialSpec, - temporal: TemporalSpec | None = None, + *, modality: str | None = None, output: OutputSpec = OutputSpec.pooled(), device: str = "auto", @@ -483,13 +503,15 @@ def get_embedding_from_data( """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 supplies a - :class:`~rs_embed.UserData` declaring what the pixels are (collection + - band names, raw provider units). 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. + 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 ---------- @@ -497,15 +519,10 @@ def get_embedding_from_data( Model identifier or alias. Precomputed models are refused (they have no imagery input). data : UserData - The imagery and its declaration. ``data.data`` must be ``[C,H,W]`` - (or ``[T,C,H,W]`` for time-series models) with raw provider values in - the declared band order. - spatial : SpatialSpec - Where the imagery is located. Required: several models condition on - geometry (e.g. lat/lon or GSD embeddings), and embedding metadata - records it. - temporal : TemporalSpec or None - When the imagery was acquired, for models that condition on time. + 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. @@ -528,22 +545,21 @@ def get_embedding_from_data( If the model cannot take user data (precomputed / no input sensor) or the declaration does not satisfy the model's sensor. SpecError - If *data*, *spatial*, or *temporal* fail validation. + If the declaration fails validation. Examples -------- - >>> emb = get_embedding_from_data( - ... "galileo", - ... data=UserData(data=chw, collection="s2", bands=("B2", "B3", "B4", ...)), + >>> data = UserData( + ... data=chw, + ... collection="s2", ... 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, - datas=[data], - spatials=[spatial], - temporal=temporal, + [data], modality=modality, output=output, device=device, @@ -553,10 +569,8 @@ def get_embedding_from_data( def get_embeddings_batch_from_data( model: str, - *, datas: list[UserData], - spatials: list[SpatialSpec], - temporal: TemporalSpec | None = None, + *, modality: str | None = None, output: OutputSpec = OutputSpec.pooled(), device: str = "auto", @@ -565,19 +579,17 @@ def get_embeddings_batch_from_data( """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 as long - as every declaration satisfies the model. + 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] - User imagery declarations, one per item. - spatials : list[SpatialSpec] - Locations aligned with *datas*. - temporal : TemporalSpec or None - Shared acquisition-time filter. + Registered imagery, one per item. modality : str or None Optional modality selector. output : OutputSpec @@ -595,31 +607,44 @@ def get_embeddings_batch_from_data( Raises ------ ModelError - If lengths mismatch, the model cannot take user data, or any + If *datas* is empty, the model cannot take user data, or any declaration does not satisfy the model's sensor. SpecError - If any spec fails validation. + 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].") - if len(datas) != len(spatials): - raise ModelError(f"datas/spatials length mismatch: {len(datas)} != {len(spatials)}") - _validate_spatial_list(spatials=spatials, temporal=temporal, output=output) sensor_eff = _resolve_user_data_sensor(model_n, modality=modality) - arrays, metas = _prepare_user_data_inputs(model_n, datas, sensor_eff) - return _run_user_input_request( - model_n=model_n, - spatials=spatials, - input_arrays=arrays, - temporal=temporal, - sensor=sensor_eff, - model_config=model_config, - output=output, - device=device, - input_metas=metas, - ) + 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], + ) + 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]]: @@ -658,6 +683,7 @@ def list_models_for_data(data: UserData) -> list[dict[str, Any]]: } 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) diff --git a/src/rs_embed/core/types.py b/src/rs_embed/core/types.py index bbebab6..17ad113 100644 --- a/src/rs_embed/core/types.py +++ b/src/rs_embed/core/types.py @@ -13,7 +13,7 @@ import numpy as np from .errors import SpecError -from .specs import FetchSpec, InputPrepSpec, SensorSpec +from .specs import FetchSpec, InputPrepSpec, SensorSpec, SpatialSpec, TemporalSpec # ── Enums ────────────────────────────────────────────────────────── @@ -148,10 +148,13 @@ class FetchResult: class UserData: """User-provided imagery with a declaration of what the pixels are. - The bring-your-own-data entrypoints (:func:`rs_embed.get_embedding_from_data` - and friends) match this declaration against a model's input sensor and - refuse the request when the data cannot satisfy it — so the declaration, - not the array shape, is the contract. + 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 @@ -166,21 +169,38 @@ class UserData: collection : str Provider collection the pixels came from, e.g. ``"COPERNICUS/S2_SR_HARMONIZED"`` or a short alias like ``"s2"``. - bands : tuple[str, ...] + 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 - bands: tuple[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 @@ -189,20 +209,26 @@ def validate(self) -> None: """ if not str(self.collection or "").strip(): raise SpecError("UserData.collection must be a non-empty collection id or alias.") - if not self.bands 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.") + 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)}." ) - 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.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.") diff --git a/src/rs_embed/tools/user_data.py b/src/rs_embed/tools/user_data.py index b0f2ffc..5062a06 100644 --- a/src/rs_embed/tools/user_data.py +++ b/src/rs_embed/tools/user_data.py @@ -20,7 +20,7 @@ import numpy as np -from ..core.errors import ModelError +from ..core.errors import ModelError, SpecError from ..core.specs import SensorSpec from ..core.types import UserData from ..providers.gee_utils import resolve_band_aliases @@ -43,6 +43,26 @@ # 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.""" @@ -57,6 +77,42 @@ def canonical_band_names(collection_id: str, bands: tuple[str, ...]) -> tuple[st 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, @@ -87,6 +143,9 @@ def match_user_data_to_sensor( 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) @@ -98,7 +157,7 @@ def match_user_data_to_sensor( "this data cannot serve the model." ) - user_bands = canonical_band_names(user_collection, tuple(data.bands)) + 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( diff --git a/tests/test_api_from_data.py b/tests/test_api_from_data.py index 0b92516..e3b8591 100644 --- a/tests/test_api_from_data.py +++ b/tests/test_api_from_data.py @@ -26,12 +26,13 @@ class _MockFromDataEmbedder(EmbedderBase): - """Captures the input_chw it receives; no I/O.""" + """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 { @@ -56,9 +57,10 @@ def get_embedding( 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}, + meta={"model": self.model_name, "output": output.mode, "temporal": temporal}, ) @@ -70,35 +72,45 @@ def describe(self): return {"type": "precomputed", "backend": ["local"], "output": ["pooled"]} +class _MockGeorefFromDataEmbedder(_MockFromDataEmbedder): + model_name = "mock_from_data_georef" + _requires_georef = True + + @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) _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)): +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", bands=bands) + 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", - data=_twelve_band_userdata(), - spatial=_POINT, - temporal=TemporalSpec.year(2022), - ) + 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 @@ -106,12 +118,21 @@ def test_superset_data_is_sliced_to_model_band_order(): 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", - data=_twelve_band_userdata(), - spatial=_POINT, - ) + 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 @@ -123,56 +144,65 @@ 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=data, spatial=_POINT) + 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=data, spatial=_POINT) + 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", - data=_twelve_band_userdata(), - spatial=_POINT, - ) + 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", data=bad, spatial=_POINT) + get_embedding_from_data("mock_from_data", bad) def test_model_kwargs_are_forwarded_as_model_config(): - get_embedding_from_data( - "mock_from_data", - data=_twelve_band_userdata(), - spatial=_POINT, - variant="large", - ) + get_embedding_from_data("mock_from_data", _twelve_band_userdata(), variant="large") assert _MockFromDataEmbedder.last_model_config == {"variant": "large"} -def test_grid_output_mode_is_passed_through(): +def test_output_mode_is_passed_through(): emb = get_embedding_from_data( "mock_from_data", - data=_twelve_band_userdata(), - spatial=_POINT, + _twelve_band_userdata(), output=OutputSpec.pooled(), ) assert emb.meta["output"] == "pooled" @@ -183,27 +213,27 @@ def test_grid_output_mode_is_passed_through(): 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=datas, - spatials=[_POINT, _POINT], - ) + 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_length_mismatch_refuses(): - with pytest.raises(ModelError, match="length mismatch"): - get_embeddings_batch_from_data( - "mock_from_data", - datas=[_twelve_band_userdata()], - spatials=[_POINT, _POINT], - ) +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", datas=[], spatials=[]) + get_embeddings_batch_from_data("mock_from_data", []) # ── list_models_for_data over the real catalog ───────────────────── @@ -227,10 +257,20 @@ def test_list_models_for_data_over_catalog_with_12_band_s2(): 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) diff --git a/tests/test_user_data_matching.py b/tests/test_user_data_matching.py index b909d51..5a762bc 100644 --- a/tests/test_user_data_matching.py +++ b/tests/test_user_data_matching.py @@ -8,17 +8,19 @@ import pytest from rs_embed.core.errors import ModelError, SpecError -from rs_embed.core.specs import SensorSpec +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): @@ -29,7 +31,7 @@ def _user_data(bands, *, collection=S2, shape_hw=(4, 4), tchw=False, fill=None): 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, bands=tuple(bands)) + return UserData(data=data, collection=collection, spatial=_POINT, bands=tuple(bands)) def _sensor(bands, *, collection=S2): @@ -45,20 +47,39 @@ def test_validate_accepts_chw_and_tchw(): def test_validate_rejects_channel_band_mismatch(): - bad = UserData(data=np.zeros((3, 4, 4), dtype=np.float32), collection=S2, bands=("B2", "B3")) + 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, bands=("B2",)).validate() + 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=" ", bands=("B2",)).validate() + 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, bands=()).validate() + 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, bands=("B2",), scale_m=0).validate() + UserData( + data=np.zeros((1, 4, 4)), collection=S2, spatial=_POINT, bands=("B2",), scale_m=0 + ).validate() # ── collection + band normalization ─────────────────────────────── @@ -78,6 +99,43 @@ def test_canonical_band_names_resolve_aliases_and_case(): 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 ─────────────────────────────────────────────────────── @@ -137,6 +195,7 @@ 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"): From 87b138c1f1d71628517f100e932ca341a4720995 Mon Sep 17 00:00:00 2001 From: Dinghye <1269226384@qq.com> Date: Fri, 21 Aug 2026 15:26:21 -0500 Subject: [PATCH 06/12] fix: clay batch prefetched-input path no longer acquires a provider --- src/rs_embed/embedders/onthefly_clay.py | 3 ++- tests/test_clay_embedder.py | 27 ++++++++++++++++++++++++- 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/src/rs_embed/embedders/onthefly_clay.py b/src/rs_embed/embedders/onthefly_clay.py index f1b40ac..34bb271 100644 --- a/src/rs_embed/embedders/onthefly_clay.py +++ b/src/rs_embed/embedders/onthefly_clay.py @@ -922,7 +922,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/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) From 71a4e3531ab025c34409c5441f619e560e97eee7 Mon Sep 17 00:00:00 2001 From: Dinghye <1269226384@qq.com> Date: Fri, 21 Aug 2026 15:26:54 -0500 Subject: [PATCH 07/12] feat: batch_size cap on user-data batch dispatch + input-size docs --- CHANGELOG.md | 6 +++++- docs/user_data.md | 16 +++++++++++++- src/rs_embed/api.py | 13 ++++++++++-- src/rs_embed/tools/runtime.py | 40 +++++++++++++++++++++++------------ tests/test_api_from_data.py | 31 +++++++++++++++++++++++++++ 5 files changed, 89 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 69db03d..f16a88f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,9 +8,13 @@ The format is based on Keep a Changelog, and the project follows Semantic Versio ## [Unreleased] +### Fixed + +- **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 items carry their own spatial/temporal; same-temporal items dispatch together, results return in input order). 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). +- **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). 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/user_data.md b/docs/user_data.md index c6db39b..38a3128 100644 --- a/docs/user_data.md +++ b/docs/user_data.md @@ -70,11 +70,13 @@ Returns one `Embedding`; `meta["user_input"]` records the declaration and the ch ### get_embeddings_batch_from_data ```python -embs = get_embeddings_batch_from_data("galileo", datas) # datas: list[UserData] +embs = get_embeddings_batch_from_data("galileo", datas, batch_size=8) # datas: list[UserData] ``` 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. +`batch_size` caps how many items reach one model forward batch — set a small value to fit 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 @@ -88,6 +90,18 @@ Runs the same matching against every catalog model without loading weights. Each --- +## Input size handling + +User arrays are fed to the model as-is spatially: each embedder resizes them to its fixed input size (224 for most image-level ViTs, 256 for clay; prithvi can pad instead via `RS_EMBED_PRITHVI_PREP=pad`). Concretely: + +- **Much larger than the model input** (e.g. 2048×2048): downsampled in one step — fine detail is lost. The provider-fetch path's `input_prep="tile"` machinery (tile at native resolution + stitch grids) does **not** apply to user data yet; if you need native-resolution detail over a large scene, split it into patches yourself and pass them as separate items. +- **Much smaller** (e.g. 16×16): upsampled to the model input size. It runs, but the information content is what your pixels carry — expect weak embeddings below roughly half the model's input size. +- **Non-square**: plain resize distorts the aspect ratio. The fetch path protects square-input models (clay, prithvi) by fetching enlarged squares and cropping back; user data has no such protection, so provide near-square patches for best fidelity. + +Rule of thumb: patches near the model's native input size at its native scale (~224–256 px at 10 m for the S2 models) are the sweet spot. + +--- + ## Refusal semantics | Situation | Result | diff --git a/src/rs_embed/api.py b/src/rs_embed/api.py index 796a3ff..f3b732e 100644 --- a/src/rs_embed/api.py +++ b/src/rs_embed/api.py @@ -574,6 +574,7 @@ def get_embeddings_batch_from_data( modality: str | None = None, output: OutputSpec = OutputSpec.pooled(), device: str = "auto", + batch_size: int | None = None, **model_kwargs: Any, ) -> list[Embedding]: """Compute embeddings for multiple user-provided inputs. @@ -596,6 +597,12 @@ def get_embeddings_batch_from_data( 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). **model_kwargs Model-specific settings, as in :func:`get_embedding`. @@ -607,8 +614,9 @@ def get_embeddings_batch_from_data( Raises ------ ModelError - If *datas* is empty, the model cannot take user data, or any - declaration does not satisfy the model's sensor. + 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. """ @@ -638,6 +646,7 @@ def get_embeddings_batch_from_data( output=output, device=device, input_metas=[metas[i] for i in indices], + batch_size=batch_size, ) for i, emb in zip(indices, embs, strict=True): results[i] = emb diff --git a/src/rs_embed/tools/runtime.py b/src/rs_embed/tools/runtime.py index 8902039..f0e7db2 100644 --- a/src/rs_embed/tools/runtime.py +++ b/src/rs_embed/tools/runtime.py @@ -815,6 +815,7 @@ def run_user_input_request( output: OutputSpec, device: str, input_metas: list[dict[str, Any]] | None = None, + batch_size: int | None = None, ) -> list[Embedding]: """Run an embedding request over user-provided inputs (no provider fetch). @@ -846,6 +847,12 @@ def run_user_input_request( 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). Returns ------- @@ -855,12 +862,15 @@ def run_user_input_request( Raises ------ ModelError - If lengths mismatch or the embedder cannot take prefetched inputs. + 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)): @@ -870,24 +880,28 @@ def run_user_input_request( ) assert_supported(embedder, backend="auto", output=output, temporal=temporal) - kwargs: dict[str, Any] = { - "spatials": spatials, - "input_chws": input_arrays, - "temporal": temporal, - "sensor": sensor, - "output": output, - "backend": "auto", - "device": device_n, - } if model_config is not None: require_model_config_support( embedder=embedder, model_config=model_config, method_name="get_embeddings_batch_from_inputs", ) - kwargs["model_config"] = model_config - with lock: - embs = embedder.get_embeddings_batch_from_inputs(**kwargs) + step = int(batch_size) if batch_size is not None else len(spatials) + embs: list[Embedding] = [] + for s0 in range(0, len(spatials), step): + kwargs: dict[str, Any] = { + "spatials": spatials[s0 : s0 + step], + "input_chws": input_arrays[s0 : s0 + step], + "temporal": temporal, + "sensor": sensor, + "output": output, + "backend": "auto", + "device": device_n, + } + if model_config is not None: + kwargs["model_config"] = model_config + with lock: + embs.extend(embedder.get_embeddings_batch_from_inputs(**kwargs)) embs = [normalize_embedding_output(emb=emb, output=output) for emb in embs] if input_metas is not None: for emb, item_meta in zip(embs, input_metas, strict=False): diff --git a/tests/test_api_from_data.py b/tests/test_api_from_data.py index e3b8591..b31c379 100644 --- a/tests/test_api_from_data.py +++ b/tests/test_api_from_data.py @@ -77,11 +77,24 @@ class _MockGeorefFromDataEmbedder(_MockFromDataEmbedder): _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 + ) + + @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) + _MockChunkTrackingEmbedder.seen_chunk_sizes = [] _MockFromDataEmbedder.last_input = None _MockFromDataEmbedder.last_sensor = None _MockFromDataEmbedder.last_model_config = None @@ -236,6 +249,24 @@ def test_batch_empty_refuses(): 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) + + # ── list_models_for_data over the real catalog ───────────────────── From 7dccbe6ac5b69c3ecfa68d12d434494ed5a84d2c Mon Sep 17 00:00:00 2001 From: Dinghye <1269226384@qq.com> Date: Fri, 21 Aug 2026 15:50:18 -0500 Subject: [PATCH 08/12] docs: model input sizes in overview+complete user-data example --- docs/models.md | 40 +++++++++++++++-------------- docs/user_data.md | 61 +++++++++++++++++++++++++++++++++++++++------ src/rs_embed/api.py | 5 ++-- 3 files changed, 78 insertions(+), 28 deletions(-) diff --git a/docs/models.md b/docs/models.md index 9f61052..2fa041c 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: inputs are resized to it before the forward pass (in the default `input_prep="tile"` fetch mode, large ROIs are instead cut into tiles of this size at native resolution and the grids stitched; user-provided data is always resized — see [User Data API](user_data.md)). 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 and manages its own tiling; 256 is its training tile size. `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 index 38a3128..9e9d7a0 100644 --- a/docs/user_data.md +++ b/docs/user_data.md @@ -51,26 +51,73 @@ Multi-frame `[T,C,H,W]` arrays are only meaningful for time-series models (galil ### 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 `B1..B8, B8A, B9, B11, B12`, raw surface-reflectance DN (`0..10000`, i.e. exactly as downloaded, not rescaled to `0..1`): + ```python -import numpy as np +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 PointBuffer, TemporalSpec +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=cube, # [12, H, W] raw S2 L2A DN - collection="s2", - spatial=PointBuffer(lon=-88.2, lat=40.1, buffer_m=640), - temporal=TemporalSpec.year(2022), + data=pixels, # [C,H,W], raw provider values + collection="s2", # which sensor product this is + spatial=BBox(minlon=left, minlat=bottom, maxlon=right, maxlat=top), + temporal=TemporalSpec.range("2022-06-01", "2022-09-01"), # acquisition window + # bands= omitted: 12 channels in canonical S2 order is the documented default. + # If your file has other bands or another order, declare them explicitly: + # bands=("B4", "B3", "B2") for an RGB-only file, etc. ) + +# 3. Embed — just name the model. Band selection is automatic: galileo slices +# out its 10 bands, an RGB model would slice B4/B3/B2, all from this one +# declaration. 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`). ### get_embeddings_batch_from_data +Continuing the example above — a whole directory of patches into one feature matrix: + ```python -embs = get_embeddings_batch_from_data("galileo", datas, batch_size=8) # datas: list[UserData] +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. diff --git a/src/rs_embed/api.py b/src/rs_embed/api.py index f3b732e..80af64e 100644 --- a/src/rs_embed/api.py +++ b/src/rs_embed/api.py @@ -549,9 +549,10 @@ def get_embedding_from_data( Examples -------- + >>> pixels = np.load("s2_patch.npy") # [12, H, W] raw S2 L2A DN (0..10000) >>> data = UserData( - ... data=chw, - ... collection="s2", + ... 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), ... ) From eaf13be47306d186746bbc5e3571fca74fd1d562 Mon Sep 17 00:00:00 2001 From: Dinghye <1269226384@qq.com> Date: Fri, 21 Aug 2026 16:06:13 -0500 Subject: [PATCH 09/12] feat: user-data path defaults to input_prep=tile with resize opt-out --- CHANGELOG.md | 2 +- docs/models.md | 2 +- docs/user_data.md | 10 ++--- src/rs_embed/api.py | 25 ++++++++++- src/rs_embed/tools/runtime.py | 81 +++++++++++++++++++++++++++++++---- tests/test_api_from_data.py | 58 +++++++++++++++++++++++++ 6 files changed, 161 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f16a88f..cb898fd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,7 +14,7 @@ The format is based on Keep a Changelog, and the project follows Semantic Versio ### 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). 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). +- **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. 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 2fa041c..1aaa929 100644 --- a/docs/models.md +++ b/docs/models.md @@ -59,7 +59,7 @@ Some detail-page filenames still use older names for compatibility, but the cano | `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: inputs are resized to it before the forward pass (in the default `input_prep="tile"` fetch mode, large ROIs are instead cut into tiles of this size at native resolution and the grids stitched; user-provided data is always resized — see [User Data API](user_data.md)). 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 and manages its own tiling; 256 is its training tile size. `anysat` and `prithvi` sizes are env-tunable (`RS_EMBED_ANYSAT_IMG`, `RS_EMBED_PRITHVI_IMG`). +**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 and manages its own tiling; 256 is its training tile size. `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 index 9e9d7a0..eb9ee54 100644 --- a/docs/user_data.md +++ b/docs/user_data.md @@ -139,13 +139,13 @@ Runs the same matching against every catalog model without loading weights. Each ## Input size handling -User arrays are fed to the model as-is spatially: each embedder resizes them to its fixed input size (224 for most image-level ViTs, 256 for clay; prithvi can pad instead via `RS_EMBED_PRITHVI_PREP=pad`). Concretely: +User data follows the package-wide `input_prep` policy, defaulting to **`"tile"`** — the same fairness semantics as the provider-fetch path: -- **Much larger than the model input** (e.g. 2048×2048): downsampled in one step — fine detail is lost. The provider-fetch path's `input_prep="tile"` machinery (tile at native resolution + stitch grids) does **not** apply to user data yet; if you need native-resolution detail over a large scene, split it into patches yourself and pass them as separate items. -- **Much smaller** (e.g. 16×16): upsampled to the model input size. It runs, but the information content is what your pixels carry — expect weak embeddings below roughly half the model's input size. -- **Non-square**: plain resize distorts the aspect ratio. The fetch path protects square-input models (clay, prithvi) by fetching enlarged squares and cropping back; user data has no such protection, so provide near-square patches for best fidelity. +- **Larger than the model's input size**: 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"` to opt into one-step downsampling instead (faster, lossy; `meta["input_prep"]` records which path ran). +- **At or below the model's input size**: nothing to tile — the array goes straight to the embedder, which resizes up if needed (prithvi can pad instead via `RS_EMBED_PRITHVI_PREP=pad`). Very small arrays (e.g. 16×16 into a 224 model) 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 rectangles cleanly (edge tiles are padded, outputs cropped back). Under `"resize"`, a plain resize distorts the aspect ratio — another reason to keep the tile default for non-square patches. -Rule of thumb: patches near the model's native input size at its native scale (~224–256 px at 10 m for the S2 models) are the sweet spot. +Rule of thumb: any patch at the model's native scale (10 m for the S2 models) is handled faithfully under the tile default; patches at or below the model's input size are also the cheapest (single forward pass). --- diff --git a/src/rs_embed/api.py b/src/rs_embed/api.py index 80af64e..0cea3ea 100644 --- a/src/rs_embed/api.py +++ b/src/rs_embed/api.py @@ -498,6 +498,7 @@ def get_embedding_from_data( 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). @@ -530,6 +531,15 @@ def get_embedding_from_data( 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`. @@ -537,7 +547,8 @@ def get_embedding_from_data( ------- Embedding Normalized embedding; ``meta['user_input']`` records the declaration - and the channel selection that was fed to the model. + and the channel selection that was fed to the model, and + ``meta['input_prep']`` records how the array was prepared. Raises ------ @@ -564,6 +575,7 @@ def get_embedding_from_data( modality=modality, output=output, device=device, + input_prep=input_prep, **model_kwargs, )[0] @@ -576,6 +588,7 @@ def get_embeddings_batch_from_data( 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. @@ -603,7 +616,14 @@ def get_embeddings_batch_from_data( 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). + ``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`. @@ -648,6 +668,7 @@ def get_embeddings_batch_from_data( 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 diff --git a/src/rs_embed/tools/runtime.py b/src/rs_embed/tools/runtime.py index f0e7db2..090512d 100644 --- a/src/rs_embed/tools/runtime.py +++ b/src/rs_embed/tools/runtime.py @@ -816,6 +816,7 @@ def run_user_input_request( 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). @@ -826,6 +827,13 @@ def run_user_input_request( 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 @@ -852,7 +860,12 @@ def run_user_input_request( 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). + 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 ------- @@ -886,12 +899,41 @@ def run_user_input_request( model_config=model_config, method_name="get_embeddings_batch_from_inputs", ) - step = int(batch_size) if batch_size is not None else len(spatials) - embs: list[Embedding] = [] - for s0 in range(0, len(spatials), step): + + 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: arrays larger than the model's tile size go through the + # shared tiler (native resolution preserved); everything else — resize + # mode, no known tile size, or an array a single tile already covers — + # keeps the efficient batched dispatch (there is nothing to tile). + tiled_set: set[int] = set() + if str(getattr(input_prep_resolved, "mode", "tile")) != "resize": + params = _resolve_tile_params(embedder, input_prep_resolved) + if params.tile_size > 0: + for i, x in enumerate(input_arrays): + if int(x.shape[-2]) > params.tile_size or int(x.shape[-1]) > params.tile_size: + tiled_set.add(i) + direct_indices = [i for i in range(len(input_arrays)) if i not in tiled_set] + + results: list[Embedding | None] = [None] * len(input_arrays) + + step = int(batch_size) if batch_size is not None else max(1, len(direct_indices)) + for s0 in range(0, len(direct_indices), step): + chunk = direct_indices[s0 : s0 + step] kwargs: dict[str, Any] = { - "spatials": spatials[s0 : s0 + step], - "input_chws": input_arrays[s0 : s0 + step], + "spatials": [spatials[i] for i in chunk], + "input_chws": [input_arrays[i] for i in chunk], "temporal": temporal, "sensor": sensor, "output": output, @@ -901,8 +943,31 @@ def run_user_input_request( if model_config is not None: kwargs["model_config"] = model_config with lock: - embs.extend(embedder.get_embeddings_batch_from_inputs(**kwargs)) - embs = [normalize_embedding_output(emb=emb, output=output) for emb in embs] + 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="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): diff --git a/tests/test_api_from_data.py b/tests/test_api_from_data.py index b31c379..c80332f 100644 --- a/tests/test_api_from_data.py +++ b/tests/test_api_from_data.py @@ -88,13 +88,41 @@ def get_embeddings_batch_from_inputs(self, *, spatials, 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 + ) + + @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) _MockChunkTrackingEmbedder.seen_chunk_sizes = [] + _MockTiledFromDataEmbedder.seen_input_shapes = [] + _MockTiledFromDataEmbedder.from_inputs_calls = 0 _MockFromDataEmbedder.last_input = None _MockFromDataEmbedder.last_sensor = None _MockFromDataEmbedder.last_model_config = None @@ -267,6 +295,36 @@ def test_invalid_batch_size_refuses(): 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) + + # ── list_models_for_data over the real catalog ───────────────────── From d19efca12a731a3a6d0808eb5ccfb6bd49dd0a11 Mon Sep 17 00:00:00 2001 From: Dinghye <1269226384@qq.com> Date: Fri, 21 Aug 2026 16:51:43 -0500 Subject: [PATCH 10/12] feat: flexible-size models consume user data natively by default --- CHANGELOG.md | 2 +- docs/models.md | 2 +- docs/user_data.md | 1 + src/rs_embed/embedders/base.py | 19 ++++ src/rs_embed/embedders/onthefly_olmoearth.py | 35 ++++++ src/rs_embed/pipelines/inference.py | 2 +- src/rs_embed/tools/runtime.py | 112 ++++++++++++++----- src/rs_embed/tools/tiling.py | 26 ++++- tests/test_api_from_data.py | 54 +++++++++ tests/test_capabilities_contract.py | 19 +++- tests/test_olmoearth.py | 39 +++++++ 11 files changed, 269 insertions(+), 42 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cb898fd..49c9ee7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,7 +14,7 @@ The format is based on Keep a Changelog, and the project follows Semantic Versio ### 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. 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). +- **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 1aaa929..24a7d3e 100644 --- a/docs/models.md +++ b/docs/models.md @@ -59,7 +59,7 @@ Some detail-page filenames still use older names for compatibility, but the cano | `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 and manages its own tiling; 256 is its training tile size. `anysat` and `prithvi` sizes are env-tunable (`RS_EMBED_ANYSAT_IMG`, `RS_EMBED_PRITHVI_IMG`). +**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 index eb9ee54..26dd90b 100644 --- a/docs/user_data.md +++ b/docs/user_data.md @@ -144,6 +144,7 @@ User data follows the package-wide `input_prep` policy, defaulting to **`"tile"` - **Larger than the model's input size**: 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"` to opt into one-step downsampling instead (faster, lossy; `meta["input_prep"]` records which path ran). - **At or below the model's input size**: nothing to tile — the array goes straight to the embedder, which resizes up if needed (prithvi can pad instead via `RS_EMBED_PRITHVI_PREP=pad`). Very small arrays (e.g. 16×16 into a 224 model) 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 rectangles cleanly (edge tiles are padded, outputs cropped back). Under `"resize"`, a plain resize distorts the aspect ratio — another reason to keep the tile default for non-square patches. +- **Flexible-size models**: `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). A 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"`. Passing an explicit `image_size` in model kwargs restores fixed-size behavior (larger inputs tile at it). Rule of thumb: any patch at the model's native scale (10 m for the S2 models) is handled faithfully under the tile default; patches at or below the model's input size are also the cheapest (single forward pass). diff --git a/src/rs_embed/embedders/base.py b/src/rs_embed/embedders/base.py index d1aebdc..69d8b06 100644 --- a/src/rs_embed/embedders/base.py +++ b/src/rs_embed/embedders/base.py @@ -201,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_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/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/tools/runtime.py b/src/rs_embed/tools/runtime.py index 090512d..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": @@ -913,42 +918,89 @@ def run_user_input_request( _model_policy, ) = resolve_model_aware_input_prep(model_n=model_n, input_prep=input_prep, output=output) - # Route per item: arrays larger than the model's tile size go through the - # shared tiler (native resolution preserved); everything else — resize - # mode, no known tile size, or an array a single tile already covers — - # keeps the efficient batched dispatch (there is nothing to tile). + # 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) - if params.tile_size > 0: - for i, x in enumerate(input_arrays): - if int(x.shape[-2]) > params.tile_size or int(x.shape[-1]) > params.tile_size: - tiled_set.add(i) - direct_indices = [i for i in range(len(input_arrays)) if i not in tiled_set] + 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) - step = int(batch_size) if batch_size is not None else max(1, len(direct_indices)) - for s0 in range(0, len(direct_indices), step): - chunk = direct_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 model_config is not None: - kwargs["model_config"] = model_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="resize" + # 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: 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/tests/test_api_from_data.py b/tests/test_api_from_data.py index c80332f..5f30f4e 100644 --- a/tests/test_api_from_data.py +++ b/tests/test_api_from_data.py @@ -113,6 +113,22 @@ def get_embeddings_batch_from_inputs(self, *, spatials, 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) @@ -120,9 +136,12 @@ def register_mocks(): 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 @@ -325,6 +344,41 @@ def test_tile_sized_inputs_keep_batched_dispatch_under_tile_default(): 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 ───────────────────── 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_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} From 0ceaf45e05f03b10c3e559f3bd549d7ed962e0c8 Mon Sep 17 00:00:00 2001 From: Dinghye <1269226384@qq.com> Date: Fri, 21 Aug 2026 16:58:12 -0500 Subject: [PATCH 11/12] docs:restyle --- docs/user_data.md | 193 ++++++++++++++++++++++++++++++++++------------ 1 file changed, 142 insertions(+), 51 deletions(-) diff --git a/docs/user_data.md b/docs/user_data.md index 26dd90b..64efb0a 100644 --- a/docs/user_data.md +++ b/docs/user_data.md @@ -1,24 +1,57 @@ # API: User-Provided Data -This page covers the bring-your-own-data API: computing embeddings from imagery you already have, instead of provider-fetched imagery. - -Related pages: [API: Embedding](api_embedding.md), [API: Specs and Data Structures](api_specs.md). +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. --- -## Concept - -The flow has two steps: - -1. **Register** each piece of imagery 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. -2. **Embed** by naming a model: `get_embedding_from_data("galileo", data)`. Nothing else is needed — the declaration already carries the full context. +## The two-step flow -Every on-the-fly model declares an input sensor (a collection plus an ordered band list). Your declaration is matched against it: - -- **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 can serve models that need 3, 6, or 10 of those bands. -- **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. +```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"] +``` -Values must be **raw provider units** for the declared collection (e.g. Sentinel-2 L2A surface-reflectance DN in `0..10000`), exactly what a provider fetch would return. Per-model normalization stays inside each embedder, so you never need to know a model's normalization. Data that looks already normalized (max ≤ 1.5 on an S2 declaration) triggers a warning. +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. --- @@ -28,22 +61,38 @@ Values must be **raw provider units** for the declared collection (e.g. Sentinel 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) + 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 + scale_m: int | None = None, # optional nominal pixel size, provenance only ) ``` -- **`spatial` is optional but supply it whenever you have it** — models whose forward pass conditions on geometry (lat/lon or GSD encodings: `clay`, `prithvi`) refuse declarations without it, because coordinates are never fabricated. All other models accept ungeoreferenced data; they just lose location provenance in the metadata. `list_models_for_data` on a spatial-less declaration reports which models refuse for this reason. -- **`temporal` travels with the data** — it is the acquisition time of *this* imagery, so it lives here rather than on the API call. Models that condition on time read it; omitting it falls back to the package default window. -- **`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. - -Collection aliases: `"s2"` / `"sentinel-2"` / `"s2-l2a"` → `COPERNICUS/S2_SR_HARMONIZED`; `"s1"` / `"sentinel-1"` → `COPERNICUS/S1_GRD`. Full collection ids pass through unchanged. Band aliases resolve the same way provider fetches resolve them (`"RED"` → `"B4"`, `"NIR_NARROW"` → `"B8A"`, …). - -Multi-frame `[T,C,H,W]` arrays are only meaningful for time-series models (galileo, prithvi, olmoearth, anysat, agrifm); single-frame models reject them. +**`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. --- @@ -51,7 +100,9 @@ Multi-frame `[T,C,H,W]` arrays are only meaningful for time-series models (galil ### 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 `B1..B8, B8A, B9, B11, B12`, raw surface-reflectance DN (`0..10000`, i.e. exactly as downloaded, not rescaled to `0..1`): +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 @@ -69,31 +120,32 @@ with rasterio.open("maize_field_2022.tif") as src: # 2. Register the imagery: the pixels plus everything that describes them. data = UserData( - data=pixels, # [C,H,W], raw provider values - collection="s2", # which sensor product this is + data=pixels, + collection="s2", spatial=BBox(minlon=left, minlat=bottom, maxlon=right, maxlat=top), - temporal=TemporalSpec.range("2022-06-01", "2022-09-01"), # acquisition window - # bands= omitted: 12 channels in canonical S2 order is the documented default. - # If your file has other bands or another order, declare them explicitly: - # bands=("B4", "B3", "B2") for an RGB-only file, etc. + temporal=TemporalSpec.range("2022-06-01", "2022-09-01"), ) -# 3. Embed — just name the model. Band selection is automatic: galileo slices -# out its 10 bands, an RGB model would slice B4/B3/B2, all from this one -# declaration. +# 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=`. +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`). +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: +Continuing the example above — a whole directory of patches into one feature +matrix: ```python from pathlib import Path @@ -120,9 +172,17 @@ 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. +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. -`batch_size` caps how many items reach one model forward batch — set a small value to fit 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. +!!! 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 @@ -133,32 +193,63 @@ 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). +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: - -- **Larger than the model's input size**: 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"` to opt into one-step downsampling instead (faster, lossy; `meta["input_prep"]` records which path ran). -- **At or below the model's input size**: nothing to tile — the array goes straight to the embedder, which resizes up if needed (prithvi can pad instead via `RS_EMBED_PRITHVI_PREP=pad`). Very small arrays (e.g. 16×16 into a 224 model) 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 rectangles cleanly (edge tiles are padded, outputs cropped back). Under `"resize"`, a plain resize distorts the aspect ratio — another reason to keep the tile default for non-square patches. -- **Flexible-size models**: `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). A 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"`. Passing an explicit `image_size` in model kwargs restores fixed-size behavior (larger inputs tile at it). - -Rule of thumb: any patch at the model's native scale (10 m for the S2 models) is handled faithfully under the tile default; patches at or below the model's input size are also the cheapest (single forward pass). +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) | +| 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 | From b7cdb2359c02cf5a0b2c413a9c84484082b363c0 Mon Sep 17 00:00:00 2001 From: Dinghye <1269226384@qq.com> Date: Fri, 21 Aug 2026 17:36:57 -0500 Subject: [PATCH 12/12] docs: unwrap hard line breaks --- docs/user_data.md | 140 ++++++++++------------------------------------ 1 file changed, 29 insertions(+), 111 deletions(-) diff --git a/docs/user_data.md b/docs/user_data.md index 64efb0a..ba3c59a 100644 --- a/docs/user_data.md +++ b/docs/user_data.md @@ -1,21 +1,11 @@ # 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. +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). +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. + 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. --- @@ -28,30 +18,14 @@ flowchart LR 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. +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"`, …). +**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. +**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. + 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. --- @@ -70,29 +44,14 @@ UserData( ) ``` -**`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. +**`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. + A `[T,C,H,W]` array is only meaningful for time-series models (`galileo`, `prithvi`, `olmoearth`, `anysat`, `agrifm`); single-frame models reject it. --- @@ -100,9 +59,7 @@ ungeoreferenced data and merely loses location provenance in the metadata. ### 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: +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 @@ -133,19 +90,13 @@ 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=`. +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)). +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: +Continuing the example above — a whole directory of patches into one feature matrix: ```python from pathlib import Path @@ -172,17 +123,10 @@ 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. +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. + `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 @@ -193,51 +137,25 @@ 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). +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. +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). + `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"`. + 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"`. ---