feat:anysat dense/tile mode support

Author: Dinghye

docs/assets/vis.png

@@ -1,6 +1,6 @@
 # AnySat (`anysat`)
 
-> Multi-frame Sentinel-2 time-series adapter that builds AnySat inputs (`s2` + `s2_dates`) from a temporal window and returns patch-grid features or pooled vectors.
+> Multi-frame Sentinel-2 time-series adapter that builds AnySat inputs (`s2` + `s2_dates`) from a temporal window and returns dense sub-patch grids by default or pooled vectors.
 
 ## Quick Facts
 
@@ -53,11 +53,13 @@ For `input_chw`, the adapter accepts either `CHW` or `TCHW` with `C=10`. A `CHW`
   <span class="pipeline-arrow">-&gt;</span> build AnySat side inputs
      <span class="pipeline-branch">s2:</span> [1,T,10,H,W]
      <span class="pipeline-branch">s2_dates:</span> [1,T] from frame-bin DOY midpoints
-  <span class="pipeline-arrow">-&gt;</span> forward with output="patch" and patch_size=sensor.scale_m
+  <span class="pipeline-arrow">-&gt;</span> forward with AnySat spatial output
+     <span class="pipeline-branch">grid path:</span> default `output="dense"` (`grid_feature_mode="dense"`)
+     <span class="pipeline-branch">pooled path:</span> default compatibility path keeps `output="patch"`; optional `pooled_source="tile"` uses native AnySat tile output
   <span class="pipeline-arrow">-&gt;</span> map [B,H,W,D] -&gt; rs-embed grid [D,H,W]
   <span class="pipeline-arrow">-&gt;</span> output projection
-     <span class="pipeline-branch">pooled:</span> spatial mean / max over grid
-     <span class="pipeline-branch">grid:</span>   model patch grid</code></pre>
+     <span class="pipeline-branch">pooled:</span> spatial mean / max over patch grid
+     <span class="pipeline-branch">grid:</span>   dense sub-patch grid by default (or patch grid when overridden)</code></pre>
 
 Important constraint:
 
@@ -73,6 +75,8 @@ Important constraint:
 | `RS_EMBED_ANYSAT_IMG` | `24` | Per-frame resize target (square) |
 | `RS_EMBED_ANYSAT_NORM` | `per_tile_zscore` | Series normalization mode |
 | `RS_EMBED_ANYSAT_MODEL_SIZE` | `base` | AnySat model size |
+| `RS_EMBED_ANYSAT_GRID_MODE` | `dense` | Grid path native AnySat spatial output (`dense` or `patch`) |
+| `RS_EMBED_ANYSAT_POOLED_SOURCE` | `patch` | Pooled path source (`patch` compatibility pooling or native `tile`) |
 | `RS_EMBED_ANYSAT_FLASH_ATTN` | `0` | Enable flash attention path if supported |
 | `RS_EMBED_ANYSAT_PRETRAINED` | `1` | Load pretrained checkpoint weights |
 | `RS_EMBED_ANYSAT_CKPT` | unset | Local checkpoint override |
@@ -85,7 +89,7 @@ Important constraint:
 
 ## Output Semantics
 
-AnySat follows the standard patch-grid pattern for multi-frame adapters. `pooled` applies spatial pooling over the patch grid, and `grid` returns `(D,H,W)` in model patch space rather than georeferenced raster pixels. The more distinctive AnySat details, such as frame packaging and `doy0_values`, are recorded in metadata rather than requiring a long per-page output section.
+AnySat now uses two spatial output paths inside the adapter. `pooled` defaults to the historical rs-embed behavior and applies spatial pooling over the AnySat `patch` grid, which preserves the previous pooled vector dimensionality; pass `pooled_source="tile"` (or set `RS_EMBED_ANYSAT_POOLED_SOURCE=tile`) to use the native AnySat tile embedding instead. `grid` defaults to AnySat `dense`, so the returned `(D,H,W)` is a denser sub-patch feature map by default; pass `grid_feature_mode="patch"` to the public API (or set `RS_EMBED_ANYSAT_GRID_MODE=patch`) to recover the older patch-grid behavior. As with other on-the-fly models, this grid is model space rather than guaranteed georeferenced raster pixels.
 
 ---
 

