---
license: cc-by-nc-sa-4.0
---

# UNI-based ABMIL models for metastasis detection

These are weakly-supervised, attention-based multiple instance learning models for binary metastasis detection (normal versus metastasis). The models were trained on the [CAMELYON16](https://camelyon16.grand-challenge.org/Data/) dataset using UNI embeddings.

If you find this model useful, please cite our corresponding [preprint](https://arxiv.org/abs/2409.03080):

```bibtex
@misc{kaczmarzyk2024explainableaicomputationalpathology,
      title={Explainable AI for computational pathology identifies model limitations and tissue biomarkers}, 
      author={Jakub R. Kaczmarzyk and Joel H. Saltz and Peter K. Koo},
      year={2024},
      eprint={2409.03080},
      archivePrefix={arXiv},
      primaryClass={q-bio.TO},
      url={https://arxiv.org/abs/2409.03080}, 
}
```

# Data

- Training set consisted of 243 whole slide images (WSIs).
    - 143 negative
    - 100 positive
        - 52 macrometastases
        - 48 micrometastases
- Validation set consisted of 27 WSIs.
    - 16 negative
    - 11 positive
        - 6 macrometastases
        - 5 micrometastases
- Test set consisted of 129 WSIs.
    - 80 negative
    - 49 positive
        - 22 macrometastases
        - 27 micrometastases

# Evaluation

Below are the classification results on the test set.

|   Seed |   Sensitivity |   Specificity |    BA |   Precision |    F1 |
|-------:|--------------:|--------------:|------:|------------:|------:|
|  0 |         0.959 |         1.000 | 0.980 |       1.000 | 0.979 |
|  1 |         0.959 |         0.988 | 0.973 |       0.979 | 0.969 |
|  2 |         1.000 |         1.000 | 1.000 |       1.000 | 1.000 |
|  3 |         0.980 |         0.950 | 0.965 |       0.923 | 0.950 |
|  4 |         0.980 |         1.000 | 0.990 |       1.000 | 0.990 |

# How to reuse the model

The model expects 128 x 128 micrometer patches, embedded with the UNI model.

```python
import torch
from abmil import AttentionMILModel

model = AttentionMILModel(in_features=1024, L=512, D=384, num_classes=2, gated_attention=True)
model.eval()
state_dict = torch.load("seed2/model_best.pt", map_location="cpu", weights_only=True)
model.load_state_dict(state_dict)

# Load a bag of features
bag = torch.ones(1000, 1024)
with torch.inference_mode():
    logits, attention = model(bag)
```

# How to train the model

Download the UNI embeddings for CAMELYON16 from https://huggingface.co/datasets/kaczmarj/camelyon16-uni and then, run the commands below.

```shell
# Seed 0
python train_classification.py --model-name AttentionMILModel --features-dir path/to/features/ --output-dir outputs/abmil-uni-128um_seed0 --csv data.csv --label-col binary_label_int --num-classes 2 --embedding-size 1024 --split-json splits.json --fold 0 --num-epochs 20 --seed 0 -L 512 -D 384 --lr 1e-4
# Seed 1
python train_classification.py --model-name AttentionMILModel --features-dir path/to/features/ --output-dir outputs/abmil-uni-128um_seed1 --csv data.csv --label-col binary_label_int --num-classes 2 --embedding-size 1024 --split-json splits.json --fold 0 --num-epochs 20 --seed 1 -L 512 -D 384 --lr 1e-4
# Seed 2
python train_classification.py --model-name AttentionMILModel --features-dir path/to/features/ --output-dir outputs/abmil-uni-128um_seed2 --csv data.csv --label-col binary_label_int --num-classes 2 --embedding-size 1024 --split-json splits.json --fold 0 --num-epochs 20 --seed 2 -L 512 -D 384 --lr 1e-4
# Seed 3
python train_classification.py --model-name AttentionMILModel --features-dir path/to/features/ --output-dir outputs/abmil-uni-128um_seed3 --csv data.csv --label-col binary_label_int --num-classes 2 --embedding-size 1024 --split-json splits.json --fold 0 --num-epochs 20 --seed 3 -L 512 -D 384 --lr 1e-4
# Seed 4
python train_classification.py --model-name AttentionMILModel --features-dir path/to/features/ --output-dir outputs/abmil-uni-128um_seed4 --csv data.csv --label-col binary_label_int --num-classes 2 --embedding-size 1024 --split-json splits.json --fold 0 --num-epochs 20 --seed 4 -L 512 -D 384 --lr 1e-4
```