Multimodal Pathology MIL
Weakly supervised multimodal pathology classification with HE patch features and protein-expression prediction features.
This repository implements a slide-level multiple instance learning (MIL) pipeline for pathology classification. Each patch already has two precomputed modalities:
HE visual feature : 1024-dProtein prediction feature : 40-d
The model aggregates patch-level information into a slide-level representation, predicts the main group, and then predicts subgroup conditionally within the predicted group.
What This Repository Provides
A hierarchical multimodal MIL training pipeline
Gated fusion between HE and protein modalities
Attention-based patch weighting for slide-level prediction
Export of high-attention patches for interpretation
Export of patch-level fused embeddings for downstream clustering
Support for standard inference and synthetic WSI grouping at test time
Model Summary
The current model combines several ideas:
Gated multimodal fusion
HE and protein features are projected into the same hidden space, then fused with a learned gate instead of naive concatenation.
Two-stage attention
The model can use pseudo-bags before whole-slide aggregation, which is helpful when a slide contains many patches.
Hierarchical classification
The model predicts group first, then routes the sample to a group-specific subgroup head.
Patch-level export
The final fused patch embedding is exported for train, validation, and inference outputs, which makes downstream clustering and visualization straightforward.
Repository Structure
1 .
2 ├── main.py
3 ├── train_hierarchical_mil.py
4 ├── train_data_loader.py
5 ├── inference_data_loader.py
6 ├── example_data/
7 ├── plot_confusion_matrix.py
8 ├── requirements.txt
9 ├── .gitignore
10 └── README.md
File Roles
main.py
Unified command-line entry point for training and test-time inference.
train_hierarchical_mil.py
Model definition, training loop, validation loop, checkpointing, metric logging, curve plotting, and train/val patch feature export.
train_data_loader.py
Training and validation data loader for aligned HE/protein patch features.
inference_data_loader.py
Inference data loader with optional synthetic WSI grouping.
plot_confusion_matrix.py
Plot a publication-style confusion matrix from grouped prediction statistics.
Data Assumptions
This repository assumes that patch-level features are already extracted .
Expected inputs:
HE feature CSV
Protein feature CSV
Training/validation data are expected to be aligned at the patch level and grouped into slides. Inference data can either be:
standard slide-organized CSVs, or
single-patch rows that are regrouped into synthetic WSI bags
This repository also includes synthetic example inputs under example_data/:
example_data/example_he_features.csvexample_data/example_protein_features.csv
These are format references only and use dummy values.
Installation
Create an environment and install the required packages:
1 conda create -n model python = 3.10 -y
2 conda activate model
3 pip install -r requirements.txt
Quick Start
1. Train
1 python main.py train \
2 --epochs 30 \
3 --batch-size 4 \
4 --experiment-name exp01
This will create a timestamped folder under experiments/.
2. Inference with main.py
1 python main.py test \
2 --checkpoint experiments/your_run/best_model.pth \
3 --he-path path/to/he_features.csv \
4 --prot-path path/to/protein_features.csv \
5 --batch-size 64 \
6 --output-name test_run
You can also test the interface quickly with the bundled example inputs:
1 python main.py test \
2 --checkpoint best_model.pth \
3 --he-path example_data/example_he_features.csv \
4 --prot-path example_data/example_protein_features.csv \
5 --batch-size 2 \
6 --output-name example_test
Synthetic WSI Grouping
Some test sets may effectively contain one patch per WSI . In that case, you can regroup patches from the same inferred class into synthetic WSI bags.
For example, to randomly group 48 patches from the same inferred group into one synthetic bag:
1 python main.py test \
2 --checkpoint experiments/your_run/best_model.pth \
3 --he-path data_test/test_he_processed.csv \
4 --prot-path data_test/test_data_processed.csv \
5 --synthetic-wsi-size 48 \
6 --synthetic-wsi-seed 42 \
7 --output-name synthetic48_test
You can keep the final smaller bag in each group with:
--keep-synthetic-remainderOutputs
Training Outputs
Each training run creates a new folder under experiments/ containing files such as:
best_model.pthlast_model.pthmetrics.csvconfig.jsontraining_curves.pngbest_val_top_patches.csvtrain_patch_features.csvval_patch_features.csv
Inference Outputs
Each inference run creates a new folder under inference_runs/ containing files such as:
test_predictions.csvtest_top_patches.csvtest_patch_features.csvtest_group_summary.csvtest_subgroup_summary.csvtest_group_accuracy_summary.csvtest_inferred_vs_pred_group.csvtest_inferred_vs_pred_subgroup.csv
Patch Feature Export
The repository exports fused patch-level embeddings from the model.
Feature columns are named:
feat_000, feat_001, ..., feat_127These exported embeddings can be used directly for:
clustering
UMAP / t-SNE visualization
prototype analysis
patch subtype exploration
Public Release Recommendations
If you plan to publish this project on GitHub or Hugging Face:
Put code and documentation in GitHub
Put model checkpoints in Hugging Face Model Hub
Optionally create a Hugging Face Space for an interactive demo
This repository already includes:
a cleaner top-level README
a minimal requirements.txt
a .gitignore that avoids uploading local artifacts, checkpoints, and private data by default
What Is Not Included
For a public release, you will usually not want to upload:
private training data
private test data
local experiment outputs
temporary logs
The provided .gitignore is set up with that assumption.
Citation
If you use this repository in academic work, consider adding a project-specific citation block here once the manuscript or preprint is available.
License
Please add your intended license before publishing, for example:
At the moment, no explicit license is declared in this repository.
