CXformer: Scalable Vision Foundation Models for Chest X-rays

CXformer is a vision transformer tailored for chest X-ray analysis, adapted from DINOv2 with clinically motivated training modifications. This repository provides code for pretraining CXformer using our optimized pipeline, as well as scripts for finetuning on downstream tasks like classification, segmentation, and report generation.

Key highlights:

• Improved training with register tokens, teacher centering, and optimized attention heads.
• Self-supervised pretraining on 600K+ CXRs from 5 global datasets.
• Strong generalization across 3 core tasks: classification, segmentation, and report generation.
• CXformer(S) matches the performance of RAD-DINO with 7× less training compute (in FLOPS).
• Models are available on Hugging Face 🤗: CXformer(B), CXformer(S)

📄 Paper: Empirical Analysis of Scaling Vision Foundation Models for Chest X-rays (MIDL 2025)
👨‍⚕️ Authors: Ahmed Al-Mahrooqi, Prateek Munjal, Ronnie Rajan, Marco AF Pimentel, Praveenkumar Kanithi
📍 Affiliation: M42, Abu Dhabi
📦 Models: CXformer(S), CXformer(B)
🧠 Base Architecture: Vision Transformers (ViT-S, ViT-B)
📊 Tasks: Image Classification, Semantic Segmentation, Report Generation

Model Checkpoints

📌 Note: Both models are trained solely on image data—no text supervision.

Datasets

Pretraining

• CheXpert
• MIMIC-CXR
• PadChest
• NIH-CXR8
• BRAX

Downstream Tasks

Getting Started

Installation

git clone https://github.com/m42-health/CXformer.git
cd CXformer
pip install -r requirements.txt

Continual Pretraining from DINOv2

Setting Up Pretrained Weights for DINOv2

Download the DINOv2 pretrained weights from https://github.com/facebookresearch/dinov2 and update the path in your config file.

For example, if you're using the following script which uses "dinov2/configs/pretrain/cxformer_small.yaml" as config file:

Make sure to set the pretrained_weights path under both the student and teacher sections like this:

student:
  pretrained_weights: /path/to/dinov2_checkpoint.pth

teacher:
  pretrained_weights: /path/to/dinov2_checkpoint.pth

sh scripts/pretrain/cxformer_slurm_submit.sh

export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:256 # to reduce fragmentation

n_nodes=1

cfg_file=dinov2/configs/pretrain/cxformer_small.yaml

PYTHONPATH=. python dinov2/run/train/train.py \
--nodes $n_nodes \
--nodelist "worker-13" \
--config-file $cfg_file \
--output-dir output_ablations_new/pretrain/cxformer_small_slurm/

Fine-Tuning on Downstream Tasks

You can fine-tune CXformer on downstream image classification tasks such as CheXpert using the provided shell script.

Run Image classification finetuning (CheXpert Dataset)

sh scripts/finetuning/image_classification/ft_cxformer_chexpert.sh

export CUDA_VISIBLE_DEVICES=0
n_epochs=100
pretrained_wt="m42-health/CXformer-small"

PYTHONPATH=. deepspeed dinov2/train/cxr_finetune.py \
  --config-file dinov2/configs/downstream/classification/cxformer_chexpert_small.yaml \
  --output-dir output_ablations_new/finetune/cxformer_chexpert \
  --exp-name ft_cxformer \
  --pretrained-weights $pretrained_wt \
  --model-type dinov2 \
  --num-epochs $n_epochs \
  --batch-size 10 \
  --num_workers 1 \
  --seed 7479 \
  --cls-n-layers 4 \
  --apply-avgpool \
  --clf_lr 5e-5 \
  --backbone_lr 5e-7

output_ablations_new/finetune/cxformer_chexpert/

Run Image Segmentation finetuning (Mimic ChexMask Dataset)

sh scripts/finetuning/segmentation/cxformer_chexmask.sh

For radiology report generation

We refer interested readers to llava repo as we straightaway used it in our work.

Results Summary

Classification (AUROC)

Segmentation (Dice Score)

Report Generation (MIMIC-CXR)

Repo Structure

CXformer/
├── configs/                              # YAML configuration files
├── models/                               # Vision backbone encoders and heads
├── dinov2/
│   ├── data/                             # Data loading utilities
│   ├── cxr_data/                         # Custom dataset classes and preprocessing logic for chest X-rays
│   └── train/
│       ├── cxr_pretrain.py              # Pretraining script
│       ├── cxr_finetune.py              # Finetuning script for image classification
│       └── cxr_segmentation.py          # Finetuning script for image segmentation
└── README.md                            # Project overview

License

This project is licensed under CC BY-NC-4.0 - see the LICENSE.md file for details.

Contribution

Contributions are welcome! Please feel free to submit a Pull Request.

Citation

@inproceedings{al2025empirical,
  title={Empirical Analysis of Scaling Vision Foundation Models for Chest X-rays},
  author={Al Mahrooqi, Ahmed and Munjal, Prateek and Rajan, Ronnie and Pimentel, Marco AF and Kanithi, Praveenkumar},
  booktitle={Medical Imaging with Deep Learning},
  year={2025}
}