Skip to content

Repository files navigation

telomere_uni

A config-driven framework for predicting slide-level labels — a biomarker, age, a pathology finding — from H&E whole-slide images (WSIs). It fine-tunes a pathology foundation model (UNI, UNI2-h, Prov-GigaPath, or any other timm/HuggingFace vision backbone) with a pluggable method (LoRA, full fine-tuning, or a frozen-feature linear probe) and a pluggable aggregation strategy (mean-pool or attention-MIL) to go from a folder of .svs files to a trained, evaluated model.

Bring your own WSIs. The pipeline was originally built on GTEx pancreas telomere length, and that remains the reference implementation, but the generic path (below) works for any cohort with WSIs and a slide-level CSV of labels — no code changes required, only a config file.

On a real benchmark task (lung congestion classification), this framework's frozen-feature linear probe (test AUC 0.807) beat two independently-reimplemented external pipelines run on the identical split — real CLAM (0.787) and real TransMIL (0.787) — and its LoRA fine-tuning (0.797) came closest of any method tried. See Benchmarks below.

Quick start — bring your own WSIs

No GPU needed to try the plumbing first:

python scripts/smoke_test_framework.py   # synthetic WSIs, full pipeline, no GPU/data required

Then, on your own data:

# 1. Scaffold a project config + folders from a slide directory.
python src/preprocess/make_project_config.py my_study \
  --wsi-dir /path/to/my/wsis \
  --raw-column biomarker_score          # or --task classification

# 2. Segmentation + patching + feature extraction (GPU). Prints the exact
#    Trident command for your project — nothing is auto-run for you here.
bash scripts/run_project.sh configs/my_study.yaml trident

# 3. Fill in your labels: one row per slide.
#    Edit data/projects/my_study/user_labels.csv
#    (slide_id, patient_id, <target column>[, age, sex])

# 4. Preflight, then preprocess -> baselines -> train.
bash scripts/run_project.sh configs/my_study.yaml preflight
bash scripts/run_project.sh configs/my_study.yaml preprocess
bash scripts/run_project.sh configs/my_study.yaml baselines
bash scripts/run_project.sh configs/my_study.yaml train      # GPU

# 5. Re-run eval on a specific checkpoint any time.
bash scripts/run_project.sh configs/my_study.yaml eval

Everything downstream of step 1 is driven entirely by configs/my_study.yaml — see Architecture for what the pipeline stages look like and what every config field does.

Full user contract (exact user_labels.csv schema, every generated file, every config field, current limitations): docs/FRAMEWORK.md

Setup

conda create -n telomere python=3.10 && conda activate telomere
conda install openslide && pip install openslide-bin
# Install torch/cuda matching your GPU, then:
pip install -r requirements-train.txt
pip install git+https://github.com/mahmoodlab/trident.git

# UNI and UNI2-h are gated on HuggingFace — request access on the model
# card, then log in once:
huggingface-cli login

Run python src/preflight.py after setup and every time you start on a new machine — it checks the environment, HuggingFace access, OpenSlide, and pipeline state in ~30s, before you're mid-pipeline.

Repo structure

configs/                 every run's config (see docs/ARCHITECTURE.md for the full catalog)
configs/template.yaml       scaffold for a new regression project
configs/template_classification.yaml   scaffold for a new classification project
data/projects/<name>/    per-project manifest/labels/splits/norm_stats (generic path)
src/preprocess/          make_project_config.py, Trident QC + patch sampling, labels, splits
src/backbones.py         backbone registry (uni, uni2-h, prov-gigapath, any timm/HF-hub id)
src/datasets.py          reads patches from WSIs on the fly (OpenSlide)
src/model.py             backbone + method (lora/full/linear_probe) + regression/classification head
src/train.py             fine-tuning loop
src/eval.py              slide-level inference + metrics + sanity plots
src/baselines.py         mean / age / age+sex / linear-probe baselines
src/preflight.py         environment + access + pipeline-state check
scripts/run_project.sh   generic pipeline entrypoint (preflight/trident/preprocess/baselines/train/eval)
scripts/smoke_test_framework.py   local end-to-end test, no GPU/data required
experiments/             one folder per run (config + checkpoints + preds)
docs/ARCHITECTURE.md     pipeline diagram, full config catalog, config version history
docs/FRAMEWORK.md        generic cohort guide (full contract) + benchmark writeup