docs/models/anysat.md

@@ -48,7 +48,7 @@ Use this table for a first-pass side-by-side comparison of input assumptions and
 | `satmaepp` | `rshf.satmaepp.SatMAEPP` | 10m | S2 RGB (`B4,B3,B2`) | raw SR -> `/10000` -> RGB `uint8`; SatMAE++ fMoW eval preprocessing (`Normalize + Resize(short side) + CenterCrop`), default channel order `bgr` | default 224; source-aligned short-side resize + center crop; no pad | token sequence -> pooled or patch-token grid | High |
 | `satmaepp_s2_10b` | SatMAE++ grouped-channel source branch (`models_mae_group_channels.py`, `base` / `large` runtime families) | 10m | S2 SR 10-band (`B2,B3,B4,B5,B6,B7,B8,B8A,B11,B12`) | clip `0..10000`; source Sentinel min/max mapping to `uint8`; `ToTensor + Resize(short side) + CenterCrop` | default 96 with patch size 8; source-style resize/crop; no pad | grouped token sequence -> pooled or group-reduced spatial token grid | High |
 | `scalemae` | `rshf.scalemae.ScaleMAE` (ViT style) | 10m | S2 RGB (`B4,B3,B2`) + `input_res_m` | raw SR -> `/10000` -> RGB `uint8`; CLIP norm tensor; pass `input_res_m` | default 224; CLIP path has `Resize + CenterCrop`; no pad | token sequence or pooled vector depending on wrapper output | Medium |
-| `anysat` | AnySat from upstream `hubconf.py` (`AnySat`, `tiny` / `small` / `base`) | 10m | S2 10-band TCHW (or CHW auto-expanded) | clip to `0..10000`; normalize mode default `per_tile_zscore`; builds per-frame `s2_dates` | resize TCHW to default 24; no crop, no pad | patch output `[D,H,W]`, pooled by spatial mean/max | Medium |
+| `anysat` | AnySat from upstream `hubconf.py` (`AnySat`, `tiny` / `small` / `base`) | 10m | S2 10-band TCHW (or CHW auto-expanded) | clip to `0..10000`; normalize mode default `per_tile_zscore`; builds per-frame `s2_dates` | resize TCHW to default 24; no crop, no pad | grid defaults to dense sub-patch output `[D,H,W]`; pooled defaults to patch-grid mean/max, optional native tile vector | Medium |
 | `galileo` | `Encoder` from official `single_file_galileo.py` | 10m | S2 10-band TCHW (or CHW auto-expanded) | clip to `0..10000`; normalize mode default `unit_scale`; constructs Galileo tensors with configurable `T` + per-frame `months`, optional NDVI channel | default 64 with patch 8; bilinear resize; no pad | pooled token vector and S2-group token grid | Medium |
 | `wildsat` | WildSAT backbone + optional image head from checkpoint | 10m | S2 RGB CHW | clip to `0..10000` then `/10000`; default normalization `minmax`; convert to `uint8` then unit tensor | default 224; resize RGB; no pad | pooled branch output and optional grid (token or feature path) | Medium-Low |
 | `prithvi` | Vendored `PrithviMAE` runtime with HF checkpoints | 30m | S2 6-band (`BLUE,GREEN,RED,NIR_NARROW,SWIR_1,SWIR_2`) | raw SR -> `/10000` -> clamp `[0,1]`; prep mode from env | default mode `resize` to 224; optional `pad` to patch multiple (legacy) | token sequence -> pooled or patch-token grid | Medium |
@@ -126,7 +126,7 @@ This table only lists env vars that materially change model input construction o
 | `satmaepp` | `RS_EMBED_SATMAEPP_ID`, `RS_EMBED_SATMAEPP_IMG`, `RS_EMBED_SATMAEPP_CHANNEL_ORDER`, `RS_EMBED_SATMAEPP_BGR` |
 | `satmaepp_s2_10b` | `RS_EMBED_SATMAEPP_S2_CKPT_REPO`, `RS_EMBED_SATMAEPP_S2_CKPT_FILE`, `RS_EMBED_SATMAEPP_S2_MODEL_FN`, `RS_EMBED_SATMAEPP_S2_IMG`, `RS_EMBED_SATMAEPP_S2_PATCH`, `RS_EMBED_SATMAEPP_S2_GRID_REDUCE`, `RS_EMBED_SATMAEPP_S2_WEIGHTS_ONLY` |
 | `scalemae` | `RS_EMBED_SCALEMAE_IMG` |
