Hi-End-MAE: Hierarchical encoder-driven masked autoencoders are stronger vision learners for medical image segmentation

Fenghe Tang^1,2, Qingsong Yao³, Wenxin Ma^1,2, Chenxu Wu^1,2, Zihang Jiang^1,2, S.Kevin Zhou^1,2

¹ School of Biomedical Engineering, University of Science and Technology of China 
² Suzhou Institute for Advanced Research, University of Science and Technology of China 
³ Stanford University, Palo Alto, California, 94305, United State

​

News

[2024/02/14] Paper and code released !

Getting Started

Prepare Environment

conda create -n HiEndMAE python=3.9
conda activate HiEndMAE
pip install torch==1.13.0 torchvision==0.14.0 torchaudio==0.13.0
pip install packaging timm==0.5.4
pip install transformers==4.34.1 typed-argument-parser
pip install numpy==1.21.2 opencv-python==4.5.5.64 opencv-python-headless==4.5.5.64
pip install 'monai[all]'
pip install monai==1.2.0

Prepare Datasets

We recommend you to convert the dataset into the nnUNet format.

└── Hi-End-MAE
    ├── data
        ├── Dataset001_BTCV
            └── imagesTr
                ├── xxx_0000.nii.gz
                ├── ...
        ├── Dataset006_FLARE2022
            └── imagesTr
                ├── xxx_0000.nii.gz
                ├── ...
        └── Other_dataset
            └── imagesTr
                ├── xxx_0000.nii.gz
                ├── ...

Start Training

Run training on multi-GPU :

# An example of training on 4 GPUs with DDP
torchrun --nproc_per_node=4 --nnodes=1 --node_rank=0 --master_addr=localhost --master_port=12351 main.py

Fine-tuning

Load pre-training weights :

# An example of Fine-tuning on BTCV (num_classes=14)
from downstream.factory import load_hi_end_mae_10k

model = load_hi_end_mae_10k(n_classes=14)

Citation

If the code, paper and weights help your research, please cite:

License

This project is released under the Apache 2.0 license. Please see the LICENSE file for more information.