Attention-based segmentation¶

This notebook walks through soma’s decoder-free segmentation method — the per-head CLS-attention variant of the dense flow:

Dataset(+masks) -> DenseTileFeatureExtractor(cls_attention) -> train (pixel classifier) -> evaluate

It is the alternative to the neural-decoder path in the dense-prediction walkthrough. Same dataset_type='segmentation' contract — same manifest, splits, spacing-aware mask reader, dense metrics, and prediction artifacts — but only the trainable component changes:

the encoder emits the ViT’s per-head CLS-token self-attention as a dense (K, grid_h, grid_w) grid (feature_kind='cls_attention') instead of the patch-feature grid, and
a lightweight per-pixel classifier (K,) → class (logistic / random forest / XGBoost / MLP) replaces the neural decoder — no decoder, no Trainer, no ``.pt`` checkpoints.

This re-implements Ramchandani et al., Benchmarking Computational Pathology Foundation Models for Semantic Segmentation (arXiv:2602.18747). The Pixel-classifier page in the docs covers the method and the one deliberate divergence (native-window vs resize-to-native extraction).

Tiny synthetic data, runs on CPU, ungated encoder — the numbers are meaningless; the point is the API. We use `phikon <https://huggingface.co/owkin/phikon>`__ at its native 224 px window (a 14×14 token grid), which avoids position-embedding interpolation.

⚠️ Scaffolding (not soma API)¶

Decoder-free segmentation consumes the same dense supervision as the neural-decoder path — a per-sample mask, not a scalar label:

dataset.csv has sample_id, image_path, mask_path; the mask is an integer-class raster the same size as the ROI.

We fabricate small 224 px ROI tiles (the dense flow consumes fixed-size tiles/ROIs, not whole WSIs) plus their integer-class masks.

[ ]:

import logging, warnings
warnings.filterwarnings('ignore')
logging.getLogger().setLevel(logging.ERROR)

import tempfile
from pathlib import Path

import numpy as np
import pandas as pd
import tifffile
from PIL import Image

WORK = Path(tempfile.mkdtemp(prefix='soma-attn-seg-tutorial-'))
ROIS = WORK / 'rois'; MASKS = WORK / 'masks'
for d in (ROIS, MASKS): d.mkdir()
rng = np.random.default_rng(0)

SIZE = 224          # phikon native window
SPACING = 0.5       # microns/pixel
NUM_CLASSES = 3     # 0 = background, 1, 2 = tissue classes

def make_roi(path):
    img = np.clip(np.stack([np.full((SIZE, SIZE), 150),
                            np.full((SIZE, SIZE), 70),
                            np.full((SIZE, SIZE), 160)], -1).astype(np.int16)
                  + rng.integers(-30, 30, (SIZE, SIZE, 3)), 0, 255).astype(np.uint8)
    tifffile.imwrite(path, img, photometric='rgb', tile=(SIZE, SIZE),
                     resolution=(20000, 20000), resolutionunit='CENTIMETER')

