gaozijun/he2st-fullpanel-6k

HE2ST full panel (6,054 genes): inference

Predicts single-cell gene expression for 6,054 genes, as log1p(counts per cell), from an H&E image plus cell masks. This repo contains only what you need to run the model: the inference code, the weights, the gene list and per-gene statistics.

• Training data: 165 Xenium slides, each supervised on all of its own panel genes (the loss is masked to the genes that slide measured). Controls, unassigned/deprecated codewords, intergenic probes and non-gene probes (mutation, TCR, HPV, CD45 isoform) were removed.
  • 107 slides are HEST/panst: 15 organs, targeted panels of 280–480 genes, plus Xenium 5K.
  • 58 samples are ST-5K gastric (Xenium Prime 5K, 16 patients).
  • Splits: 125 train / 22 val / 18 test slides.
• Checkpoint: step 130,000, the best validation PCC (0.2919).
• Architecture:
  • Virchow2 ViT-H and a ConvNeXt-B distilled from H-optimus-1, both fine-tuned.
  • A JAFAR attention upsampler to 112 x 112.
  • Per-cell mask pooling plus the tile mean, followed by an MLP.
  • A 32-bin confidence head, which gives a 95% error quantile per cell and gene.

Files

file	content
model.safetensors	weights (fp32, 796M parameters)
config.json	model settings (200 um field of view, 448 px input, 112 px label map)
genes.json	the 6,054 output genes, in output order
gene_stats.npz	per-gene mu, sd of the training targets (the network predicts z-scores; pred * sd + mu = log1p)
he2st/	model code (model.py), tile-level predictor (predict.py), whole-slide pipeline (wsi.py)
scripts/infer_wsi.py	command-line whole-slide inference

Install

huggingface-cli login                      # the repo is private
huggingface-cli download gaozijun/he2st-fullpanel-6k --local-dir he2st-fullpanel-6k
cd he2st-fullpanel-6k
pip install -r requirements.txt            # torch, timm>=1.0, openslide, ...

No access to Virchow2 or H-optimus-1 is needed. Both trunks are built locally and all weights come from model.safetensors.

Whole slide: WSI + cell polygons -> .h5ad

python scripts/infer_wsi.py \
    --model . \
    --wsi slide.svs \
    --cells cells.geojson \
    --out slide_he2st.h5ad \
    --expand_um 4          # only if --cells are nuclei (CellViT, HoVer-Net, StarDist, ...)

--wsi: the H&E image.

• Files that OpenSlide can read (.svs, .ndpi, pyramidal .tif, ...) are read tile by tile.
• Other TIFF / OME-TIFF files are read into memory with tifffile. An example is the untiled, planar (CYX) he_registered_s0.ome.tiff of the ST-5K samples.

--cells: polygons in level-0 pixel coordinates of the WSI.

• GeoJSON:
  • Polygon features are one cell each.
  • MultiPolygon features are split into one cell per polygon. CellViT/CellViT++ write one MultiPolygon per cell class; the class name is kept as obs['cell_type'].
• Parquet with a WKB geometry column, one row per cell, for example HEST *_xenium_cell_seg.parquet.

--mpp: microns per level-0 pixel. It must be right, because the model sees a fixed 200 um field of view.

• By default it is read from the slide: openslide.mpp-x, aperio.MPP, the TIFF resolution, or OME PhysicalSizeX.
• Pass it explicitly when the file does not record it. For example, the ST-5K OME-TIFFs need --mpp 0.2125.

--expand_um: what matters is whether the polygons are whole cells or nuclei, not the file name.

• Training masks were Xenium cell boundaries, so use the default --expand_um 0 for Xenium cell boundaries. Examples: HEST xenium_seg/*_xenium_cell_seg.parquet, or ST-5K cell_boundaries_morph_s0.geojson.
• CellViT's cells.geojson contains nuclei despite its name. For nuclei, use --expand_um 4. The growth is Voronoi-limited, so neighbouring cells never overlap.
• Raw Xenium output (cell_boundaries.parquet) is in microns on the DAPI image, not in H&E pixels. It must be registered to the H&E first (HEST's xenium_seg files already are).

Validated on held-out TENX115:

cell masks	per-gene PCC (mean)	HVG50
Xenium cell boundaries, --expand_um 0	0.333	0.575
Xenium nuclei, --expand_um 4	0.326	0.565

Other options:

• --genes A,B,C or --genes genes.txt: output a subset of genes.
• --q95: also store the error bar (layer q95).
• --bs, --workers, --device: batch size, data-loader workers, device.

Output: an AnnData file.

• X: float16 predictions in log1p(counts per cell); cells x genes.
• obs: cell_id, x, y (centroid, level-0 px), n_pix, and cell_type if available. n_pix is the cell's pixels on the 112 px map, 1.786 um each; cells with n_pix == 0 get NaN.
• obsm['spatial']: the centroids.
• layers['q95'] (with --q95): the 95% quantile of |prediction - truth| in log1p units.

The full output is dense: cells x 6,054 x 2 bytes, for example 6 GB for 500k cells. Use --genes if you need fewer genes.

How it tiles: every cell is predicted once, from the 200 um tile centred on the 100 um grid block that contains its centroid. A predicted cell is therefore always at least 50 um from the tile edge. Polygons are drawn on the label map exactly as in training: vertices mapped to the 112 px map, rounded, and filled with cv2.fillPoly.

Examples on the held-out test slides (reproduces the per-slide test numbers below within ±0.004 PCC):

# HEST Xenium slide (OpenSlide path, pixel size read from the TIFF)
python scripts/infer_wsi.py --model . --wsi HEST/wsis/TENX115.tif \
    --cells HEST/xenium_seg/TENX115_xenium_cell_seg.parquet --out TENX115.h5ad
# ST-5K gastric sample (OME-TIFF read with tifffile; pixel size not stored in the file)
python scripts/infer_wsi.py --model . --wsi S053949_2_2/he_registered_s0.ome.tiff \
    --cells S053949_2_2/cell_boundaries_morph_s0.geojson --mpp 0.2125 --out S053949_2_2.h5ad

Tile level (Python)

from he2st import HE2STPredictor
p = HE2STPredictor.from_pretrained(".", device="cuda")        # or "gaozijun/he2st-fullpanel-6k"
# images: uint8 (B, H, W, 3) RGB tiles covering 200 um x 200 um (any pixel size, resized to 448 inside)
# labels: int (B, 112, 112) cell label maps of the same fields of view (0 = background, cells 1..n)
preds = p.predict(images, labels, genes=["EPCAM", "CD3E", "MKI67"])     # list of (n_b, 3) arrays, log1p
preds, q95 = p.predict(images, labels, return_q95=True)                 # plus error bars, same layout

Held-out test results

Each cell of every test slide is predicted once. Metrics are computed per slide, on the genes that slide measured.

• Targeted panels (6 slides; breast, bowel, lung, skin; 280–419 genes): per-gene PCC 0.33–0.41, HVG50 0.43–0.57.
• Xenium 5K, all ~5,000 genes:
  • Gastric, 11 samples from 4 held-out patients: mean per-gene PCC 0.14. For genes detected in more than 20% of cells it is 0.46. Most 5K genes are very sparse in gastric cells.
  • Lung (TENX189): mean per-gene PCC 0.25; 0.59 for genes detected in more than 20% of cells.
  • Pseudo-bulk (slide-mean profile across genes): r = 0.88–0.98.

Licences

The weights include fine-tuned Virchow2 (CC-BY-NC-ND-4.0) and a ConvNeXt distilled from H-optimus-1. Keep this repo private / non-commercial unless those terms are cleared.