Fig.1: Different paradigms for brain lesion segmentation and an overview of the proposed framework:
a) the traditional paradigm;
b) the foundation model paradigm;
c) the proposed mixture of modality experts (MoME) framework for constructing the foundation model.
Fig.2:Detailed analysis of the MoME result on seen datasets. a) A radar chart that compares the average Dice score of foundation models from the perspectives of different modalities and lesion types. b) t-SNE plots of latent spaces for nnU-Net and MoME, where each dot represents a brain image.
Please also cite these paper [MICCAI link] and [TMI link] if you are using MoME for your research!
@inproceedings{zhang2024foundation,
title={A foundation model for brain lesion segmentation with mixture of modality experts},
author={Zhang, Xinru and Ou, Ni and Basaran, Berke Doga and Visentin, Marco and Qiao, Mengyun and Gu, Renyang and Ouyang, Cheng and Liu, Yaou and Matthews, Paul M and Ye, Chuyang and others},
booktitle={International Conference on Medical Image Computing and Computer-Assisted Intervention},
pages={379--389},
year={2024},
organization={Springer}
}
@ARTICLE{10879789,
author={Zhang, Xinru and Ou, Ni and Doga Basaran, Berke and Visentin, Marco and Qiao, Mengyun and Gu, Renyang and Matthews, Paul M. and Liu, Yaou and Ye, Chuyang and Bai, Wenjia},
journal={IEEE Transactions on Medical Imaging},
title={A Foundation Model for Lesion Segmentation on Brain MRI With Mixture of Modality Experts},
year={2025},
volume={44},
number={6},
pages={2594-2604},
keywords={Image segmentation;Lesions;Brain modeling;Imaging;Foundation models;Biomedical imaging;Magnetic resonance imaging;Training;Context modeling;Three-dimensional displays;Foundation model;mixture of experts;brain imaging;brain lesion segmentation},
doi={10.1109/TMI.2025.3540809}}
conda create --name MoME python=3.9.18
conda activate MoME
git clone https://github.com/ZhangxinruBIT/MoME.git
cd MoME/MoME_foundation
pip install -e .
Additionally, since we have made significant changes to the third-party library named dynamic_network_architectures, we did not take the original one.
Data Preprocessing
We perform affine registration, skull stripping, and brain imaging cropping, along with their respective annotations. But make sure you have downloaded the Advanced Normalization Tools (ANTs) before. The following command can be executed after you convert the dataset to the format required by nnU-NetV2.
cd MoME/Codes_prepro
python one-step.py -dataset_path -ss #If all the brain images in your dataset include the skull, please add the `-ss` flag. If none of the images include the skull, you can omit the flag. Do not mix brain images with and without skulls in the same directory, as their preprocessing steps are different.
Datasplit
Please set your own datasplit, for me we save the information in the datasplit.json. And then modify the do_split in nnUNetTrainer with the code like:
if self.fold == 'MoME':
tr_keys = DATA['train']['BraTS'] + DATA['train']['ATLAS'] +DATA['train']['OASIS']\
+DATA['train']['ISLES']+DATA['train']['WMH2017']+DATA['train']['MSSEG']
val_keys = DATA['val']['BraTS'] + DATA['val']['ATLAS'] +DATA['val']['OASIS']\
+DATA['val']['ISLES']+DATA['val']['WMH2017']+DATA['val']['MSSEG']
to better manage the data split.
For traning and inference, since we implement within the nnU-NetV2, the well-introduced usage can be follwed at nnU-Net. Please ensure that the dataset format meets the requirements expected by nnU-Net with our preprocessed data (Rename the directories to resemble imagesTr and labelsTr under Dataset_XXXX.).
Experiment planning and preprocessing
nnUNetv2_plan_and_preprocess -d DATASET_ID --verify_dataset_integrity
Training
Before utilizing the MoME training, make sure you have acquired pretrained modality experts trained with nnUNet using corresponding modality images. Our pretrained modality experts are available at Hugging Face Pretrained Experts.
You can adjust the path to align with the your own specialized modality nnUNets in nnUNetTrainer.initialize(from line 214 to 245). Alternatively, you can utilize our well-trained experts and configure the path accordingly.
nnUNet_train XXX 3d_fullres MoME
Inference
When performing inference, ensure you have the pretrained checkpoint list similar to Pretrained_Experts.
nnUNetv2_predict -i INPUT_FOLDER -o OUTPUT_FOLDER -d DATASET_NAME_OR_ID -f MoME -c 3d_fullres -chk checkpoint_best.pth
Additionally, the final MoME model has been released on Hugging Face at MoME_CHECKPOINT.
The MoME model has been extended to MoME+ to handle combined multiple modalities as input. To this end, we develop a novel dispatch network with a soft assignment strategy and prior constraints, and employ random dropout of modalities during training to mimic the heterogeneous input.
MoME+ and MoME can share the same Conda environment. However, each time you switch between MoME and MoME+, you need to navigate to the corresponding folder and run pip install -e . to properly link the package.
Assuming you have already set up the MoME virtual environment, follow these steps to activate it and configure MoME+:
conda activate MoME
cd MoME/MoME_plus
pip install -e .
This ensures that the correct package is linked within the environment. Always remember to repeat this process whenever switching between MoME and MoME+.
Data Preprocessing and Datasplit can follow the same instructions described in the MoME section.
Besides, we made some modifications to the original nnUNet requirements of the data structure. Specifically, for the BraTS2021 dataset which contains T1w, T1ce, T2w, and FLAIR modalities, we do not use 0000, 0001, 0002, 0003 to differentiate modalities for the same patient. Instead, we directly use the modality keywords. Additionally, the structures of imagesTr and imagesTs differ slightly.
For example, taking BraTS2021_00000 as training case and BraTS2021_00002 as test case, all modality filenames for the images and labels should follow this format:
├── imagesTr
│ ├── BraTS2021_00000_t1_0000.nii.gz
│ ├── BraTS2021_00000_t1ce_0000.nii.gz
│ ├── BraTS2021_00000_t2_0000.nii.gz
│ ├── BraTS2021_00000_flair_0000.nii.gz
└── labelsTr
│ ├── BraTS2021_00000_t1.nii.gz
│ ├── BraTS2021_00000_t1ce.nii.gz
│ ├── BraTS2021_00000_t2.nii.gz
│ ├── BraTS2021_00000_flair.nii.gz
└── imagesTs
│ ├── BraTSt1
└── BraTS2021_00002_t1_0000.nii.gz
│ ├── BraTSt1ce
└── BraTS2021_00002_t1ce_0000.nii.gz
│ ├── BraTSt2
└── BraTS2021_00002_t2_0000.nii.gz
│ ├── BraTSflair
└── BraTS2021_00002_flair_0000.nii.gz
Ensure that all files are named in this format to align with the modified preprocessing, training, and inference pipelines.
Experiment planning and preprocessing
nnUNetv2_plan_and_preprocess -d DATASET_ID --verify_dataset_integrity
Training
Before running MoME+ training, ensure that you have pretrained modality experts with nnU-Net using the corresponding modality-specific images. In our paper, MoME+ is evaluated on the BraTS dataset for a fair comparison, so each modality expert is pretrained exclusively on BraTS data within the nnUNet framework. You can find our pretrained expert models here.
nnUNet_train XXX 3d_fullres MoME_plus
We employ random dropout of modalities during training to mimic the heterogeneous input. So that during the inference the input data can be heterogeneous.
Make sure to pretrain the modality expert model first, and then update the path in nnUNetTrainer.py (line 215、221、229、236).
Inference
For assessing the extended multi-input MoME+ model, we used the BraTS2021 dataset, which contains T1w, T1ce, T2w, and FLAIR modalities, and simulated different scenarios of missing modalities. Specifically, for the four modalities, each can either be available or unavailable, with the constraint that at least one modality must be present. This results in a total of
nnUNetv2_predict -i INPUT_FOLDER/BraTSt1 -o OUTPUT_FOLDER -d DATASET_NAME_OR_ID -f MoME_plus -c 3d_fullres -chk checkpoint_best.pth --MultiMod 1 X X X
nnUNetv2_predict -i INPUT_FOLDER/BraTSt1ce -o OUTPUT_FOLDER -d DATASET_NAME_OR_ID -f MoME_plus -c 3d_fullres -chk checkpoint_best.pth --MultiMod 0 1 X X
nnUNetv2_predict -i INPUT_FOLDER/BraTSt2 -o OUTPUT_FOLDER -d DATASET_NAME_OR_ID -f MoME_plus -c 3d_fullres -chk checkpoint_best.pth --MultiMod 0 0 1 X
nnUNetv2_predict -i INPUT_FOLDER/BraTSflair -o OUTPUT_FOLDER -d DATASET_NAME_OR_ID -f MoME_plus -c 3d_fullres -chk checkpoint_best.pth --MultiMod 0 0 0 1
This structure ensures that the correct modality is selected based on the --MultiMod parameter.
The final MoME+ model has been released on Hugging Face at MoME_plus_CHECKPOINT.