MuPaD-HE2ST

Predict the expression of 50 genes in every cell of an H&E slide, by finetuning MuPaD-512 on Xenium breast cancer slides from HEST-1k.

Code: github.com/jinxixiang/MUPAD → 6-HE2ST/

Model

Expression is painted onto each cell's footprint, giving a 50-channel gene map per 100 µm tile (512 px). A 50-channel VAE (stage 1) compresses it to a 4-channel latent; the SiT-XL trunk (stage 2) predicts that latent from H&E in a single forward pass:

z_he  = RGB_VAE(he)                          H&E latent, 4 extra patch-embed channels
cells = paint(PCA(MUSK(cell crops)))         per-cell MUSK features, 16 channels
ctx   = MUSK(he)[CLS]                        cross-attention
z_st  = -SiT([0 ; z_he ; cells], t=1, ctx)   one-pass readout of a flow-matching model
gene map = ST_VAE.decode(z_st)  ->  mean over each cell's footprint

Each cell is read from the tile it is most central in, over four half-stride grids.

Contents

Laid out to drop straight into exps/:

stage1_vae/
  best/model.safetensors        50-channel ST VAE (335 MB)
  gene_stats.json               per-gene scale and positive fraction (train split)
  best_val.json                 val pixel PCC of the reconstruction (the ceiling)
stage2_sit/
  best.safetensors              SiT trunk, 1.41B params, fp32 (5.6 GB)
  cell_pca.npz                  PCA of the per-cell MUSK features (train cells)
  latent_stats.json             ST / H&E latent mean and std
  best_val.json                 val scores at the selected step

Requires

xiangjx/MuPaD-512 for the trunk architecture, the H&E VAE and MUSK.

Usage

git clone https://github.com/jinxixiang/MUPAD && cd MUPAD/6-HE2ST
pip install -r ../requirements.txt && pip install -e .
huggingface-cli download xiangjx/MuPaD-512 --local-dir pretrained/MuPaD-512
huggingface-cli download xiangjx/MuPaD-HE2ST --local-dir exps

# two processed example slides (NCBI917 train, NCBI915 test)
huggingface-cli download xiangjx/MuPaD-examples --repo-type dataset --include "he2st/*" --local-dir examples
export HEST_ROOT=$PWD/examples/he2st/HEST1k HE2ST_CACHE=$PWD/examples/he2st/cache \
       HE2ST_MUSK_FEATURES=$PWD/examples/he2st/features/musk

python scripts/export_mupad_cells.py --stage 2 --slides NCBI915 --out-dir preds/example
python scripts/score.py --pred-dir preds/example --slides NCBI915 --out eval_results/example.csv

New slides need the tile cache and per-cell MUSK features first (scripts/preprocess_slides.py, python -m he2st.embed_musk); see the code README.

Genes

The 50-gene panel is he2st/panel_50.json: 20 curated breast cancer genes (ERBB2, ESR1, PGR, AR, GATA3, FOXA1, EPCAM, TACSTD2, ANKRD30A, FASN, KRT14, ACTA2, MKI67, TOP2A, PDGFRB, MMP2, CXCL12, PECAM1, CD3E, CD68) and 30 more chosen for detectability and presence on other organs' Xenium panels.

Data

Trained on 27 breast cancer Xenium slides from HEST-1k v1.3.1 (18 subjects); validated on 3 held-out subjects. The split is subject-disjoint, frozen in he2st/splits_breast.json.

License

MIT, following the REPA / SiT lineage the trunk derives from. The training data (HEST-1k) is CC BY-NC-SA 4.0.