haozhongma/WSI_Quantify

WSI Quantify

Eight-class histopathology tile classification and whole-slide image (WSI) quantification for H&E-stained slides. This repository brings together the ViT-B/16 checkpoint and the project code for training, WSI inference, and downstream lymphocyte–stroma ratio (LSR) survival analyses.

• Code upstream: HaozhongMa/WSI_Quantify.
• Code snapshot: db1ccf5, copied on 8 October 2026.
• Weights: 8class_final_vit.pth, 343,284,050 bytes (327.4 MiB).

Repository contents

File	Purpose
8class_final_vit.pth	Torchvision ViT-B/16 state dictionary with an eight-class head
8class-WSI-quantify.py	WSI tile inference, probability fusion, tissue visualization, and count/area metrics
8class-model-train.py	Supervised tile classifier training and checkpoint evaluation
8class-model-train-for-lm.py	Frozen CONCH, TITAN, and UNI encoders with linear or MLP probes
survival_ml_lrepeated_splits.py	Repeated-split survival analysis comparing clinical features with and without LSR
environments.yaml	Upstream Python 3.10 / PyTorch 2.5.1 / CUDA 12.1 Conda environment
requirements.txt	Upstream pinned Python dependencies for the CUDA 12.1 environment

Training images, patient-level clinical tables, and additional encoder/probe checkpoints must be supplied separately.

Download and installation

Download the code and weights with the Hugging Face CLI:

python -m pip install huggingface_hub
hf download haozhongma/WSI_Quantify --local-dir WSI_Quantify
cd WSI_Quantify

For the supplied CUDA environment, use an NVIDIA GPU on Linux:

conda env create -f environments.yaml
conda activate wsi-quantify-pytorch121

Alternatively, in a Python 3.10 environment:

python -m pip install -r requirements.txt

The supplied dependency files select CUDA 12.1 builds of PyTorch and are not portable CPU/macOS environment specifications. For a minimal CPU inference environment, install the native PyTorch builds and the WSI dependencies instead:

python -m pip install torch==2.5.1 torchvision==0.20.1 \
  numpy scipy pillow tqdm openslide-python openslide-bin huggingface_hub

OpenSlide requires its native library as well as the Python binding. environments.yaml includes openslide, and requirements.txt includes openslide-bin. If the binary wheel is unavailable on your platform, install OpenSlide through your system package manager; on macOS, use brew install openslide. The WSI script uses CPU when CUDA is unavailable.

If you already have the code, download only the checkpoint:

hf download haozhongma/WSI_Quantify 8class_final_vit.pth --local-dir .

Model and class order

The checkpoint is a PyTorch state dictionary, loaded with torchvision.models.vit_b_16(weights=None) and an eight-output classification head. It is not a Transformers AutoModel checkpoint.

The WSI classifier uses the following output order, matching the expected alphabetical ImageFolder class directories:

Index	Directory	Tissue class
0	ADI	Adipose
1	BAC	Background
2	DEB	Debris
3	LYM	Lymphocytes
4	MUS	Muscle
5	NOR	Normal
6	STR	Stroma
7	TUM	Tumour

Use RGB tiles of 224 × 224 pixels. WSI inference reads level-0 tiles, converts values to [0, 1], and applies ImageNet normalization: mean [0.485, 0.456, 0.406], standard deviation [0.229, 0.224, 0.225]. It does not resample slides to a fixed physical resolution. Slide MPP therefore matters when comparing slides from different scanners.

Classify one tile

import torch
from PIL import Image
from torchvision import models, transforms

classes = ["ADI", "BAC", "DEB", "LYM", "MUS", "NOR", "STR", "TUM"]
model = models.vit_b_16(weights=None)
model.heads.head = torch.nn.Linear(model.heads.head.in_features, len(classes))
state_dict = torch.load("8class_final_vit.pth", map_location="cpu", weights_only=True)
model.load_state_dict(state_dict)
model.eval()

# Supply an RGB tile already cropped to 224 x 224 pixels.
tile = Image.open("tile.png").convert("RGB")
preprocess = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
])
with torch.inference_mode():
    probabilities = model(preprocess(tile).unsqueeze(0)).softmax(dim=1)[0]
print(dict(zip(classes, probabilities.tolist())))

Quantify a whole slide

Supply an OpenSlide-readable WSI, such as slides/example.svs:

python 8class-WSI-quantify.py \
  --path slides/example.svs \
  --out outputs/wsi \
  --model_type vit_b_16 \
  --model 8class_final_vit.pth \
  --gpu_id 0 \
  --batch_size 64 \
  --tile_size 224 \
  --step 112

--model is explicit here because the upstream script otherwise looks for checkpoints/8class_vit_b_16.pth. Reduce --batch_size if GPU memory is limited. Use --num_workers 0 for single-process slide reading, and --no_save_vis when only the metrics CSV is needed.

The script first selects tiles using a low-resolution tissue mask, predicts class probabilities, and averages overlapping tile predictions on a stride grid. It then identifies tumour-related lymphocytes by proximity to tumour regions. These receive derived label 8 (Tumor_Relate_Lymphocytes); label 8 is a spatial post-processing result, not a ninth model output.

Important spatial settings:

Option	Default	Meaning
--tile_size	224	Level-0 tile width and height; keep 224 for this checkpoint
--step	Half the tile size	Stride in level-0 pixels; 112 with the default tile size
--mpp	Slide metadata	Explicit microns-per-pixel override
--default_mpp	0.25	Fallback when slide metadata has no MPP
--trl_radius_um	500	Distance used to assign tumour-related lymphocytes
--min_tumor_component_tiles	20	Minimum connected tumour component size in grid cells
--tumor_closing_radius	1	Morphological closing radius in grid steps

For slides without MPP metadata, pass the known value with --mpp so physical areas and the lymphocyte proximity radius use the correct scale.

Outputs and LSR definitions

For each slide, the script writes:

• <slide>_counts_<nine-count-values>.jpg: tissue label visualization, unless --no_save_vis is set.
• <slide>_counts_<nine-count-values>_metrics.csv: one row of tissue counts, areas, and LSR metrics; always written.

Counts refer to stride-grid cells, not individual cells or overlapping inference tiles. The estimated area of each grid cell is ((step * mpp) / 1000) ** 2 mm². The CSV includes WSI_name, model_type, total_count, grid_cell_area_mm2, per-class *_count and *_area_mm2, LSR, and NEW_LSR.

Let L be the count of remaining Lymphocytes, TRL the count of Tumor_Relate_Lymphocytes, and S the Stroma count after post-processing:

LSR     = (L + TRL) / (S + L + TRL)
NEW_LSR = TRL / (S + TRL)

The implementation returns 0.0 when a denominator is zero. These are lymphocyte fractions within the respective lymphocyte–stroma compartments; they are not lymphocyte-to-tumour ratios.

Train or evaluate tile classifiers

Prepare an ImageFolder dataset with the same eight class directories:

data/8class_patch/
  train/
    ADI/  BAC/  DEB/  LYM/  MUS/  NOR/  STR/  TUM/
  test/
    ADI/  BAC/  DEB/  LYM/  MUS/  NOR/  STR/  TUM/

Train a ViT-B/16 classifier:

python 8class-model-train.py \
  --mode train \
  --data-root data/8class_patch/train \
  --output outputs/model_train \
  --model-arch vit_b_16 \
  --gpu-ids 0

Evaluate the supplied checkpoint on your labelled test tiles:

python 8class-model-train.py \
  --mode test \
  --data-root data/8class_patch/test \
  --output outputs/model_eval \
  --model-arch vit_b_16 \
  --model-path 8class_final_vit.pth \
  --gpu-ids 0

Use --gpu-ids cpu to force CPU in the supervised training/evaluation script. It also supports resnet50, efficientnet_b0, efficientnet_b3, and swin_t; their checkpoints are not included here.

Frozen encoder probes

CONCH, TITAN, and UNI workflows need their own encoder weights. Place these in local model directories and train the classifier head on your tile dataset:

python 8class-model-train-for-lm.py \
  --encoder conch \
  --train-data-root data/8class_patch/train \
  --test-data-root data/8class_patch/test \
  --conch-model-dir models/CONCH \
  --output-dir outputs/conch_probe \
  --device cuda:0

--head-type linear selects a linear head instead of the default MLP. --encoder all runs all three encoders sequentially; supply --titan-model-dir models/TITAN and --uni-model-dir models/UNI2-h as well. The WSI script supports ViT-B/16, ResNet50, and a CONCH MLP head.

For WSI inference with a separately trained CONCH MLP probe:

python 8class-WSI-quantify.py \
  --path slides/example.svs \
  --out outputs/conch_wsi \
  --model_type conch \
  --model checkpoints/conch_mlp_probe.pt \
  --conch_model_dir models/CONCH \
  --gpu_id 0

Downstream survival analysis

Install the additional survival-analysis packages, which are not listed in the upstream dependency files:

python -m pip install scikit-survival xgboost

Provide train.csv and validation.csv in --data-root, or under its LSR/ subdirectory. Default columns are FileName_HE, OS (months), OSstatus (1 = event, 0 = censored), Age, Gender, Tumor Size, Stage, Tumor Site, Histologic Grade, and LSR. Merge the slide-derived feature with your patient-level clinical data before running this analysis; the WSI metrics CSV alone does not contain these clinical columns.

python survival_ml_lrepeated_splits.py \
  --data-root data/clinical/LSR \
  --output-dir outputs/survival_ml \
  --time-points 12 36 60 \
  --seeds 1 2 3 4 5

The script pools the two input CSVs and uses repeated random holdout splits (default test fraction 0.2). It compares NO_LSR with WITH_LSR using CoxPH, gradient boosting survival, random survival forest, survival SVM, and XGBoost with a Cox objective. It reports C-index and time-dependent AUC; IBS is computed for models that provide survival functions (CoxPH, GBS, RSF).

Outputs are survival_ml_metrics.csv, survival_ml_delta.csv, cindex_ibs_delta_summary.csv, and summary.json. Use --lsr-col NEW_LSR if your patient-level table contains that feature and you intend to analyse it.

Path configuration

Use explicit command-line paths as shown above, or configure the supported environment variables: WSI_DATA_ROOT, WSI_OUTPUT_DIR, WSI_TRAIN_ROOT, WSI_TEST_ROOT, RESNET50_MODEL_PATH, VIT_B_16_MODEL_PATH, CONCH_MODEL_DIR, CONCH_HEAD_PATH, TITAN_MODEL_DIR, and UNI_MODEL_DIR.

For example, point the WSI script's default ViT path at the downloaded weight:

export VIT_B_16_MODEL_PATH="$PWD/8class_final_vit.pth"

Use each script's --help to inspect its complete set of options.