Test-Time Instance Selection for Improved Whole Slide Image Analysis

Official implementation of Test-Time Instance Selection for Improved Whole Slide Image Analysis at the 2nd MICCAI Workshop on Efficient Medical AI, MICCAI 2026.

Test-Time Instance Selection (TTIS) is a training-free, plug-and-play framework for whole slide image (WSI) classification. At inference time, TTIS groups patch features using hierarchical spatio-semantic clustering, selects representative subsets using farthest point sampling, and combines centroid-based, attention-based, and class-guided slide representations. It can be applied to a trained multiple instance learning (MIL) model without retraining or changing its architecture.

This repository supports ABMIL, CLAM-SB, R2T-MIL, S4MIL, TransMIL, and WiKG. The paper evaluates these models on CAMELYON16, TCGA-BRCA, TCGA-ESCA, TCGA-NSCLC, and TCGA-RCC using UNI patch features and MUSK for class-guided sampling.

Repository structure

configs/
├── models/                         Model, dataset, and training configurations
├── cluster/                        Offline clustering configurations
└── prompts/musk/                   Class prompts used by MUSK
src/
├── datasets/                       Dataset loaders
├── models/                         MIL backbones
│   └── model_interface.py          Train/test routing, checkpoints, and metrics
└── ttis/
    ├── clustering.py               Hierarchical spatio-semantic clustering
    ├── sampling.py                 Random sampling and TTIS patch selection
    ├── aggregation.py              Multi-view weighting
    └── zeroshot_classifier.py      MUSK prompt scoring
wsi_dataset_annotation/             Dataset metadata and cross-validation splits
train.py                            Training and evaluation entry point
extract_offline_clusters.py         Offline cluster extraction entry point

Installation

The tested environment uses Python 3.12, PyTorch 2.7.1, and CUDA 11.8.

conda create -n ttis_cu118 python=3.12 -y
conda activate ttis_cu118
pip install -r requirements.txt

A CUDA-capable GPU is required for offline clustering and TTIS inference. MUSK weights are downloaded from Hugging Face during the first TTIS run if they are not already cached.

Data preparation

The precomputed features used in the paper are not included in this repository. We used the TRIDENT WSI processing pipeline for tissue segmentation, patching, and UNI feature extraction. Follow the TRIDENT instructions to process the raw WSIs, then organize each slide as <slide_id>.pt:

{
    "features": torch.Tensor,  # [num_patches, feature_dim]
    "coords": torch.Tensor     # [num_patches, 2]
}

Feature rows and coordinate rows must remain aligned. Slide identifiers must match the metadata and fold files under wsi_dataset_annotation/.

Replace the /path/to/... value in data_dir and verify label_dir and csv_dir in the selected model and clustering configurations before running the code. Paper experiments use:

configs/models/<backbone>/uni_v1/<dataset>.yaml
configs/cluster/uni_v1/<dataset>.yaml
configs/prompts/musk/<dataset>.yaml

Valid backbone directories are abmil, clamsb, rrtmil, s4mil, transmil, and wikg. Dataset configuration names are camel, brca, esca, nsclc, and rcc.

Reproduction

The commands below use ABMIL on TCGA-BRCA fold 0. Use the corresponding model configuration, checkpoint, clusters, and fold for other experiments.

Train

Training and validation use the complete slide. TTIS, offline clusters, and MUSK are not used during training.

python train.py \
    --stage train \
    --config configs/models/abmil/uni_v1/brca.yaml \
    --fold 0 \
    --gpus 0

The best validation checkpoint is saved under logs/<model>/<encoder>/<dataset>/fold_<n>/checkpoints/.

Test

Select one of the following sampling modes. All modes use the same trained MIL checkpoint.

Full-slide baseline

Use --sampling none. Offline clusters and MUSK are not required.

python train.py \
    --stage test \
    --sampling none \
    --config configs/models/abmil/uni_v1/brca.yaml \
    --model_path /path/to/checkpoint.ckpt \
    --fold 0 \
    --gpus 0

Random sampling

Use --sampling random. Offline clusters and MUSK are not required.

python train.py \
    --stage test \
    --sampling random \
    --config configs/models/abmil/uni_v1/brca.yaml \
    --model_path /path/to/checkpoint.ckpt \
    --fold 0 \
    --gpus 0

Repeat the random-sampling evaluation following the protocol described in the paper.

TTIS

TTIS requires a trained checkpoint and offline clusters from the same dataset fold.

First, extract the clusters:

python extract_offline_clusters.py \
    --config_path configs/cluster/uni_v1/brca.yaml \
    --fold 0 \
    --gpus 0

Then, run TTIS inference using the generated cluster JSON:

python train.py \
    --stage test \
    --sampling ttis \
    --config configs/models/abmil/uni_v1/brca.yaml \
    --model_path /path/to/checkpoint.ckpt \
    --cluster_path cluster_results/uni_v1/fold_0/tcga_brca_h_kmeans_256.json \
    --fold 0 \
    --gpus 0

Server command examples are available in train.sh, test.sh, test_random.sh, extract_offline_clusters.sh, and test_ttis.sh.

Outputs

Checkpoints and CSV metrics are written under logs/<model>/<encoder>/<dataset>/fold_<n>/. Test results are separated by sampling method:

fold_<n>/test
fold_<n>/test_random
fold_<n>/test_ttis

The code reports accuracy, macro F1, macro precision, macro recall, and macro AUROC. HF_TOKEN, WANDB_API_KEY, and WANDB_USER_NAME are placeholders in src/constants.py. W&B remains disabled until WANDB_MODE is set to online.

Acknowledgements

This work was supported by a grant of the National Research Foundation of Korea (NRF) (No. RS-2025-00558322 and RS-2024-00397293).

Citation

If you find this work useful, please cite:

@inproceedings{nguyen2026ttis,
  title={Test-Time Instance Selection for Improved Whole Slide Image Analysis},
  author={Nguyen, Quoc Anh and Park, Sunhong and Kwak, Jin Tae},
  booktitle={2nd MICCAI Workshop on Efficient Medical AI},
  year={2026}
}