Skip to content

Repository files navigation

Prompt-Guided Foundation Model Tuning for Pathology Image Classification

Yi Lin1,*, Zhengjie Zhu1,*, Kwang-Ting Cheng1, Hao Chen1,†

1The Hong Kong University of Science and Technology  ·  *Equal contribution  ·  †Corresponding author

Medical Image Analysis, 2026


Official implementation of PAMT (Prompt-guided Adaptive Model Transformation), from our paper in Medical Image Analysis. PAMT adapts frozen pre-trained foundation models to histopathology whole-slide image (WSI) classification through dual-level reprogramming.

Input Reprogramming

  • RPS — Representative Patch Sampling: keep only the top-K most informative patches per slide (via attention scores), drastically reducing cost and enabling end-to-end training.
  • PVP — Prototypical Visual Prompts: cluster the sampled patches and add a cluster-specific learnable prompt to each patch, capturing intra-slide tissue heterogeneity.

Model Reprogramming

  • AMT — Adaptive Model Transformation: lightweight adapter modules inserted into the (frozen) backbone, fine-tuned together with the visual prompts and the MIL classifier for domain-specific feature learning.

In the paper, PAMT is validated on 14 public datasets across multiple cancer types and backbones (ResNet-50 and ViT-based foundation models such as PLIP and CONCH).

Scope of this release: this repository provides the ResNet-50 (CNN) WSI-classification pipeline — Representative Patch Sampling, Prototypical Visual Prompts, and the adapter-based AMT backbone (resnet50_prompt), plus the six MIL aggregators used as baselines. The ViT / foundation-model backbones and the additional PEFT baselines (VPT, LoRA, AdaptFormer, SSF, …) reported in the paper will be released separately at a later date.

Project Structure

PAMT/
├── dataset/          # Dataset loaders (feature / H5 / WSI-bag, with padding variants)
├── labels/           # Label CSVs and label-conversion utilities (BRACS, TCGA, TUPAC16)
├── mil_models/       # MIL aggregators: att_mil, model_clam, wikg, mean_max_mil, dtfd/
├── models/           # Backbones & prompters: resnet_baseline, resnet_prompt (AMT), prompters
├── optimizer/        # Custom optimizers/schedulers: radam, lookahead, cosine_lr
├── scripts/          # Pipeline scripts (see Pipeline below)
├── utils/            # Metrics, evaluation, clustering, IO helpers
├── requirements.txt
├── setup.py
└── LICENSE

Installation

conda create -n pamt python=3.9 -y
conda activate pamt
pip install -r requirements.txt

Notes on a few dependencies:

  • faiss (used in utils/clustering.py) needs a GPU build; install via conda: conda install -c pytorch faiss-gpu.
  • openslide-python requires the OpenSlide C library (conda install -c conda-forge openslide or your system package manager).
  • The clam MIL method imports topk.svm.SmoothTop1SVM from CLAM's topk package (https://github.com/mahmoodlab/CLAM); it is optional and only needed if you run --mil_method clam.

Supported Datasets & Models

Datasets: the WSI-classification scripts here support Camelyon16, TCGA-NSCLC, TCGA-RCC, BRACS and TUPAC16 (the full paper evaluates 14 public datasets).

MIL aggregators (--mil_method): mean, max, abmil (attention MIL), clam, wikg. DTFD-MIL is trained through the dedicated 3_train_res_dtfd.py script.

Pipeline

The pipeline goes from raw WSIs to final metrics in four stages. All paths shown as /path/to/... are placeholders — pass your own via the CLI flags.

Phase 1 — Top-patch selection (DTFD attention). Train a DTFD model on pre-extracted features to score patches and keep the top-K per slide:

python scripts/1_get_top_patch_dtfd.py --num_cls 2 --topk 128 --seed 32

Phase 2 — Crop the selected patches and extract features for clustering:

python scripts/2_crop_patch_reduce.py \
    --dataset Camelyon16 \
    --wsi_dir /path/to/WSI/CAMELYON16 \
    --coord_pkl /path/to/topk_coords \
    --save_dir /path/to/Camelyon16/patches_topk

python scripts/2_extract_feature_for_cluster.py \
    --dataset Camelyon16 \
    --test_data /path/to/Camelyon16/patches_topk \
    --split_csv_path ./labels/... \
    --save_dir /path/to/Camelyon16/reduced_features

Phase 3a — Build prototypes and match patches (RPS):

python scripts/3_cluster.py --dataset Camelyon16 --clustering \
    --num_prototypes 4 --num_shift_vectors 200 \
    --feature_path /path/to/reduced_features --save_path /path/to/prototypes
python scripts/3_cluster.py --dataset Camelyon16 --matching \
    --existing_prototypes_path /path/to/prototypes --new_save_path /path/to/matched

Phase 3b — Train PAMT. Train the prompt-adapted ResNet backbone with a MIL aggregator:

python scripts/3_train_res_mils.py \
    --dataset Camelyon16 \
    --mil_method abmil \
    --res resnet50_prompt \
    --pretrained_res imagenet \
    --num_cls 2 \
    --test_data /path/to/Camelyon16/patches_topk \
    --val_data  /path/to/Camelyon16/patches_topk \
    --split_csv_path ./labels/... \
    --train_match_path /path/to/matched \
    --prompt_block_num 6 --reduction 16 \
    --log_dir ./exp

For the DTFD-MIL variant use scripts/3_train_res_dtfd.py with the same data flags.

Key training flags: --dataset, --mil_method, --res (resnet50_prompt enables AMT), --pretrained_res (imagenet or a checkpoint path), --num_cls, --prompt_block_num / --reduction (adapter config), --fine_tune / --fine_tune_type. Run a script with -h for the full list.

License

This project is licensed under the MIT License — see LICENSE.

Citation

If you use this code, please cite our paper (add volume/pages/DOI once available):

@article{lin2026prompt,
  title     = {Prompt-Guided Foundation Model Tuning for Pathology Image Classification},
  author    = {Lin, Yi and Zhu, Zhengjie and Cheng, Kwang-Ting and Chen, Hao},
  journal   = {Medical Image Analysis},
  year      = {2026},
  publisher = {Elsevier}
}

Yi Lin and Zhengjie Zhu contributed equally.

About

The official code for "Prompt-Guided Foundation Model Tuning for Pathology Image Analysis"

Resources

Stars

6 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages