Dense prediction

This notebook walks through soma’s dense flow — per-pixel and per-point prediction on a frozen foundation-model token grid:

Dataset(+masks/points) -> DenseTileFeatureExtractor -> train (decoder + head) -> evaluate

It mirrors the slide-level walkthrough, but the encoder now emits a token grid per tile instead of one vector per slide, and a decoder (not a MIL aggregator) upsamples that grid. We work segmentation end to end, then show that detection is the same flow with a different head + point supervision.

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)

Dense supervision lives in per-sample files, not a scalar label:

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

• detection — sample_id, image_path, points_path; the points file is a CSV of x, y, class in ROI-pixel coordinates.

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

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-dense-tutorial-'))
ROIS = WORK / 'rois'; MASKS = WORK / 'masks'; POINTS = WORK / 'points'
for d in (ROIS, MASKS, POINTS): 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)

def make_points(path):
    # a few cells per class, in ROI-pixel coordinates
    pts = [(56, 56, 0), (112, 112, 1), (160, 160, 1)]
    pd.DataFrame(pts, columns=['x', 'y', 'class']).to_csv(path, index=False)

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')
    make_points(POINTS / f'{sid}.csv')

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
# task-specific manifests below (masks for segmentation, points for detection).
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))

segmentation manifest:
sample_id                                       image_path                                         mask_path
    roi00 /tmp/soma-dense-tutorial-xblg322u/rois/roi00.tif /tmp/soma-dense-tutorial-xblg322u/masks/roi00.png
    roi01 /tmp/soma-dense-tutorial-xblg322u/rois/roi01.tif /tmp/soma-dense-tutorial-xblg322u/masks/roi01.png
    roi02 /tmp/soma-dense-tutorial-xblg322u/rois/roi02.tif /tmp/soma-dense-tutorial-xblg322u/masks/roi02.png

1. Extract dense token grids

DenseTileFeatureExtractor reads each ROI at spacing_um and runs the frozen encoder to produce a (feature_dim, gh, gw) grid per sample, stored in a DenseFeatureStore. With phikon at 224 px and patch-16 that’s a 14×14 grid.

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

extractor = DenseTileFeatureExtractor(
    Dataset(extract_csv),
    EncoderConfig(name='phikon'),
    target_size=SIZE,
    spacing_um=SPACING,
    backend='openslide',
    cache=CacheConfig(enabled=False),
)
dense_store = extractor.run(str(WORK / 'dense'))
print('dense store ready for', len(dense_store.available_samples), 'ROIs')

Note: Environment variable`HF_TOKEN` is set and is the current active token independently from the token you've just configured.

Dense extraction mode: whole (single padded forward)
dense store ready for 8 ROIs

2. Train segmentation (decoder + head)

train(dataset_type='segmentation', ...) builds a decoder (lightweight_conv upsamples the token grid) plus a parameter-free segmentation head that crops to the mask and scores Dice / IoU. SegmentationManifest is the dense counterpart of Dataset.

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

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

seg_result = train(
    feature_store=dense_store,
    dataset=seg_manifest,
    splits=seg_splits,
    dataset_type='segmentation',
    decoder=DecoderConfig(name='lightweight_conv'),
    task=TaskConfig(name='segmentation', params={'num_classes': NUM_CLASSES}),
    training=TrainingConfig(epochs=3, batch_size=2, learning_rate=1e-3, seed=0),
    evaluation=EvalConfig(metrics=['mean_dice', 'mean_iou']),
    # our PNG masks carry no spacing metadata; declare the ROI spacing so they
    # register against the grids extracted at the same spacing.
    preprocessing=PreprocessingConfig(requested_spacing_um=SPACING),
    run_dir=str(WORK / 'runs' / 'segmentation'),
)
print('segmentation run dir:', seg_result.run_dir)

segmentation run dir: /tmp/soma-dense-tutorial-xblg322u/runs/segmentation

3. Switch to detection — same grid, point supervision

Detection reuses the same dense extraction; only the supervision and head change. TaskConfig('detection') renders each annotated point as a peak Gaussian, the decoder smooths the grid into a peak heatmap, and the head recovers points (local-maxima + NMS) scored with F1 at a matching distance δ. match_distance and sigma are given in µm.

from soma.dataset import DetectionManifest

det_csv = WORK / 'det.csv'
pd.DataFrame({'sample_id': ids,
              'image_path': [str(ROIS / f'{s}.tif') for s in ids],
              'points_path': [str(POINTS / f'{s}.csv') for s in ids]}).to_csv(det_csv, index=False)

det_manifest = DetectionManifest(det_csv)
det_splits = Splits(splits_csv, det_manifest)

det_result = train(
    feature_store=dense_store,
    dataset=det_manifest,
    splits=det_splits,
    dataset_type='detection',
    decoder=DecoderConfig(name='lightweight_conv'),
    task=TaskConfig(name='detection', params={
        'num_classes': NUM_CLASSES,
        'match_distance': 2.0,   # microns
        'sigma': 0.7,            # microns
    }),
    training=TrainingConfig(epochs=3, batch_size=2, learning_rate=1e-3, seed=0),
    evaluation=EvalConfig(metrics=['mean_f1', 'f1_per_class']),
    preprocessing=PreprocessingConfig(requested_spacing_um=SPACING, requested_tile_size_px=SIZE),
    run_dir=str(WORK / 'runs' / 'detection'),
)
print('detection run dir:', det_result.run_dir)

detection run dir: /tmp/soma-dense-tutorial-xblg322u/runs/detection

4. The one-shot Pipeline equivalent

As with the slide-level flow, Pipeline collapses extract + train + evaluate into a single config-driven call.

(Shown for reference, not executed.)

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

config = PipelineConfig(
    dataset_csv=str(seg_csv),
    splits_csv=str(splits_csv),
    output_root='output/segmentation',
    dataset_type='segmentation',
    preprocessing=PreprocessingConfig(
        backend='openslide', requested_tile_size_px=224, requested_spacing_um=0.5,
    ),
    encoder=EncoderConfig(name='phikon'),
    decoder=DecoderConfig(name='lightweight_conv'),
    task=TaskConfig(name='segmentation', params={'num_classes': 3}),
    training=TrainingConfig(epochs=3, batch_size=2, learning_rate=1e-3),
    evaluation=EvalConfig(metrics=['mean_dice', 'mean_iou']),
    cache=CacheConfig(enabled=True),
)
results = Pipeline(config).run()