Architecture & configs

Every stage of the pipeline — patch sampling, backbone, fine-tuning method, aggregation, training hyperparameters, output paths — is driven by one YAML config file; nothing is hardcoded. docs/ARCHITECTURE.md has the full pipeline diagram, a catalog of every config in configs/ and what it's for, and how the config set evolved from the first pancreas run to the current pluggable-backbone framework.

Benchmarks

Lung congestion classification (GTEx, n=463, identical train/val/test split across every row — see docs/FRAMEWORK.md for the full setup and how each comparator was run):

Approach Test AUC
Probe (frozen UNI + LogisticRegression) 0.807
method: lora (this framework, default hyperparameters) 0.797
method: full, tuned LR 0.789
method: linear_probe (this framework, live-pixel path) 0.787
Real CLAM (mahmoodlab/CLAM) 0.787
Real TransMIL (szc19990412/TransMIL) 0.787
Prov-GigaPath backbone + method: linear_probe 0.781
ImageNet ViT backbone + method: linear_probe (non-pathology plumbing check) 0.761

Reference implementation: GTEx telomere/age (legacy)

The framework's original use case: predicting bulk telomere length or age from GTEx pancreas H&E slides, patch-level regression averaged to slide level. Kept working and still the best-validated path end to end.

Switching target (telomere vs age)

Both targets are z-scored at split time and stored in data/splits.csv (log_z_telomere and z_age), and data/norm_stats.json carries per-target mean/std. Change two lines in configs/v1_lora.yaml:

target:
  column: log_z_telomere   # telomere model
  log_space: true

or

target:
  column: z_age            # age model
  log_space: false

Pipeline (run order — from the repo root)

# 0. Preflight (see Setup above).
python src/preflight.py

# 1. Download a small GTEx pancreas subset.
python src/preprocess/download_gtex.py

# 2. Trident: segmentation + patching + UNI feature extraction (on the cluster).
python run_batch_of_slides.py --task all \
    --wsi_dir <gtex_svs_dir> --job_dir trident_output \
    --patch_encoder uni_v1 --mag 20 --patch_size 256 --overlap 0

# 3. QC tissue masks + patch counts (mandatory).
#    - stages 20 random viz JPGs in data/qc/masks_sample/
#    - writes data/qc/patch_counts.csv + patch_counts.png
#    Eyeball the masks; list any clearly bad slide_ids in
#    data/excluded_slides.txt (one per line). If >=10% are bad,
#    talk to a collaborator before retuning Trident.
python src/preprocess/qc_masks.py

# 4. Subsample ~200 patches/slide -> data/manifest.csv
#    (skips slides in data/excluded_slides.txt and any with <50 patches).
python src/preprocess/sample_patches.py

# 5. Build labels.csv from telomere + phenotype (donor parsed from slide_id).
python src/preprocess/make_labels.py \
    --telomere_file <demanelis.csv> --phenotype_file <gtex_subj_pheno.txt> \
    --tissue Pancreas

# 6. Patient-level train/val/test split + per-target z-score stats.
python src/preprocess/make_splits.py

# 7. Baselines (mean / age / age+sex / linear probe) - run BEFORE LoRA.
python src/baselines.py --config configs/v1_lora.yaml

# 8. LoRA fine-tuning.
python src/train.py --config configs/v1_lora.yaml

# 9. Slide-level evaluation on the test set.
#    Run dirs are auto-suffixed by target: experiments/<name>_telomere or
#    experiments/<name>_age. Point --checkpoint at whichever you trained.
python src/eval.py --config configs/v1_lora.yaml \
    --checkpoint experiments/v1_lora_telomere/best.pt --split test

Notes:

  • Everything is config-driven via configs/v1_lora.yaml; never hardcode params.
  • Splits are by patient to avoid leakage. Seeds are fixed for reproducibility.
  • The linear-probe Pearson r on test is your headline baseline. If LoRA does not beat it, LoRA isn't doing its job.

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages