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
• Output: CLS embedding and patch embeddings
• Training objective: global and local cosine feature distillation from high-magnification teacher features

Usage

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

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 the related XMag project and the original teacher/student foundation models where appropriate.