def make_mask(path):
    m = np.zeros((SIZE, SIZE), np.uint8)
    m[SIZE // 4:SIZE // 2, SIZE // 4:SIZE // 2] = 1
    m[SIZE // 2:3 * SIZE // 4, SIZE // 2:3 * SIZE // 4] = 2
    Image.fromarray(m).save(path)

ids = [f'roi{i:02d}' for i in range(8)]
for sid in ids:
    make_roi(ROIS / f'{sid}.tif')
    make_mask(MASKS / f'{sid}.png')

split = ['train'] * 4 + ['tune'] * 2 + ['test'] * 2
splits_csv = WORK / 'splits.csv'
pd.DataFrame({'sample_id': ids, 'split': split, 'fold': 0}).to_csv(splits_csv, index=False)

img_paths = [str(ROIS / f'{s}.tif') for s in ids]

# Feature extraction only needs the images; supervision lives in the
# segmentation manifest below (the mask raster per sample).
extract_csv = WORK / 'extract.csv'
pd.DataFrame({'sample_id': ids, 'image_path': img_paths,
              'label': 0}).to_csv(extract_csv, index=False)

seg_csv = WORK / 'seg.csv'
pd.DataFrame({'sample_id': ids, 'image_path': img_paths,
              'mask_path': [str(MASKS / f'{s}.png') for s in ids]}).to_csv(seg_csv, index=False)
print('segmentation manifest:')
print(pd.read_csv(seg_csv).head(3).to_string(index=False))

1. Extract per-head CLS-attention grids¶

This is the one extraction difference from the neural-decoder path. We point DenseTileFeatureExtractor at feature_kind='cls_attention', so instead of the ViT patch-feature grid it captures the CLS-token self-attention of the chosen block(s) — one (grid_h × grid_w) map per head. With phikon at 224 px and patch-16 that’s a 14×14 grid; the channel count K is num_blocks × num_heads (per-head is preserved — head specialization is the signal the classifier exploits).

AttentionConfig(blocks=[-1]) takes the last block (the paper’s choice); include_registers=False keeps only the CLS query row. Everything else — the sliding/stitching, cache, and DenseFeatureStore — is shared with the decoder path, because an attention grid is structurally just another (K, gh, gw) dense grid.

[ ]:

from soma import (
    Dataset, DenseTileFeatureExtractor, EncoderConfig, AttentionConfig,
    CacheConfig, PreprocessingConfig,
)

extractor = DenseTileFeatureExtractor(
    Dataset(extract_csv),
    EncoderConfig(name='phikon'),
    target_size=SIZE,
    spacing_um=SPACING,
    backend='openslide',
    cache=CacheConfig(enabled=False),
    # feature_kind='cls_attention' is what makes this the decoder-free path:
    # the encoder emits per-head CLS-attention maps, not the patch grid.
    preprocessing=PreprocessingConfig(
        feature_kind='cls_attention',
        attention=AttentionConfig(blocks=[-1], include_registers=False),
        requested_spacing_um=SPACING,
    ),
)
attn_store = extractor.run(str(WORK / 'attention'))
print('attention store ready for', len(attn_store.available_samples), 'ROIs')

2. Train a per-pixel classifier (no decoder)¶

train(dataset_type='segmentation', pixel_classifier=...) selects the decoder-free path: it upsamples the per-head attention grid to the mask’s resolution and fits a per-pixel classifier (K,) → class on class-stratified sampled pixels, then predicts every pixel at evaluation. pixel_classifier and decoder are mutually exclusive under dataset_type='segmentation'.

We use logistic (multinomial logistic regression) because it is the lightest and has no extra dependency; swap in random_forest, xgboost, or mlp by name. The classifiers own their own training loop (no torch Trainer, no .pt fold checkpoints on this path). TrainingConfig.max_train_pixels is the class-stratified pixel budget for the fit.

[ ]:

from soma.dataset import SegmentationManifest
from soma import (
    Splits, PixelClassifierConfig, TaskConfig, TrainingConfig, EvalConfig, train,
)

seg_manifest = SegmentationManifest(seg_csv)
seg_splits = Splits(splits_csv, seg_manifest)

seg_result = train(
    feature_store=attn_store,
    dataset=seg_manifest,
    splits=seg_splits,
    dataset_type='segmentation',
    # the decoder-free trainable component — XOR `decoder=...`
    pixel_classifier=PixelClassifierConfig(
        name='logistic',                # logistic | random_forest | xgboost | mlp
        params={'max_iter': 200},
    ),
    task=TaskConfig(name='segmentation', params={'num_classes': NUM_CLASSES}),
    training=TrainingConfig(epochs=1, batch_size=2, seed=0, max_train_pixels=50_000),
    evaluation=EvalConfig(metrics=['mean_dice', 'mean_iou']),
    # our PNG masks carry no spacing metadata; declare the ROI spacing so they
    # register against the attention grids extracted at the same spacing. The
    # cross-default also resolves feature_kind=cls_attention from the classifier.
    preprocessing=PreprocessingConfig(
        feature_kind='cls_attention', requested_spacing_um=SPACING,
    ),
    run_dir=str(WORK / 'runs' / 'attention_segmentation'),
)
print('attention-segmentation run dir:', seg_result.run_dir)

3. The one-shot `Pipeline` equivalent¶

As with the other walkthroughs, Pipeline collapses extract + train + evaluate into a single config-driven call. The only things that select the decoder-free path are preprocessing.feature_kind='cls_attention' (auto when a pixel_classifier is set) and naming a pixel_classifier instead of a decoder.

(Shown for reference, not executed.)

from soma import (
    Pipeline, PipelineConfig, PreprocessingConfig, EncoderConfig, AttentionConfig,
    PixelClassifierConfig, TaskConfig, TrainingConfig, EvalConfig, CacheConfig,
)

config = PipelineConfig(
    dataset_csv=str(seg_csv),
    splits_csv=str(splits_csv),
    output_root='output/attention-segmentation',
    dataset_type='segmentation',
    preprocessing=PreprocessingConfig(
        backend='openslide', requested_tile_size_px=224, requested_spacing_um=0.5,
        feature_kind='cls_attention',         # auto when pixel_classifier is set
        attention=AttentionConfig(blocks=[-1], include_registers=False),
        dense_window_size=None,               # null=whole; =native input for native-window mode
    ),
    encoder=EncoderConfig(name='phikon'),
    pixel_classifier=PixelClassifierConfig(name='logistic'),   # XOR decoder=...
    task=TaskConfig(name='segmentation', params={'num_classes': 3}),
    training=TrainingConfig(max_train_pixels=2_000_000),
    evaluation=EvalConfig(metrics=['mean_dice', 'mean_iou']),
    cache=CacheConfig(enabled=True),
)
results = Pipeline(config).run()

Attention-based segmentation¶

⚠️ Scaffolding (not soma API)¶

1. Extract per-head CLS-attention grids¶

2. Train a per-pixel classifier (no decoder)¶

3. The one-shot Pipeline equivalent¶

3. The one-shot `Pipeline` equivalent¶