AMAES: Masked Autoencoder Pretraining on 3D Brain MRI

Official Pytorch implementation of AMAES from the paper

AMAES: Augmented Masked Autoencoder Pretraining on Public Brain MRI Data for 3D-Native Segmentation
ADSMI @ MICCAI 2024
Asbjørn Munk*, Jakob Ambsdorf*, Sebastian Llambias, Mads Nielsen

Pioneer Centre for AI & University of Copenhagen

* Equal Contribution

Efficient pretraining for 3D segmentation models using MAE and augmentation reversal on a large domain-specific dataset.

For more information on the paper see amaes.asbn.dk.

[NEW!] The codebase has been rewritten for the FOMO25 challenge baseline codebase. Check it our here.

Overview

Method

🧠BRAINS-45K dataset

All models are pretrained on 🧠BRAINS-45K, the largest pretraining dataset available for brain MRI.

We have released a new dataset of 60K brain MRI callled FOMO60k, which is available for download on Hugging Face 🤗.

[New 2026!] We have just released a dataset of 319K Brain MRI, called FOMO300K, which is available for download on Hugging Face 🤗.

Model checkpoints

All checkpoints have been pretrained on 🧠BRAINS-45K for 100 epochs using AMAES.

Model	Parameters	Checkpoint
	M	Zenodo
U-Net XL	90	Download
U-Net B	22	Download
MedNeXt-L	55	Download
MedNeXt-M	21	Download

All models were pretrained on 2xH100 GPUs with 80GB of memory.

Running the code

1. Install Poetry.
2. Create environment by calling poetry install.

Setup data

AMAES is using the Yucca library for handling 3D medical data.

Guide on how to setup data comming soon.

Pretraining

To pretrain using AMAES run

poetry run src/pretrain.py --base_path=<path to base data directory>

Finetuning

To finetune using AMAES, run

poetry run src/train.py --base_path=<path to base data directory> --pretrained_weights_path="<path_to_checkpoint>" --model=<model_to_instantiate>

Note that the checkpoint must match the model provided. For instance, to finetune unet_xl_lw_dec_fullaug.pth use --model=unet_xl.

Citation

Please use

@article{munk2024amaes,
  title={AMAES: Augmented Masked Autoencoder Pretraining on Public Brain MRI Data for 3D-Native Segmentation},
  author={Munk, Asbjørn and Ambsdorf, Jakob and Llambias, Sebastian and Nielsen, Mads},
  journal={MICCAI Workshop on Advancing Data Solutions in Medical Imaging AI (ADSMI 2024), MICCAI 2024},
  year={2024}
}