-| `anysat` | `RS_EMBED_ANYSAT_IMG`, `RS_EMBED_ANYSAT_NORM`, `RS_EMBED_ANYSAT_FRAMES` |
+| `anysat` | `RS_EMBED_ANYSAT_IMG`, `RS_EMBED_ANYSAT_NORM`, `RS_EMBED_ANYSAT_FRAMES`, `RS_EMBED_ANYSAT_GRID_MODE`, `RS_EMBED_ANYSAT_POOLED_SOURCE` |
 | `galileo` | `RS_EMBED_GALILEO_IMG`, `RS_EMBED_GALILEO_PATCH`, `RS_EMBED_GALILEO_NORM`, `RS_EMBED_GALILEO_INCLUDE_NDVI`, `RS_EMBED_GALILEO_FRAMES`, `RS_EMBED_GALILEO_MONTH` |
 | `wildsat` | `RS_EMBED_WILDSAT_IMG`, `RS_EMBED_WILDSAT_NORM` |
 | `prithvi` | `RS_EMBED_PRITHVI_PREP`, `RS_EMBED_PRITHVI_IMG`, `RS_EMBED_PRITHVI_PATCH_MULT` |

docs/models_reference.md

@@ -656,7 +656,7 @@
     "    ],  # or you can spcify as [\"remoteclip\", \"terrafm\"]\n",
     "    target=ExportTarget.per_item(\"exports\", names=[\"p1\", \"p2\"]),\n",
     "    output=OutputSpec.grid(),\n",
-    "    config=ExportConfig(save_inputs=True, input_prep='tile'),\n",
+    "    config=ExportConfig(save_inputs=True, input_prep=\"tile\"),\n",
     "    backend=\"gee\",\n",
     ")"
    ]

examples/playground.ipynb

@@ -101,6 +101,38 @@ def _normalize_anysat_model_size(model_size: Any) -> str:
     return resolved
 
 
