README.md

metadata

title: BrainAnytime Demo
emoji: 🧠
colorFrom: purple
colorTo: yellow
sdk: gradio
sdk_version: 6.14.0
python_version: '3.10'
app_file: app.py
pinned: false
license: mit
arxiv: 2605.13059
models:
  - Simmonstt/BrainAnytime

BrainAnytime

Official implementation of BrainAnytime: Anatomy-Aware Cross-Modal Pretraining for Brain Image Analysis with Arbitrary Modality Availability.

Paper

BrainAnytime: Anatomy-Aware Cross-Modal Pretraining for Brain Image Analysis with Arbitrary Modality Availability

• arXiv: 2605.13059
• PDF: https://arxiv.org/pdf/2605.13059

Congrats: This paper has been early accepted (top 9%) by MICCAI 2026.

Pretrained Weights

Finetuned checkpoints are available on Hugging Face: Simmonstt/BrainAnytime.

Pretrained weights are also available at Google Drive.

Overview

BrainAnytime is a self-supervised pretraining framework for multi-modal 3D brain imaging (T1, T2, Flair, PET) that handles arbitrary missing modality combinations at both training and inference time.

Key Features

• Multi-modal Masked Autoencoder (MultiMAE3D): Shared ViT encoder with per-modality input/output adapters, supporting 4 modalities (T1, T2, Flair, PET)
• Cross-Modal Mutual Prediction: EMA teacher-student framework for MRI-PET cross-level feature alignment
• Anatomy-Aware Adaptive Masking: Three-phase curriculum masking guided by AAL116 brain atlas and AD-relevant region priors
• Missing Modality Robustness: Handles arbitrary missing modality combinations via attention masking and observed indicators

Project Structure

BrainAnytime/
├── models/
│   ├── multimae3d.py              # MultiMAE3D model architecture
│   └── multimae3d_utils.py        # Patchify, masking, positional embeddings
├── anatomy_masking.py             # Anatomy-aware adaptive masking module
├── pretrain_dataloader_v2.py      # Multi-modal pretraining data loader
├── train_multimae.py              # Pretraining script (single/multi-GPU DDP)
├── finetune_main.py               # Downstream finetuning (CN vs AD, CN vs MCI, MMSE, AGE)
├── test_main.py                   # Test-only evaluation
└── altas/
    └── AAL116_standard.nii.gz     # AAL116 brain atlas (128x128x128)

Requirements

• Python >= 3.8
• PyTorch >= 1.12
• torchio
• nibabel
• timm
• einops
• tensorboardX
• scikit-learn
• pandas
• scipy
• tqdm

Data Preparation

Organize your data as follows:

./data/
├── Match_data_path/
│   └── pretraining_processed/     # Pretraining Excel files
│       ├── modality_data_A4.xlsx
│       ├── modality_data_ADNIDOD.xlsx
│       ├── modality_data_AIBL.xlsx
│       ├── modality_data_BraTS.xlsx
│       └── modality_data_NACC.xlsx
├── Pretrain/                      # Preprocessed NIfTI files for pretraining
└── Downstream/
    └── ADNI/                      # Downstream task data
        └── ADNI_Division/
            ├── modality_data_train.xlsx
            ├── modality_data_val.xlsx
            └── modality_data_test.xlsx

Each Excel file should contain columns for subject IDs and file paths to the corresponding NIfTI images for each modality.

Attention

To save training time, we preprocess the 3D multimodal image by following skull stripping, MN152 template co-registration, min–max normalization, and resampling to 128 × 128 ×128 in advance. The data loader only contains data augmentation during training.

Usage

Pretraining

# Single GPU
python train_multimae.py --batch_size 4

# Multi-GPU DDP (8 GPUs)
torchrun --nproc_per_node=8 train_multimae.py \
    --batch_size 16 \
    --enable_cross_modal \
    --use_anatomy_masking \
    --atlas_path altas/AAL116_standard.nii.gz

Downstream Finetuning

# Finetune on all tasks (3 seeds each)
python finetune_main.py \
    --pretrained ./pretrain_checkpoints/multimae/best_model.pth

# Specific task only
python finetune_main.py \
    --pretrained ./pretrain_checkpoints/multimae/best_model.pth \
    --tasks "CN vs AD"

Testing

# Test all tasks for finetune mode
python test_main.py --mode finetune

# Test a specific task
python test_main.py --mode finetune --tasks "CN vs AD"

Downstream Tasks

Task	Type	Metric
CN vs AD	Classification	ACC, AUC, Sensitivity, Specificity, F1
CN vs MCI	Classification	ACC, AUC, Sensitivity, Specificity, F1
MMSE	Regression	MAE, RMSE, Pearson
AGE	Regression	MAE, RMSE, Pearson

License

This project is released for academic research purposes only.

Citation

If you use BrainAnytime in your research, please cite:

@misc{yang2026brainanytimeanatomyawarecrossmodalpretraining,
      title={BrainAnytime: Anatomy-Aware Cross-Modal Pretraining for Brain Image Analysis with Arbitrary Modality Availability},
      author={Guangqian Yang and Tong Ding and Wenlong Hou and Yue Xun and Ye Du and Qian Niu and Shujun Wang},
      year={2026},
      eprint={2605.13059},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2605.13059},
}

Paper page: https://arxiv.org/abs/2605.13059