slide2vec 5.0.1__tar.gz → 5.1.1__tar.gz

{slide2vec-5.0.1 → slide2vec-5.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: slide2vec
-Version: 5.0.1
+Version: 5.1.1
 Summary: Embedding of whole slide images with Foundation Models
 Author-email: Clément Grisi <clement.grisi@radboudumc.nl>
 License-Expression: Apache-2.0

{slide2vec-5.0.1 → slide2vec-5.1.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "slide2vec"
-version = "5.0.1"
+version = "5.1.1"
 description = "Embedding of whole slide images with Foundation Models"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -167,7 +167,7 @@ no_implicit_reexport = true
 max-line-length = 160
 
 [tool.bumpver]
-current_version = "5.0.1"
+current_version = "5.1.1"
 version_pattern = "MAJOR.MINOR.PATCH"
 commit = false       # We do version bumping in CI, not as a commit
 tag = false          # Git tag already exists — we don't auto-tag

{slide2vec-5.0.1 → slide2vec-5.1.1}/slide2vec/__init__.py

@@ -11,7 +11,7 @@ from slide2vec.api import (
 from slide2vec.artifacts import HierarchicalEmbeddingArtifact, SlideEmbeddingArtifact, TileEmbeddingArtifact
 
 
-__version__ = "5.0.1"
+__version__ = "5.1.1"
 
 __all__ = [
     "Model",

{slide2vec-5.0.1 → slide2vec-5.1.1}/slide2vec/runtime/dense_regions.py

@@ -5,7 +5,10 @@ The dense counterpart of the pooled coordinate path (``compute_tile_embeddings_f
 each sampled ROI is read **spacing-aware** from the slide, run through the encoder's
 normalization-only dense transform (``get_dense_transform`` — NOT the pooled transform,
 which crops), padded up to the encoder's patch multiple, and encoded via
-``encode_tiles_dense`` into a ``(d, grid_h, grid_w)`` token grid.
+``encode_tiles_dense`` into a ``(d, grid_h, grid_w)`` token grid. ``iter_regions_dense``
+**streams** these grids — yielding one per coordinate, in coordinate order, holding at most
+one ``batch_size`` chunk resident — so host memory is bounded by ``batch_size`` rather than
+by a slide's ROI count.
 
 This is the extraction half of soma's slide-manifest segmentation path: slide2vec reads
 regions + encodes (it already owns the region reader and the dense encode); soma sources
@@ -18,23 +21,49 @@ the finest pyramid level ``<=`` the requested µm/px is read and downscaled to t
 the same spacing. The ``wsi`` is injected (any object exposing ``read_region_at_spacing``),
 so the loop is unit-testable offline with a fake reader + a random-weight encoder.
 
-Whole-tile only (one padded forward per region). Sliding-window dense extraction over
-coordinates (``window_size`` < input) is a deferred follow-up — large ROIs that exceed the
-encoder's comfortable field are out of scope for the first increment.
+Both dense modes run through one primitive (:func:`~slide2vec.runtime.dense_sliding.encode_dense_sliding`):
+``window_size=None`` is a single whole-tile forward (byte-identical to the legacy
+whole-region encode), and a ``window_size`` smaller than the encoded tile slides the
+encoder's native field over the padded tile and blends the per-window token grids with a
+separable raised-cosine map — letting a native-field encoder (e.g. 224-px Virchow2/phikon)
+serve a larger ROI without interpolating its position embeddings.
 """
 
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import Callable, Sequence
+from typing import Callable, Iterator, Sequence
 
 import numpy as np
 import torch
 import torch.nn.functional as F
 from PIL import Image
 
+from slide2vec.runtime.batching import autocast_dtype
+from slide2vec.runtime.dense_sliding import encode_dense_sliding
 from slide2vec.runtime.slide_encode import slide_encode_autocast_ctx
 
+
+def _resolve_output_dtype(output_dtype: "torch.dtype | None", precision: str) -> "torch.dtype":
+    """Resolve the dtype emitted grids are materialized in.
+
+    Defaults (``output_dtype is None``) to the model's compute precision — fp16 runs
+    yield fp16 grids — so the engine no longer force-upcasts everything to float32.
+    numpy has no bfloat16, so a bf16 compute precision widens to float32 (its lossless
+    container) and an *explicit* ``torch.bfloat16`` request is rejected: the grid is
+    materialized via ``.numpy()`` and bfloat16 cannot cross that boundary.
+    """
+    if output_dtype is None:
+        compute = autocast_dtype(torch, precision)  # float16 | bfloat16 | None(fp32)
+        return torch.float16 if compute == torch.float16 else torch.float32
+    if output_dtype == torch.bfloat16:
+        raise ValueError(
+            "output_dtype=torch.bfloat16 cannot be materialized as a numpy grid; "
+            "request torch.float16 or torch.float32"
+        )
+    return output_dtype
+
+
 _PAD_MODES = {"reflect", "constant", "zero", "replicate"}
 
 
@@ -136,7 +165,7 @@ def _resolve_encode_fn(
     )
 
 
-def encode_regions_dense(
+def iter_regions_dense(
     *,
     model,
     device: torch.device | str,
@@ -147,14 +176,23 @@ def encode_regions_dense(
     tolerance: float = 0.05,
     pad_mode: str = "reflect",
     image_pad_value: float | None = None,
+    window_size: int | None = None,
+    overlap: float = 0.0,
     feature_kind: str = "patch_features",
     attention_blocks: tuple[int, ...] = (-1,),
     attention_include_registers: bool = False,
     batch_size: int = 1,
     precision: str = "fp32",
+    output_dtype: "torch.dtype | None" = None,
     dense_transform: Callable | None = None,
-) -> np.ndarray:
-    """Encode slide regions at ``coordinates`` into dense grids; return ``(N, d, gh, gw)``.
+) -> Iterator[np.ndarray]:
+    """Stream slide regions at ``coordinates`` into dense grids, one per coordinate.
+
+    Yields one ``(d, grid_h, grid_w)`` grid per coordinate, in coordinate order, in the
+    model's compute ``precision`` by default (fp16 runs yield fp16 grids; see
+    ``output_dtype``). Regions are read and encoded one ``batch_size`` chunk at a time, so
+    resident host memory is bounded by ``batch_size`` rather than by a slide's ROI count
+    (the loop holds at most one batch of grids resident — no per-slide accumulation).
 
     Injectable core: takes a constructed dense-capable ``model`` (with
     ``encode_tiles_dense`` / ``encode_tiles_attention`` / ``patch_size`` /
@@ -162,19 +200,38 @@ def encode_regions_dense(
     ``read_region_at_spacing(location, requested_spacing_um, size, *, tolerance,
     interpolation)``, so it runs offline in tests with random weights + a fake reader.
 
+    Arguments are validated and geometry is resolved **eagerly** (before any region is
+    read): an invalid ``pad_mode`` or ``feature_kind`` raises at the call site, not on the
+    first ``next()``. Iteration itself is lazy — reads advance one batch at a time.
+
     Args:
         coordinates: ``(x, y)`` top-left locations in **level-0** pixel space (the hs2p
             tiling convention; passed straight to ``read_region_at_spacing``).
         requested_spacing_um: µm/px to read each region at.
         target_size: supervision tile size (int or ``(h, w)``); the region is read at this
             size at ``requested_spacing_um`` and the token grid registers to it.
-
-    Returns a ``float32`` array of dense grids in coordinate order. ``feature_kind``
-    selects ``encode_tiles_dense`` (patch grid) vs ``encode_tiles_attention`` (CLS-attention
-    grid); both produce a ``(C, gh, gw)`` grid and share this path.
+        window_size: encoder field-of-view chunk fed through the backbone per forward.
+            ``None`` (default) is one whole-tile forward, byte-identical to the
+            whole-region encode; a value smaller than the encoded tile slides the encoder
+            over patch-aligned windows and blends the token grids (raised-cosine map). The
+            output grid is always the whole geometry's ``(grid_h, grid_w)`` either way —
+            sliding is internal to extraction.
+        overlap: fractional window overlap in ``[0, 1)`` for the sliding path (ignored when
+            ``window_size is None``); the stride is ``window * (1 - overlap)``.
+        output_dtype: torch dtype the grids are materialized in. ``None`` (default) follows
+            the compute ``precision`` — fp16 → fp16, fp32 → fp32, bf16 → fp32 (numpy has no
+            bfloat16). Pass e.g. ``torch.float32`` to force a lossless cache regardless of
+            precision; an explicit ``torch.bfloat16`` is rejected (cannot cross ``.numpy()``).
+
+    Yields grids in coordinate order in ``output_dtype``; empty ``coordinates`` yields nothing.
+    ``feature_kind`` selects ``encode_tiles_dense`` (patch grid) vs
+    ``encode_tiles_attention`` (CLS-attention grid); both produce a ``(C, gh, gw)`` grid and
+    share this path. Each yielded grid is a standalone contiguous copy, so it does not pin
+    the rest of its batch's memory alive.
     """
     if pad_mode not in _PAD_MODES:
         raise ValueError(f"unsupported pad_mode {pad_mode!r}; expected one of {sorted(_PAD_MODES)}")
+    resolved_output_dtype = _resolve_output_dtype(output_dtype, precision)
     geometry = compute_dense_geometry(target_size=target_size, patch_size=model.patch_size)
     if dense_transform is None:
         dense_transform = model.get_dense_transform()
@@ -185,11 +242,8 @@ def encode_regions_dense(
         attention_include_registers=attention_include_registers,
     )
     target_h, target_w = geometry.target_size
-
     coords = [(int(x), int(y)) for x, y in coordinates]
-    grid_h, grid_w = geometry.grid_shape
-    if not coords:
-        return np.empty((0, 0, grid_h, grid_w), dtype=np.float32)
+    step = max(1, int(batch_size))
 
     def _read_padded(location: tuple[int, int]) -> torch.Tensor:
         region = wsi.read_region_at_spacing(
@@ -215,15 +269,33 @@ def encode_regions_dense(
             tensor, geometry, pad_mode=pad_mode, image_pad_value=image_pad_value
         )
 
-    grids: list[np.ndarray] = []
-    with torch.inference_mode(), slide_encode_autocast_ctx(device, precision):
-        for start in range(0, len(coords), max(1, int(batch_size))):
-            chunk = coords[start : start + max(1, int(batch_size))]
-            batch = torch.stack([_read_padded(loc) for loc in chunk]).to(device, non_blocking=True)
-            out = encode_fn(batch)
-            if out.ndim != 4:
-                raise ValueError(
-                    f"{feature_kind} encode returned a {out.ndim}-D tensor; expected (B, d, gh, gw)."
+    def _stream() -> Iterator[np.ndarray]:
+        with torch.inference_mode(), slide_encode_autocast_ctx(device, precision):
+            for start in range(0, len(coords), step):
+                chunk = coords[start : start + step]
+                batch = torch.stack([_read_padded(loc) for loc in chunk]).to(
+                    device, non_blocking=True
+                )
+                # Every batch goes through the one windowed primitive: window_size=None
+                # short-circuits to a single whole-tile forward (byte-identical to the
+                # whole-region encode), so there is no separate whole-region branch.
+                out = encode_dense_sliding(
+                    model,
+                    batch,
+                    geometry=geometry,
+                    window_size=window_size,
+                    overlap=overlap,
+                    encode_fn=encode_fn,
                 )
-            grids.append(out.detach().float().cpu().numpy())
-    return np.concatenate(grids, axis=0)
+                if out.ndim != 4:
+                    raise ValueError(
+                        f"{feature_kind} encode returned a {out.ndim}-D tensor; expected (B, d, gh, gw)."
+                    )
+                batch_np = out.detach().to(resolved_output_dtype).cpu().numpy()
+                for i in range(batch_np.shape[0]):
+                    # Standalone C-contiguous copy: a per-row view would pin the whole
+                    # batch alive (the blended sliding output is contiguous, so a view of
+                    # it would not copy). ``.copy()`` always copies, in C order.
+                    yield batch_np[i].copy()
+
+    return _stream()

slide2vec-5.1.1/slide2vec/runtime/dense_sliding.py

@@ -0,0 +1,185 @@
+"""Sliding-window dense encoding — ``window_size`` + ``overlap`` as a free knob.
+
+The ``whole`` path feeds the full padded tile through the encoder in one forward,
+interpolating the positional embeddings to the larger grid. That is one end of a
+single mechanism; the other end is running the encoder over smaller **windows** and
+stitching the per-window token grids. Three sizes that are usually conflated —
+
+* **native size** (e.g. 224) — sets the pos-embed table; not a hard input limit
+  (``dynamic_img_size`` lets a ViT process a larger field at the correct mpp);
+* **window size** ``W`` — how big a chunk goes through the ViT in one forward;
+* **input size** — the padded ``encoded_size`` we want dense features for.
+
+``whole`` is ``W >= input`` (one window, zero stitching); native sliding is ``W = 224``;
+the useful middle is ``W = 512`` slid over a larger input. So this is **one**
+parametrized path, not a separate mode: :func:`encode_dense_sliding` takes
+``window_size`` (``None`` ⇒ ``whole``) and ``overlap``, and the ``whole`` case falls out
+as the degenerate single window — which we short-circuit to the exact same
+``encode_tiles_dense(batch)`` call, so it stays **byte-identical** to the whole-region
+path (the parity anchor).
+
+Stitching happens in **token space** (the grid the decoder/head consume), so the output
+is always ``(B, d, grid_h, grid_w)`` for ``geometry.grid_shape`` regardless of
+``window_size`` — sliding is purely internal to extraction. Windows and strides are kept
+patch-aligned, so each window maps cleanly onto a block of tokens; overlapping windows
+are blended with a separable raised-cosine importance map (the standard frozen-backbone
+dense-inference recipe, cf. MONAI ``sliding_window_inference``) to remove the
+block-boundary seams naive non-overlapping tiling would introduce.
+
+Ported from soma's ``soma/dense/sliding.py`` (the window/blend math is encoder
+featurization that belongs in slide2vec); adapted to slide2vec's own
+:class:`~slide2vec.runtime.dense_regions.DenseGridGeometry`.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING, Callable
+
+import torch
+
+if TYPE_CHECKING:
+    from slide2vec.runtime.dense_regions import DenseGridGeometry
+
+__all__ = [
+    "cover_origins",
+    "encode_dense_sliding",
+    "resolve_window_geometry",
+]
+
+
+def _round_up(value: int, multiple: int) -> int:
+    return ((value + multiple - 1) // multiple) * multiple
+
+
+def _round_to(value: float, multiple: int) -> int:
+    return max(multiple, int(round(value / multiple)) * multiple)
+
+
+def cover_origins(extent: int, size: int, stride: int) -> list[int]:
+    """Start offsets of ``size``-wide windows that fully cover ``[0, extent)``.
+
+    Walks ``[0, extent - size]`` in ``stride`` steps and, if the last step leaves a gap,
+    appends one final start flush to the far edge (``extent - size``) so coverage is
+    complete with no partial tail. ``extent``/``size``/``stride`` are patch multiples,
+    so every start is too — the edge-flush ``extent - size`` is a difference of patch
+    multiples.
+    """
+    if size >= extent:
+        return [0]
+    starts = list(range(0, extent - size + 1, stride))
+    if starts[-1] + size < extent:
+        starts.append(extent - size)  # shift the last window flush to the edge
+    return starts
+
+
+def _window_starts(extent: int, win: int, stride: int) -> list[int]:
+    """Patch-aligned encoder-window starts — the token-space use of :func:`cover_origins`."""
+    return cover_origins(extent, win, stride)
+
+
+def resolve_window_geometry(
+    geometry: DenseGridGeometry, *, window_size: int | None, overlap: float
+) -> tuple[tuple[int, int], tuple[int, int], list[int], list[int]]:
+    """Resolve per-dim window size, stride, and start offsets (all patch-aligned).
+
+    ``window_size`` is rounded up to the patch multiple and clamped to the encoded
+    extent; because ``round_up`` is monotonic, ``window_size >= target_size`` always
+    clamps to the full extent ⇒ a single window ⇒ the ``whole`` path. ``stride`` is
+    ``window * (1 - overlap)`` rounded to the patch multiple and clamped to
+    ``[patch, window]``.
+    """
+    enc_h, enc_w = geometry.encoded_size
+    ph, pw = geometry.patch_size
+    if window_size is None:
+        return (enc_h, enc_w), (enc_h, enc_w), [0], [0]
+
+    win_h = min(_round_up(int(window_size), ph), enc_h)
+    win_w = min(_round_up(int(window_size), pw), enc_w)
+    keep = 1.0 - float(overlap)
+    stride_h = min(win_h, _round_to(win_h * keep, ph))
+    stride_w = min(win_w, _round_to(win_w * keep, pw))
+    starts_h = _window_starts(enc_h, win_h, stride_h)
+    starts_w = _window_starts(enc_w, win_w, stride_w)
+    return (win_h, win_w), (stride_h, stride_w), starts_h, starts_w
+
+
+def _hann_1d(n: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+    """Strictly-positive raised-cosine weights of length ``n`` (uniform if ``n == 1``).
+
+    ``0.5 - 0.5*cos(2*pi*(i+1)/(n+1))`` is > 0 for every ``i in [0, n)`` (no zeros at
+    the edges), so the accumulated weight map never hits zero where a window covers.
+    """
+    if n <= 1:
+        return torch.ones(n, device=device, dtype=dtype)
+    i = torch.arange(1, n + 1, device=device, dtype=dtype)
+    return 0.5 - 0.5 * torch.cos(2.0 * math.pi * i / (n + 1))
+
+
+def encode_dense_sliding(
+    encoder,
+    batch: torch.Tensor,
+    *,
+    geometry: DenseGridGeometry,
+    window_size: int | None,
+    overlap: float = 0.0,
+    encode_fn: Callable[[torch.Tensor], torch.Tensor] | None = None,
+) -> torch.Tensor:
+    """Encode a padded ``(B, C, enc_h, enc_w)`` batch into ``(B, d, grid_h, grid_w)``.
+
+    ``window_size is None`` (or any window that covers the whole encoded input) is the
+    degenerate single-window case: it short-circuits to one full-tile forward,
+    byte-identical to the whole-region path. Otherwise the encoder runs over
+    patch-aligned overlapping windows and the per-window token grids are blended with a
+    separable raised-cosine importance map. The stitch math runs in fp32 (sub-grids are
+    upcast before accumulation) so blended regions don't accumulate autocast-dtype error.
+
+    ``encode_fn`` is the per-window encode callable ``(B, C, wh, ww) -> (B, d, th, tw)``;
+    it defaults to ``encoder.encode_tiles_dense`` (the patch-feature grid). The attention
+    path passes ``encoder.encode_tiles_attention`` (partial-applied with its
+    block/register knobs) so a CLS-attention grid stitches through the identical
+    raised-cosine blending — the output is just ``(B, K, grid)`` instead of ``(B, d, grid)``.
+    """
+    if encode_fn is None:
+        encode_fn = encoder.encode_tiles_dense
+    (win_h, win_w), _, starts_h, starts_w = resolve_window_geometry(
+        geometry, window_size=window_size, overlap=overlap
+    )
+    if len(starts_h) == 1 and len(starts_w) == 1:
+        # Single window == the whole encoded tile: identical forward to the whole-region path.
+        return encode_fn(batch)
+
+    ph, pw = geometry.patch_size
+    grid_h, grid_w = geometry.grid_shape
+    wth, wtw = win_h // ph, win_w // pw
+    # Raised-cosine weights where windows overlap; uniform along any dim that is not
+    # actually tiled (a single window there) — avoids needless edge attenuation.
+    fdtype = torch.float32
+    wh = (
+        _hann_1d(wth, batch.device, fdtype)
+        if len(starts_h) > 1
+        else torch.ones(wth, device=batch.device, dtype=fdtype)
+    )
+    ww = (
+        _hann_1d(wtw, batch.device, fdtype)
+        if len(starts_w) > 1
+        else torch.ones(wtw, device=batch.device, dtype=fdtype)
+    )
+    weight = torch.outer(wh, ww)  # (wth, wtw)
+
+    acc: torch.Tensor | None = None
+    wsum = torch.zeros(1, 1, grid_h, grid_w, device=batch.device, dtype=fdtype)
+    for sh in starts_h:
+        th = sh // ph
+        for sw in starts_w:
+            tw = sw // pw
+            window = batch[:, :, sh : sh + win_h, sw : sw + win_w]
+            sub = encode_fn(window).to(fdtype)  # (B, d, wth, wtw)
+            if acc is None:
+                acc = torch.zeros(
+                    sub.shape[0], sub.shape[1], grid_h, grid_w, device=batch.device, dtype=fdtype
+                )
+            acc[:, :, th : th + wth, tw : tw + wtw] += sub * weight
+            wsum[:, :, th : th + wth, tw : tw + wtw] += weight
+    assert acc is not None  # at least one window always runs
+    return acc / wsum

{slide2vec-5.0.1 → slide2vec-5.1.1}/slide2vec.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: slide2vec
-Version: 5.0.1
+Version: 5.1.1
 Summary: Embedding of whole slide images with Foundation Models
 Author-email: Clément Grisi <clement.grisi@radboudumc.nl>
 License-Expression: Apache-2.0

{slide2vec-5.0.1 → slide2vec-5.1.1}/slide2vec.egg-info/SOURCES.txt

@@ -54,6 +54,7 @@ slide2vec/runtime/artifacts_collect.py
 slide2vec/runtime/batching.py
 slide2vec/runtime/cpu_budget.py
 slide2vec/runtime/dense_regions.py
+slide2vec/runtime/dense_sliding.py
 slide2vec/runtime/distributed.py
 slide2vec/runtime/distributed_stage.py
 slide2vec/runtime/embedding.py
@@ -84,6 +85,7 @@ tests/test_architecture_runtime_split.py
 tests/test_attention_extraction.py
 tests/test_dense_extraction.py
 tests/test_dense_regions.py
+tests/test_dense_sliding.py
 tests/test_encoder_registry.py
 tests/test_hs2p_package_cutover.py
 tests/test_output_consistency.py