+def _normalize_anysat_grid_feature_mode(value: Any) -> str:
+    raw = str(value).strip().lower()
+    aliases = {
+        "dense": "dense",
+        "subpatch": "dense",
+        "subpatches": "dense",
+        "patch": "patch",
+        "patches": "patch",
+    }
+    resolved = aliases.get(raw)
+    if resolved is None:
+        raise ModelError(
+            f"Unknown AnySat grid_feature_mode={value!r} (expected one of: dense, patch)."
+        )
+    return resolved
+
+
+def _normalize_anysat_pooled_source(value: Any) -> str:
+    raw = str(value).strip().lower()
+    aliases = {
+        "patch": "patch",
+        "grid": "patch",
+        "tile": "tile",
+        "native": "tile",
+        "model": "tile",
+    }
+    resolved = aliases.get(raw)
+    if resolved is None:
+        raise ModelError(f"Unknown AnySat pooled_source={value!r} (expected one of: patch, tile).")
+    return resolved
+
+
 def _resolve_anysat_runtime_config(
     *,
     model_config: dict[str, Any] | None,
@@ -129,6 +161,16 @@ def _resolve_anysat_runtime_config(
 
     norm_mode = os.environ.get("RS_EMBED_ANYSAT_NORM", "per_tile_zscore").strip()
 
+    grid_feature_mode_v = model_config_value(model_config, "grid_feature_mode")
+    if grid_feature_mode_v is None:
+        grid_feature_mode_v = os.environ.get("RS_EMBED_ANYSAT_GRID_MODE", "dense")
+    grid_feature_mode = _normalize_anysat_grid_feature_mode(grid_feature_mode_v)
+
+    pooled_source_v = model_config_value(model_config, "pooled_source")
+    if pooled_source_v is None:
+        pooled_source_v = os.environ.get("RS_EMBED_ANYSAT_POOLED_SOURCE", "patch")
+    pooled_source = _normalize_anysat_pooled_source(pooled_source_v)
+
     ckpt_path = os.environ.get("RS_EMBED_ANYSAT_CKPT")
     ckpt_path = ckpt_path or None
 
@@ -156,6 +198,8 @@ def _resolve_anysat_runtime_config(
         "image_size": int(image_size),
         "n_frames": int(n_frames),
         "norm_mode": norm_mode,
+        "grid_feature_mode": grid_feature_mode,
+        "pooled_source": pooled_source,
         "ckpt_path": ckpt_path,
         "hf_repo": hf_repo,
         "hf_filename": hf_filename,
@@ -385,11 +429,12 @@ def _prepare_anysat_s2_input(
     }
 
 
-def _anysat_patch_features(
+def _anysat_spatial_features(
     model: Any,
     s2_input: dict[str, Any],
     *,
     patch_size_m: int,
+    feature_mode: str,
 ) -> tuple[np.ndarray, dict[str, Any]]:
     ensure_torch()
     import torch
@@ -399,27 +444,73 @@ def _anysat_patch_features(
             f"AnySat patch_size must be a positive multiple of 10 (meters), got {patch_size_m}"
         )
 
+    native_output = str(feature_mode).strip().lower()
+    if native_output not in {"patch", "dense"}:
+        raise ModelError(
+            f"Unknown AnySat feature_mode={feature_mode!r} (expected 'patch' or 'dense')."
+        )
+
     with torch.no_grad():
-        out = model(s2_input, patch_size=int(patch_size_m), output="patch")
+        out = model(s2_input, patch_size=int(patch_size_m), output=native_output)
 
     if not hasattr(out, "ndim") or int(out.ndim) != 4:
         raise ModelError(
-            f"AnySat output='patch' returned unexpected shape/type: {type(out)} {getattr(out, 'shape', None)}"
+            "AnySat output="
+            f"{native_output!r} returned unexpected shape/type: {type(out)} {getattr(out, 'shape', None)}"
         )
 
-    # AnySat patch output: [B,H,W,D]
+    # AnySat spatial outputs are [B,H,W,D] for both patch and dense.
     if int(out.shape[0]) != 1:
         raise ModelError(f"AnySat embedder expects B=1 per call, got {tuple(out.shape)}")
     arr = out[0].detach().float().cpu().numpy().astype(np.float32)  # [H,W,D]
     grid = np.transpose(arr, (2, 0, 1)).astype(np.float32)  # [D,H,W]
     meta = {
-        "patch_output_hw": (int(arr.shape[0]), int(arr.shape[1])),
+        "feature_mode": native_output,
+        "native_output_hw": (int(arr.shape[0]), int(arr.shape[1])),
         "feature_dim": int(arr.shape[2]),
         "patch_size_m": int(patch_size_m),
     }
+    if native_output == "patch":
+        meta["patch_output_hw"] = meta["native_output_hw"]
+    else:
+        meta["dense_output_hw"] = meta["native_output_hw"]
     return grid, meta
 
 
+def _anysat_tile_features(
+    model: Any,
+    s2_input: dict[str, Any],
+    *,
+    patch_size_m: int,
+) -> tuple[np.ndarray, dict[str, Any]]:
+    ensure_torch()
+    import torch
+
+    if patch_size_m <= 0 or (patch_size_m % 10) != 0:
+        raise ModelError(
+            f"AnySat patch_size must be a positive multiple of 10 (meters), got {patch_size_m}"
+        )
+
+    with torch.no_grad():
+        out = model(s2_input, patch_size=int(patch_size_m), output="tile")
+
+    if not hasattr(out, "ndim") or int(out.ndim) != 2:
+        raise ModelError(
+            "AnySat output='tile' returned unexpected shape/type: "
+            f"{type(out)} {getattr(out, 'shape', None)}"
+        )
+    if int(out.shape[0]) != 1:
+        raise ModelError(f"AnySat embedder expects B=1 per call, got {tuple(out.shape)}")
+
+    vec = out[0].detach().float().cpu().numpy().astype(np.float32)
+    meta = {
+        "pooled_source": "tile",
+        "feature_dim": int(vec.shape[0]),
+        "patch_size_m": int(patch_size_m),
+    }
+    return vec, meta
+
+
 @register("anysat")
 class AnySatEmbedder(EmbedderBase):
     DEFAULT_FETCH_WORKERS = 8
@@ -458,19 +549,31 @@ def describe(self) -> dict[str, Any]:
                 "cloudy_pct": self.input_spec.cloudy_pct,
                 "composite": self.input_spec.composite,
                 "normalization": "per_tile_zscore",
+                "grid_feature_mode": "dense",
+                "pooled_source": "patch",
             },
             "model_config": {
                 "variant": {
                     "type": "string",
                     "default": "base",
                     "choices": ["base"],
-                }
+                },
+                "grid_feature_mode": {
+                    "type": "string",
+                    "default": "dense",
+                    "choices": ["dense", "patch"],
+                },
+                "pooled_source": {
+                    "type": "string",
+                    "default": "patch",
+                    "choices": ["patch", "tile"],
+                },
             },
             "notes": [
                 "AnySat expects S2 time-series + day-of-year dates.",
                 "This adapter builds T frames by splitting TemporalSpec.range into equal sub-windows.",
                 "Loads AnySat from a vendored local runtime and optional Hugging Face checkpoint.",
-                "grid output maps AnySat output='patch' to [D,H,W].",
+                "grid output defaults to AnySat output='dense'; pooled defaults to patch-grid pooling but can opt into native tile output.",
             ],
         }
 
@@ -514,6 +617,8 @@ def get_embedding(
         flash_attn = bool(runtime_cfg["flash_attn"])
         image_size = int(runtime_cfg["image_size"])
         norm_mode = str(runtime_cfg["norm_mode"])
+        grid_feature_mode = str(runtime_cfg["grid_feature_mode"])
+        pooled_source = str(runtime_cfg["pooled_source"])
         patch_size_m = int(ss.scale_m)
 
         if input_chw is None:
@@ -557,11 +662,22 @@ def get_embedding(
             norm_mode=norm_mode,
             device=dev,
         )
-        grid, fmeta = _anysat_patch_features(
-            model,
-            s2_input,
-            patch_size_m=patch_size_m,
-        )
+        feature_mode = grid_feature_mode if output.mode == "grid" else "patch"
+        pooled_vec = None
+        if output.mode == "pooled" and pooled_source == "tile":
+            pooled_vec, fmeta = _anysat_tile_features(
+                model,
+                s2_input,
+                patch_size_m=patch_size_m,
+            )
+            grid = None
+        else:
+            grid, fmeta = _anysat_spatial_features(
+                model,
+                s2_input,
+                patch_size_m=patch_size_m,
+                feature_mode=feature_mode,
+            )
 
         meta = build_meta(
             model=self.model_name,
@@ -582,6 +698,8 @@ def get_embedding(
                 "model_size": model_size,
                 "flash_attn": bool(flash_attn),
                 "normalization": norm_mode,
+                "grid_feature_mode": grid_feature_mode,
+                "pooled_source": pooled_source,
                 "start": t.start,
                 "end": t.end,
                 "n_frames": int(raw_tchw.shape[0]),
@@ -595,21 +713,29 @@ def get_embedding(
         )
 
         if output.mode == "pooled":
+            if pooled_vec is not None:
+                ometa = {
+                    **meta,
+                    "pooling": "tile_native",
+                    "requested_pooling": output.pooling,
+                    "pooled_shape": tuple(pooled_vec.shape),
+                }
+                return Embedding(data=pooled_vec, meta=ometa)
             if output.pooling == "max":
                 vec = np.max(grid, axis=(1, 2)).astype(np.float32)
             else:
                 vec = np.mean(grid, axis=(1, 2)).astype(np.float32)
             ometa = {
                 **meta,
-                "pooling": f"patch_{output.pooling}",
+                "pooling": f"{feature_mode}_{output.pooling}",
                 "pooled_shape": tuple(vec.shape),
             }
             return Embedding(data=vec, meta=ometa)
 
         if output.mode == "grid":
             gmeta = {
                 **meta,
-                "grid_kind": "patch_tokens",
+                "grid_kind": "dense_subpatch" if feature_mode == "dense" else "patch_tokens",
                 "grid_hw": (int(grid.shape[1]), int(grid.shape[2])),
                 "grid_shape": tuple(grid.shape),
             }