File size: 2,828 Bytes
3f09d41 c849cf3 3f09d41 c849cf3 60a84fd c849cf3 0fcabb3 c849cf3 0fcabb3 c849cf3 0fcabb3 c849cf3 0fcabb3 c849cf3 25ae462 97a430c 25ae462 97a430c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | ---
library_name: transformers
pipeline_tag: image-feature-extraction
tags:
- pathology
- histopathology
- foundation-model
- distillation
- dinov2
- xmag
---
# XMag
XMag is a pathology image encoder distilled from a high-magnification foundation model to a low-magnification student model.
- Teacher: frozen UNIv2
- Student: DINOv2 ViT-B/14
- Input: RGB pathology patch, `224 x 224`@5x or 2um mpp
- Output: CLS embedding and patch embeddings
- Training objective: global and local cosine feature distillation from high-magnification teacher features
## Usage
```python
import torch
from PIL import Image
from torchvision import transforms
from transformers import AutoModel
eval_transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
])
model = AutoModel.from_pretrained("AI4PATH/XMAG", trust_remote_code=True)
model.eval()
image = Image.open("patch.png").convert("RGB")
pixel_values = eval_transform(image).unsqueeze(0)
with torch.no_grad():
outputs = model(pixel_values)
cls_embedding = outputs["cls_embedding"] # (1, 768)
patch_embeddings = outputs["patch_embeddings"] # (1, 256, 768)
```
`pixel_values` should be normalized RGB tensors with shape `(B, 3, 224, 224)`.
The model does not apply preprocessing internally. Resize, `ToTensor()`, and ImageNet mean/std normalization should be done before calling the model.
## Model Details
The student sees a low-magnification `224 x 224` view of the tissue region. During training, the frozen UNIv2 teacher sees the corresponding high-magnification `896 x 896` region split into `4 x 4` subpatches. The student is trained to match both:
- A global teacher representation, computed by averaging the 16 teacher local features.
- Local teacher representations, aligned to pooled student patch-token blocks.
This release contains the student EMA backbone weights only. Projection heads used during distillation are not included.
## Requirements
```bash
pip install torch torchvision transformers huggingface_hub
```
The model code uses `torch.hub` to instantiate the DINOv2 ViT-B/14 backbone. The first load may need internet access to fetch the DINOv2 hub code, unless it is already cached.
## Citation
If you use this model, please cite:
Su, Z., Akbar, A. R., & Niazi, M. K. K. (2025). Streamline pathology foundation model by cross-magnification distillation. arXiv preprint arXiv:2509.23097. Available at arXiv:2509.23097.
```
@article{su2025streamline,
title={Streamline pathology foundation model by cross-magnification distillation},
author={Su, Ziyu and Akbar, Abdul Rehman and Sajjad, Usama and Parwani, Anil V and Niazi, Muhammad Khalid Khan},
journal={arXiv preprint arXiv:2509.23097},
year={2025}
}
```
|