From 299f7f0276f9ab4ec6d43a27bc11a028f24c709d Mon Sep 17 00:00:00 2001 From: clemsgrs Date: Tue, 30 Jun 2026 15:32:30 +0200 Subject: [PATCH 1/3] Add GPFM and GenBio-PathFM as supported tile encoders GPFM (majiabo/GPFM, MIT) is a ViT-L/14 DINOv2 tile encoder. Its weights ship as a standalone GPFM.pth checkpoint rather than a timm hf-hub config, so build the timm arch unpretrained and load_state_dict the checkpoint (MOOZY pattern). Verified on the real weights: GPFM.pth is a bare DINOv2 state_dict that loads strict=True with zero missing/unexpected keys and a 1024-d pooled forward. Metadata: dim 1024, input_size 224, patch_size 14, supported_spacing_um 0.25 (40x), precision fp32. Dense-capable. GenBio-PathFM (genbio-ai/genbio-pathfm) is a 1.1B-param custom HF AutoModel (trust_remote_code), so register it as a custom TileEncoder with non-ImageNet normalization. The backbone is a single-channel ViT (in_chans=1, embed_dim 1536): its canonical forward splits RGB into 3 single-channel images and concatenates the per-channel CLS tokens -> 4608-d (verified on real weights; matches the card's advertised feature dim). No dense path, so no patch_size. Metadata: dim 4608, input_size 224, supported_spacing_um 0.5, precision fp32. Add focused metadata-contract tests plus heavy real-weight forward-shape tests, and recognise non-dense tile encoders in the patch_size contract. Also drop the stale gated-weights/HF_TOKEN note from the mSTAR docs row. --- README.md | 2 +- docs/models.rst | 20 ++++-- slide2vec/encoders/models/__init__.py | 4 ++ slide2vec/encoders/models/genbio.py | 90 +++++++++++++++++++++++++++ slide2vec/encoders/models/gpfm.py | 89 ++++++++++++++++++++++++++ tests/test_encoder_registry.py | 31 +++++++++ tests/test_gpfm_genbio_heavy.py | 53 ++++++++++++++++ tests/test_patch_size_metadata.py | 22 ++++++- tests/test_regression_core.py | 2 + 9 files changed, 305 insertions(+), 8 deletions(-) create mode 100644 slide2vec/encoders/models/genbio.py create mode 100644 slide2vec/encoders/models/gpfm.py create mode 100644 tests/test_gpfm_genbio_heavy.py diff --git a/README.md b/README.md index 83d44a0..7d6a46a 100644 --- a/README.md +++ b/README.md @@ -119,7 +119,7 @@ The package writes explicit artifact directories: ### Supported Models -`slide2vec` currently ships preset configs for 18 tile-level models and 3 slide-level models. +`slide2vec` currently ships preset configs for 20 tile-level models and 3 slide-level models. For the full catalog and preset names, see [`docs/models.md`](docs/models.md). ## CLI diff --git a/docs/models.rst b/docs/models.rst index 35ed977..65c5a33 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -72,7 +72,12 @@ Tile-level encoders - `mSTAR `_ - 1024 - ``0.5`` - - Xu et al. (2024); gated weights, requires ``HF_TOKEN`` + - Xu et al. (2024) + * - ``gpfm`` + - `GPFM `_ + - 1024 + - ``0.25`` + - Ma et al. (2024); MIT license * - ``uni`` - `UNI `_ - 1024 @@ -118,6 +123,11 @@ Tile-level encoders - 3072 - ``0.25``, ``0.5``, ``1.0``, ``2.0`` - Karasikov et al. (2025) + * - ``genbio-pathfm`` + - `GenBio-PathFM `_ + - 4608 + - ``0.5`` + - GenBio AI (2024); custom Community License; non-ImageNet normalization Dense tile grids ~~~~~~~~~~~~~~~~ @@ -127,7 +137,7 @@ Dense tile feature extraction is available on tile encoders that implement instead of the pooled ``(B, D)`` tensor returned by ``encode_tiles``. The following built-in tile presets are covered by the dense encoder interface: -``conch``, ``conchv15``, ``gigapath``, ``h0-mini``, ``h-optimus-0``, +``conch``, ``conchv15``, ``gigapath``, ``gpfm``, ``h0-mini``, ``h-optimus-0``, ``h-optimus-1``, ``hibou-b``, ``hibou-l``, ``lunit``, ``midnight``, ``mstar``, ``musk``, ``phikon``, ``phikonv2``, ``prost40m``, ``uni``, ``uni2``, ``virchow``, and ``virchow2``. @@ -158,9 +168,9 @@ prefix-token self-attention as a spatial grid ``(B, K, h, w)`` — see the contract and knobs. The following built-in tile presets are covered: ``conch``, ``conchv15``, -``gigapath``, ``h0-mini``, ``h-optimus-0``, ``h-optimus-1``, ``hibou-b``, -``hibou-l``, ``lunit``, ``midnight``, ``mstar``, ``phikon``, ``phikonv2``, -``prost40m``, ``uni``, ``uni2``, ``virchow``, and ``virchow2``. +``gigapath``, ``gpfm``, ``h0-mini``, ``h-optimus-0``, ``h-optimus-1``, +``hibou-b``, ``hibou-l``, ``lunit``, ``midnight``, ``mstar``, ``phikon``, +``phikonv2``, ``prost40m``, ``uni``, ``uni2``, ``virchow``, and ``virchow2``. Notes: diff --git a/slide2vec/encoders/models/__init__.py b/slide2vec/encoders/models/__init__.py index 964eca1..a012419 100644 --- a/slide2vec/encoders/models/__init__.py +++ b/slide2vec/encoders/models/__init__.py @@ -5,7 +5,9 @@ from . import ( conch, + genbio, gigapath, + gpfm, hibou, hoptimus, lunit, @@ -23,7 +25,9 @@ __all__ = [ "conch", + "genbio", "gigapath", + "gpfm", "hibou", "hoptimus", "lunit", diff --git a/slide2vec/encoders/models/genbio.py b/slide2vec/encoders/models/genbio.py new file mode 100644 index 0000000..2100671 --- /dev/null +++ b/slide2vec/encoders/models/genbio.py @@ -0,0 +1,90 @@ +"""GenBio-PathFM tile encoder. + +GenBio-PathFM (GenBio AI, 2024; ``genbio-ai/genbio-pathfm``) is a 1.1B-param ViT +histopathology tile encoder (JEDI = JEPA + DINO training on public data). It is +loaded via HF ``AutoModel(trust_remote_code=True)`` (auto_map -> +``GenBioPathFMModel``), so it is a custom :class:`TileEncoder`, not a +``TimmTileEncoder``. Weights are openly downloadable (safetensors); released +under the custom **GenBio AI Community License** (not OSI-approved — read it for +acceptable-use / commercial restrictions; slide2vec only wraps user-downloaded +weights and does not redistribute them). + +Output dim (4608, verified on real weights): the backbone is a *single-channel* +ViT (``in_chans=1``, ``embed_dim=1536``). The model's canonical ``forward`` takes +an RGB ``[B, 3, H, W]`` tensor, treats each colour channel as a separate +single-channel image, encodes all three, and concatenates the three per-channel +CLS tokens -> ``[B, embed_dim * 3] = [B, 4608]``. This per-channel-CLS +concatenation is the model's intrinsic design (matching the HF card's advertised +feature dimension of 4608), not an ad-hoc CLS+patch pooling; ``encode_tiles`` +therefore returns the model's default ``forward`` output directly. (Confirmed by +running a ``(1, 3, 224, 224)`` dummy through the real weights -> shape +``(1, 4608)``; see ``tests/test_gpfm_genbio_heavy.py``.) + +Normalization is **non-ImageNet** (``config.json`` ``image_mean`` / ``image_std``) +and must be set explicitly. No dense (spatial-grid) path is exposed in this +encoder, so it declares no ``patch_size``. +""" + +from typing import Callable + +import torch +from torch import Tensor +from torchvision.transforms import v2 + +from slide2vec.encoders.base import ( + TileEncoder, + preferred_default_device, + resolve_requested_output_variant, +) +from slide2vec.encoders.registry import register_encoder + +_HF_REPO_ID = "genbio-ai/genbio-pathfm" +# Non-ImageNet normalization from the model's config.json (image_mean / image_std). +_GENBIO_MEAN = (0.697, 0.575, 0.728) +_GENBIO_STD = (0.188, 0.240, 0.187) +# embed_dim (1536) x 3 colour channels — see module docstring. +_GENBIO_ENCODE_DIM = 4608 + + +@register_encoder( + "genbio-pathfm", + output_variants={"default": {"encode_dim": _GENBIO_ENCODE_DIM}}, + default_output_variant="default", + input_size=224, + supported_spacing_um=0.5, # 20x; card states no magnification, house default + precision="fp32", # upstream runs plain fp32, no autocast + source="genbio-ai/genbio-pathfm", +) +class GenBioPathFM(TileEncoder): + def __init__(self, *, output_variant: str | None = None): + from transformers import AutoModel + + self._model = AutoModel.from_pretrained(_HF_REPO_ID, trust_remote_code=True).eval() + self._device = preferred_default_device() + self._output_variant = resolve_requested_output_variant(output_variant) + + def get_transform(self) -> Callable: + return v2.Compose([ + v2.ToImage(), + v2.Resize(224, interpolation=v2.InterpolationMode.BICUBIC, antialias=True), + v2.CenterCrop(224), + v2.ToDtype(torch.float32, scale=True), + v2.Normalize(mean=_GENBIO_MEAN, std=_GENBIO_STD), + ]) + + def encode_tiles(self, batch: Tensor) -> Tensor: + # Canonical forward: per-channel CLS tokens concatenated -> (B, 4608). + return self._model(batch) + + @property + def encode_dim(self) -> int: + return _GENBIO_ENCODE_DIM + + @property + def device(self) -> torch.device: + return self._device + + def to(self, device: torch.device | str) -> "GenBioPathFM": + self._device = torch.device(device) + self._model = self._model.to(self._device) + return self diff --git a/slide2vec/encoders/models/gpfm.py b/slide2vec/encoders/models/gpfm.py new file mode 100644 index 0000000..74a0dd4 --- /dev/null +++ b/slide2vec/encoders/models/gpfm.py @@ -0,0 +1,89 @@ +"""GPFM tile encoder (Generalizable Pathology Foundation Model). + +GPFM (Ma et al., 2024; ``birkhoffkiki/GPFM``) is a ``ViT-L/14`` DINOv2 tile +encoder (embedding dim 1024, patch_size 14). Weights are openly available under +the MIT license at ``majiabo/GPFM`` but ship as a standalone ``GPFM.pth`` +checkpoint, so — unlike ``hf-hub:`` timm presets — we build the timm arch +unpretrained and ``load_state_dict`` the downloaded checkpoint (the MOOZY +loading pattern). + +The published ``GPFM.pth`` is a bare DINOv2 ``state_dict`` (``cls_token``, +``pos_embed``, ``blocks.*``, …) that loads into +``vit_large_patch14_dinov2.lvd142m`` with ``strict=True`` and zero missing / +unexpected keys (verified on the real weights; see +``tests/test_gpfm_genbio_heavy.py``). The loader below still defensively unwraps +common checkpoint wrappers (``{"model": ...}`` / ``{"teacher": ...}`` / +``{"student": ...}`` / ``{"state_dict": ...}``) and strips ``module.`` / +``backbone.`` prefixes so a re-exported checkpoint keeps loading strictly. +""" + +from typing import Mapping + +import torch +from huggingface_hub import hf_hub_download + +from slide2vec.encoders.base import TimmTileEncoder +from slide2vec.encoders.registry import register_encoder + +_HF_REPO_ID = "majiabo/GPFM" +_HF_CHECKPOINT = "GPFM.pth" +_CHECKPOINT_WRAPPER_KEYS = ("model", "teacher", "student", "state_dict", "teacher_backbone") +_STATE_DICT_PREFIXES = ("module.", "backbone.") + + +def _unwrap_gpfm_state_dict(payload: Mapping[str, object]) -> dict[str, torch.Tensor]: + """Reduce a GPFM checkpoint payload to a bare ``state_dict``. + + Unwraps a single common wrapper key if present, then strips ``module.`` / + ``backbone.`` prefixes off every key. A no-op on the published bare + ``GPFM.pth`` (no wrapper, no prefixes), so it preserves strict loading there + while tolerating a re-wrapped/prefixed re-export. + """ + state: object = payload + for wrapper in _CHECKPOINT_WRAPPER_KEYS: + if isinstance(state, Mapping) and wrapper in state and isinstance(state[wrapper], Mapping): + state = state[wrapper] + break + if not isinstance(state, Mapping): + raise ValueError( + f"Unexpected GPFM checkpoint payload: expected a state_dict mapping, got {type(state)}" + ) + cleaned: dict[str, torch.Tensor] = {} + for key, value in state.items(): + name = key + for prefix in _STATE_DICT_PREFIXES: + if name.startswith(prefix): + name = name[len(prefix):] + cleaned[name] = value + return cleaned + + +@register_encoder( + "gpfm", + output_variants={"default": {"encode_dim": 1024}}, + default_output_variant="default", + input_size=224, + patch_size=14, + supported_spacing_um=0.25, # 40x; ~512px tiles resized to 224 (per card) + precision="fp32", # upstream runs plain fp32, no autocast + source="majiabo/GPFM", +) +class GPFM(TimmTileEncoder): + def __init__(self, *, output_variant: str | None = None): + super().__init__( + "vit_large_patch14_dinov2.lvd142m", + output_variant=output_variant, + pretrained=False, + img_size=224, + init_values=1e-5, + dynamic_img_size=True, + ) + checkpoint_path = hf_hub_download(repo_id=_HF_REPO_ID, filename=_HF_CHECKPOINT) + payload = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + if not isinstance(payload, Mapping): + raise ValueError( + f"Invalid GPFM checkpoint payload: expected a dict, got {type(payload)}" + ) + state_dict = _unwrap_gpfm_state_dict(payload) + self._model.load_state_dict(state_dict, strict=True) + self._model.eval() diff --git a/tests/test_encoder_registry.py b/tests/test_encoder_registry.py index 72bbaa9..3ad80c5 100644 --- a/tests/test_encoder_registry.py +++ b/tests/test_encoder_registry.py @@ -12,7 +12,9 @@ "virchow2", "conch", "conchv15", + "genbio-pathfm", "gigapath", + "gpfm", "h-optimus-0", "h-optimus-1", "h0-mini", @@ -131,6 +133,35 @@ def test_mstar_slide_alias_resolves_to_mstar(): assert canonicalize_model_name("mstar-slide") == "mstar" +def test_gpfm_metadata_contract(): + info = encoder_registry.info("gpfm") + assert info["level"] == "tile" + assert info["input_size"] == 224 + assert info["patch_size"] == 14 + assert info["supported_spacing_um"] == pytest.approx(0.25) + assert info["precision"] == "fp32" + assert info["source"] == "majiabo/GPFM" + assert info["output_variants"]["default"]["encode_dim"] == 1024 + assert info["default_output_variant"] == "default" + + +def test_genbio_pathfm_metadata_contract(): + info = encoder_registry.info("genbio-pathfm") + assert info["level"] == "tile" + assert info["input_size"] == 224 + # GenBio is a custom AutoModel without a recoverable dense patch grid, so it + # declares no patch_size (not in the dense coverage lists). + assert info["patch_size"] is None + assert info["supported_spacing_um"] == pytest.approx(0.5) + assert info["precision"] == "fp32" + assert info["source"] == "genbio-ai/genbio-pathfm" + # GenBio is a single-channel ViT (in_chans=1): forward() splits RGB into 3 + # single-channel images and concatenates their per-channel CLS tokens, so the + # canonical pooled output is embed_dim*3 = 1536*3 = 4608 (verified on real weights). + assert info["output_variants"]["default"]["encode_dim"] == 4608 + assert info["default_output_variant"] == "default" + + def test_resolve_preprocessing_requirements_for_tile_encoder(): from slide2vec.encoders.registry import resolve_preprocessing_requirements diff --git a/tests/test_gpfm_genbio_heavy.py b/tests/test_gpfm_genbio_heavy.py new file mode 100644 index 0000000..4318e3c --- /dev/null +++ b/tests/test_gpfm_genbio_heavy.py @@ -0,0 +1,53 @@ +"""Real-weight forward-shape contracts for GPFM and GenBio-PathFM. + +These load multi-GB foundation-model weights on CPU (minutes each), so they are +marked ``heavy`` and excluded from the PR suite (``-m 'not heavy'``); they run on +the scheduled/manual heavy workflow. Each skips cleanly when weights / optional +deps / network are unavailable so a developer run never hard-fails. + +They pin the two registration choices that can only be confirmed against the real +checkpoints: + +* GPFM: ``GPFM.pth`` loads into the timm DINOv2 ViT-L/14 (``strict=True``) and the + pooled forward returns the registered 1024-d embedding. +* GenBio-PathFM: the single-channel backbone's canonical forward returns the + registered 4608-d (= embed_dim*3) per-channel-CLS concatenation. +""" + +from __future__ import annotations + +import pytest +import torch + +from slide2vec.encoders import encoder_registry + + +@pytest.mark.heavy +def test_gpfm_loads_strict_and_forward_is_1024_d(): + try: + encoder = encoder_registry.require("gpfm")(output_variant="default") + except Exception as exc: # network / weights / optional dep unavailable + pytest.skip(f"gpfm weights unavailable: {type(exc).__name__}: {exc}") + encoder = encoder.to("cpu") + assert encoder.encode_dim == 1024 + assert encoder.patch_size == (14, 14) + batch = torch.randn(1, 3, 224, 224) + with torch.inference_mode(): + out = encoder.encode_tiles(batch) + assert out.shape == (1, 1024) + + +@pytest.mark.heavy +def test_genbio_pathfm_forward_is_4608_d(): + try: + encoder = encoder_registry.require("genbio-pathfm")(output_variant="default") + except Exception as exc: # network / weights / optional dep unavailable + pytest.skip(f"genbio-pathfm weights unavailable: {type(exc).__name__}: {exc}") + encoder = encoder.to("cpu") + assert encoder.encode_dim == 4608 + batch = torch.randn(1, 3, 224, 224) + with torch.inference_mode(): + out = encoder.encode_tiles(batch) + # Single-channel ViT (in_chans=1): forward concatenates the 3 per-channel CLS + # tokens (embed_dim 1536) -> 4608. + assert out.shape == (1, 4608) diff --git a/tests/test_patch_size_metadata.py b/tests/test_patch_size_metadata.py index e60c3e7..570cc4f 100644 --- a/tests/test_patch_size_metadata.py +++ b/tests/test_patch_size_metadata.py @@ -28,6 +28,7 @@ "virchow": (14, 14), "virchow2": (14, 14), "gigapath": (14, 14), + "gpfm": (14, 14), "h-optimus-0": (14, 14), "h-optimus-1": (14, 14), "h0-mini": (14, 14), @@ -44,8 +45,14 @@ "hibou-l": (14, 14), } +# Non-dense tile encoders: real tile encoders with no recoverable patch grid +# (e.g. custom AutoModel backbones whose forward does not expose a spatial +# patch-token sequence), so they declare no ``patch_size`` and opt out of the +# dense interface. +NON_DENSE_TILE_ENCODERS = ["genbio-pathfm"] + # Non-dense encoders: no recoverable patch grid, so no declared patch_size. -NON_DENSE_ENCODERS = ["gigapath-slide", "titan", "prism"] +NON_DENSE_ENCODERS = ["gigapath-slide", "titan", "prism"] + NON_DENSE_TILE_ENCODERS def test_normalize_patch_size_int_and_tuple(): @@ -95,11 +102,22 @@ def test_resolve_patch_size_raises_for_non_dense_encoder(name): def test_every_tile_encoder_declares_patch_size(): - """All registered tile-level encoders are dense-capable and declare patch_size.""" + """Dense-capable tile encoders declare patch_size; non-dense ones opt out. + + Most tile encoders are dense-capable ViTs and must declare a ``patch_size``. + A few (``NON_DENSE_TILE_ENCODERS``, e.g. custom AutoModel backbones) expose no + recoverable patch grid and deliberately declare none. Every tile encoder must + fall into exactly one of those two tracked buckets. + """ for info in encoder_registry.list_with_metadata(): if info["level"] != "tile": continue name = info["name"] + if name in NON_DENSE_TILE_ENCODERS: + assert info.get("patch_size") is None, ( + f"non-dense tile encoder '{name}' must not declare a patch_size" + ) + continue assert info.get("patch_size") is not None, ( f"tile encoder '{name}' is dense-capable but declares no patch_size" ) diff --git a/tests/test_regression_core.py b/tests/test_regression_core.py index 62b5723..63b21ec 100644 --- a/tests/test_regression_core.py +++ b/tests/test_regression_core.py @@ -189,7 +189,9 @@ def test_list_models_can_filter_by_level(): assert list_models("tile") == [ "conch", "conchv15", + "genbio-pathfm", "gigapath", + "gpfm", "h-optimus-0", "h-optimus-1", "h0-mini", From b7b0f82c48a0795408a2f9174c9730cd01d0f7ae Mon Sep 17 00:00:00 2001 From: clemsgrs Date: Tue, 30 Jun 2026 15:59:51 +0200 Subject: [PATCH 2/3] Register GPFM at 0.5 um/px (20x-equivalent) effective spacing GPFM trains on 512px tiles at 0.25 um/px (128 um FOV) resized to 224, so the model effectively sees ~0.5 um/px (20x), matching how UNI's 256px@0.5 -> 224 is registered. 0.25 would tile half the trained FOV. --- docs/models.rst | 2 +- slide2vec/encoders/models/gpfm.py | 2 +- tests/test_encoder_registry.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/models.rst b/docs/models.rst index 65c5a33..efd27c8 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -76,7 +76,7 @@ Tile-level encoders * - ``gpfm`` - `GPFM `_ - 1024 - - ``0.25`` + - ``0.5`` - Ma et al. (2024); MIT license * - ``uni`` - `UNI `_ diff --git a/slide2vec/encoders/models/gpfm.py b/slide2vec/encoders/models/gpfm.py index 74a0dd4..567bdab 100644 --- a/slide2vec/encoders/models/gpfm.py +++ b/slide2vec/encoders/models/gpfm.py @@ -64,7 +64,7 @@ def _unwrap_gpfm_state_dict(payload: Mapping[str, object]) -> dict[str, torch.Te default_output_variant="default", input_size=224, patch_size=14, - supported_spacing_um=0.25, # 40x; ~512px tiles resized to 224 (per card) + supported_spacing_um=0.5, # 512px@0.25um (128um FOV) resized to 224 => ~0.5um/px effective (20x), as for UNI precision="fp32", # upstream runs plain fp32, no autocast source="majiabo/GPFM", ) diff --git a/tests/test_encoder_registry.py b/tests/test_encoder_registry.py index 3ad80c5..fc67f7e 100644 --- a/tests/test_encoder_registry.py +++ b/tests/test_encoder_registry.py @@ -138,7 +138,7 @@ def test_gpfm_metadata_contract(): assert info["level"] == "tile" assert info["input_size"] == 224 assert info["patch_size"] == 14 - assert info["supported_spacing_um"] == pytest.approx(0.25) + assert info["supported_spacing_um"] == pytest.approx(0.5) assert info["precision"] == "fp32" assert info["source"] == "majiabo/GPFM" assert info["output_variants"]["default"]["encode_dim"] == 1024 From 1d23080f84dd64ce2f0a0ed2e027979246bcc08e Mon Sep 17 00:00:00 2001 From: clemsgrs Date: Tue, 30 Jun 2026 16:22:09 +0200 Subject: [PATCH 3/3] Add dense feature extraction for GenBio-PathFM and trim transforms GenBio-PathFM now supports dense (2D-grid) extraction via the model's forward_with_patches path (per-channel patch grids concatenated along the feature dim -> (B, 4608, 14, 14) at 224; patch size 16). Attention-map extraction stays unsupported: the backbone uses fused SDPA with no exposed weights and encodes the three colour channels independently, so there is no single coherent CLS-over-patches attention to recover. Also align the GenBio pooled transform with its model card (Resize((224, 224)) + custom Normalize; drop the CenterCrop, a no-op after a square resize) and inline the patch_size/encode_dim literals to match the house style. --- docs/models.rst | 20 ++++--- slide2vec/encoders/models/genbio.py | 81 ++++++++++++++++++++++++++--- tests/test_encoder_registry.py | 6 +-- tests/test_gpfm_genbio_heavy.py | 15 +++++- tests/test_patch_size_metadata.py | 10 ++-- 5 files changed, 109 insertions(+), 23 deletions(-) diff --git a/docs/models.rst b/docs/models.rst index efd27c8..248fe6f 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -77,7 +77,7 @@ Tile-level encoders - `GPFM `_ - 1024 - ``0.5`` - - Ma et al. (2024); MIT license + - Ma et al. (2024) * - ``uni`` - `UNI `_ - 1024 @@ -127,7 +127,7 @@ Tile-level encoders - `GenBio-PathFM `_ - 4608 - ``0.5`` - - GenBio AI (2024); custom Community License; non-ImageNet normalization + - GenBio AI (2024) Dense tile grids ~~~~~~~~~~~~~~~~ @@ -137,10 +137,10 @@ Dense tile feature extraction is available on tile encoders that implement instead of the pooled ``(B, D)`` tensor returned by ``encode_tiles``. The following built-in tile presets are covered by the dense encoder interface: -``conch``, ``conchv15``, ``gigapath``, ``gpfm``, ``h0-mini``, ``h-optimus-0``, -``h-optimus-1``, ``hibou-b``, ``hibou-l``, ``lunit``, ``midnight``, ``mstar``, -``musk``, ``phikon``, ``phikonv2``, ``prost40m``, ``uni``, ``uni2``, -``virchow``, and ``virchow2``. +``conch``, ``conchv15``, ``genbio-pathfm``, ``gigapath``, ``gpfm``, ``h0-mini``, +``h-optimus-0``, ``h-optimus-1``, ``hibou-b``, ``hibou-l``, ``lunit``, +``midnight``, ``mstar``, ``musk``, ``phikon``, ``phikonv2``, ``prost40m``, +``uni``, ``uni2``, ``virchow``, and ``virchow2``. Notes: @@ -158,6 +158,9 @@ Notes: - H-Optimus dense extraction at non-native input sizes requires ``dynamic_img_size=True`` and ``allow_non_recommended_settings=True`` when constructing the encoder. +- ``genbio-pathfm`` is a single-channel ViT: its dense grid (via + ``forward_with_patches``) concatenates the three per-colour-channel patch grids + along the feature dim, so ``d`` = 4608 matches the pooled output dimension. Dense attention maps ~~~~~~~~~~~~~~~~~~~~~ @@ -177,6 +180,11 @@ Notes: - ``musk`` is **not** covered: its BEiT3 backbone uses a non-timm attention module, so attention extraction raises ``NotImplementedError`` (dense patch-token extraction is still available). +- ``genbio-pathfm`` is **not** covered: it computes attention with a fused + ``scaled_dot_product_attention`` (no materialized weights) and encodes the three + colour channels as independent single-channel images, so there is no single + coherent CLS-over-patches attention to extract; ``encode_tiles_attention`` raises + ``NotImplementedError`` (dense patch-token extraction is still available). - ``hibou-b`` / ``hibou-l`` carry register tokens; pass ``include_registers=True`` to add their query rows as extra channels. - ``conch`` / ``conchv15`` recover attention from their inner timm ViT trunk, the diff --git a/slide2vec/encoders/models/genbio.py b/slide2vec/encoders/models/genbio.py index 2100671..f3aa452 100644 --- a/slide2vec/encoders/models/genbio.py +++ b/slide2vec/encoders/models/genbio.py @@ -21,8 +21,23 @@ ``(1, 4608)``; see ``tests/test_gpfm_genbio_heavy.py``.) Normalization is **non-ImageNet** (``config.json`` ``image_mean`` / ``image_std``) -and must be set explicitly. No dense (spatial-grid) path is exposed in this -encoder, so it declares no ``patch_size``. +and must be set explicitly. + +Dense (spatial-grid) extraction is supported via the model's +``forward_with_patches`` (a DINOv2-style ``x_norm_patchtokens`` path): it returns +the fused per-channel patch tokens ``(B, T, 4608)`` — the three single-channel +patch-token grids concatenated along the feature dim, with the prefix tokens +(CLS + storage tokens) already stripped — which fold straight into a +``(B, 4608, h, w)`` grid. The patch size is 16 (a 224 tile -> a 14x14 = 196 grid). + +Attention-map extraction is **not** supported. The backbone computes attention +with a fused ``F.scaled_dot_product_attention`` (no materialized weights, no +``output_attentions``), and — more fundamentally — it encodes the three colour +channels as three independent single-channel images, so there is no single +coherent CLS-over-patches attention to extract: any "attention grid" would be +three separate grayscale-channel attentions. Recovering it would need a bespoke +per-channel recompute path that diverges from the shared timm/HF attention +helpers, so GenBio deliberately opts out of ``encode_tiles_attention``. """ from typing import Callable @@ -34,6 +49,7 @@ from slide2vec.encoders.base import ( TileEncoder, preferred_default_device, + reshape_tokens_to_grid, resolve_requested_output_variant, ) from slide2vec.encoders.registry import register_encoder @@ -42,15 +58,16 @@ # Non-ImageNet normalization from the model's config.json (image_mean / image_std). _GENBIO_MEAN = (0.697, 0.575, 0.728) _GENBIO_STD = (0.188, 0.240, 0.187) -# embed_dim (1536) x 3 colour channels — see module docstring. -_GENBIO_ENCODE_DIM = 4608 @register_encoder( "genbio-pathfm", - output_variants={"default": {"encode_dim": _GENBIO_ENCODE_DIM}}, + # encode_dim 4608 = embed_dim (1536) x 3 colour channels; patch_size 16 (a 224 + # tile -> a 14x14 = 196 patch-token grid) — see module docstring. + output_variants={"default": {"encode_dim": 4608}}, default_output_variant="default", input_size=224, + patch_size=16, supported_spacing_um=0.5, # 20x; card states no magnification, house default precision="fp32", # upstream runs plain fp32, no autocast source="genbio-ai/genbio-pathfm", @@ -64,10 +81,21 @@ def __init__(self, *, output_variant: str | None = None): self._output_variant = resolve_requested_output_variant(output_variant) def get_transform(self) -> Callable: + # Mirrors the model card: Resize to an exact 224x224 (tuple form, as the + # card's ``Resize((224, 224))``) + custom Normalize. No CenterCrop: the card + # has none, and after a square resize it would be a no-op anyway. + return v2.Compose([ + v2.ToImage(), + v2.Resize((224, 224), interpolation=v2.InterpolationMode.BICUBIC, antialias=True), + v2.ToDtype(torch.float32, scale=True), + v2.Normalize(mean=_GENBIO_MEAN, std=_GENBIO_STD), + ]) + + def get_dense_transform(self) -> Callable: + # Normalization only — no Resize/CenterCrop (see TileEncoder.get_dense_transform), + # so the dense grid stays registered to the full source tile. return v2.Compose([ v2.ToImage(), - v2.Resize(224, interpolation=v2.InterpolationMode.BICUBIC, antialias=True), - v2.CenterCrop(224), v2.ToDtype(torch.float32, scale=True), v2.Normalize(mean=_GENBIO_MEAN, std=_GENBIO_STD), ]) @@ -76,9 +104,46 @@ def encode_tiles(self, batch: Tensor) -> Tensor: # Canonical forward: per-channel CLS tokens concatenated -> (B, 4608). return self._model(batch) + @property + def patch_size(self) -> tuple[int, int]: + return (16, 16) + + def encode_tiles_dense(self, batch: Tensor) -> Tensor: + """Encode tiles into a dense spatial grid. (B, C, H, W) -> (B, d, h, w). + + Uses the model's ``forward_with_patches`` (DINOv2-style), which returns the + fused per-channel patch tokens ``(B, T, 4608)`` — the three single-channel + patch-token grids concatenated along the feature dim, with the prefix + tokens (CLS + storage tokens) already stripped — then folds that token + sequence back into its spatial grid (``num_prefix_tokens=0``). ``H, W`` must + be divisible by the patch size; non-224 inputs rely on the backbone's + DINOv2 positional-embedding interpolation. + """ + if batch.ndim != 4: + raise ValueError( + "encode_tiles_dense expects a (B, C, H, W) batch, got shape " + f"{tuple(batch.shape)}." + ) + _, _, height, width = batch.shape + patch_h, patch_w = self.patch_size + if height % patch_h != 0 or width % patch_w != 0: + raise ValueError( + f"Dense extraction for '{type(self).__name__}' requires input " + f"divisible by the patch size: got {height}x{width}, patch " + f"{patch_h}x{patch_w}. Pad the tile up to a patch multiple first." + ) + _, patch_tokens = self._model.forward_with_patches(batch) + return reshape_tokens_to_grid( + patch_tokens, + grid_h=height // patch_h, + grid_w=width // patch_w, + num_prefix_tokens=0, + encoder_name=type(self).__name__, + ) + @property def encode_dim(self) -> int: - return _GENBIO_ENCODE_DIM + return 4608 @property def device(self) -> torch.device: diff --git a/tests/test_encoder_registry.py b/tests/test_encoder_registry.py index fc67f7e..8ac1eaf 100644 --- a/tests/test_encoder_registry.py +++ b/tests/test_encoder_registry.py @@ -149,9 +149,9 @@ def test_genbio_pathfm_metadata_contract(): info = encoder_registry.info("genbio-pathfm") assert info["level"] == "tile" assert info["input_size"] == 224 - # GenBio is a custom AutoModel without a recoverable dense patch grid, so it - # declares no patch_size (not in the dense coverage lists). - assert info["patch_size"] is None + # GenBio exposes a dense patch grid via forward_with_patches (single-channel + # ViT, patch size 16: a 224 tile -> a 14x14 grid). + assert info["patch_size"] == 16 assert info["supported_spacing_um"] == pytest.approx(0.5) assert info["precision"] == "fp32" assert info["source"] == "genbio-ai/genbio-pathfm" diff --git a/tests/test_gpfm_genbio_heavy.py b/tests/test_gpfm_genbio_heavy.py index 4318e3c..63e44c6 100644 --- a/tests/test_gpfm_genbio_heavy.py +++ b/tests/test_gpfm_genbio_heavy.py @@ -9,9 +9,12 @@ checkpoints: * GPFM: ``GPFM.pth`` loads into the timm DINOv2 ViT-L/14 (``strict=True``) and the - pooled forward returns the registered 1024-d embedding. + pooled forward returns the registered 1024-d embedding; its dense path folds a + 224 tile into a ``(1, 1024, 16, 16)`` patch-token grid (patch size 14). * GenBio-PathFM: the single-channel backbone's canonical forward returns the - registered 4608-d (= embed_dim*3) per-channel-CLS concatenation. + registered 4608-d (= embed_dim*3) per-channel-CLS concatenation; its dense path + (``forward_with_patches``) folds a 224 tile into a ``(1, 4608, 14, 14)`` grid + (patch size 16). """ from __future__ import annotations @@ -34,7 +37,10 @@ def test_gpfm_loads_strict_and_forward_is_1024_d(): batch = torch.randn(1, 3, 224, 224) with torch.inference_mode(): out = encoder.encode_tiles(batch) + dense = encoder.encode_tiles_dense(batch) assert out.shape == (1, 1024) + # patch 14: 224 / 14 = 16 -> a 16x16 grid of 1024-d patch tokens. + assert dense.shape == (1, 1024, 16, 16) @pytest.mark.heavy @@ -45,9 +51,14 @@ def test_genbio_pathfm_forward_is_4608_d(): pytest.skip(f"genbio-pathfm weights unavailable: {type(exc).__name__}: {exc}") encoder = encoder.to("cpu") assert encoder.encode_dim == 4608 + assert encoder.patch_size == (16, 16) batch = torch.randn(1, 3, 224, 224) with torch.inference_mode(): out = encoder.encode_tiles(batch) + dense = encoder.encode_tiles_dense(batch) # Single-channel ViT (in_chans=1): forward concatenates the 3 per-channel CLS # tokens (embed_dim 1536) -> 4608. assert out.shape == (1, 4608) + # Dense via forward_with_patches: the 3 per-channel patch grids concatenated + # along the feature dim; patch 16: 224 / 16 = 14 -> a 14x14 grid of 4608-d tokens. + assert dense.shape == (1, 4608, 14, 14) diff --git a/tests/test_patch_size_metadata.py b/tests/test_patch_size_metadata.py index 570cc4f..2b64c82 100644 --- a/tests/test_patch_size_metadata.py +++ b/tests/test_patch_size_metadata.py @@ -27,6 +27,7 @@ "uni2": (14, 14), "virchow": (14, 14), "virchow2": (14, 14), + "genbio-pathfm": (16, 16), "gigapath": (14, 14), "gpfm": (14, 14), "h-optimus-0": (14, 14), @@ -46,10 +47,11 @@ } # Non-dense tile encoders: real tile encoders with no recoverable patch grid -# (e.g. custom AutoModel backbones whose forward does not expose a spatial -# patch-token sequence), so they declare no ``patch_size`` and opt out of the -# dense interface. -NON_DENSE_TILE_ENCODERS = ["genbio-pathfm"] +# (e.g. a custom AutoModel backbone whose forward exposes no spatial patch-token +# sequence), so they declare no ``patch_size`` and opt out of the dense interface. +# Currently empty — every built-in tile encoder is dense-capable — but the bucket +# is kept so a future non-dense tile encoder has a tracked home. +NON_DENSE_TILE_ENCODERS: list[str] = [] # Non-dense encoders: no recoverable patch grid, so no declared patch_size. NON_DENSE_ENCODERS = ["gigapath-slide", "titan", "prism"] + NON_DENSE_TILE_ENCODERS