README.md

---
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}
}
```