CineMA: A Foundation Model for Cine Cardiac MRI 🎥🫀

Overview

CineMA is a vision foundation model for Cine cardiac magnetic resonance (CMR) imaging based on Masked-Autoencoder. CineMA has been pre-trained on UK Biobank data and fine-tuned on multiple clinically relevant tasks such as ventricle and myocaridum segmentation, ejection fraction (EF) regression, cardiovascular disease (CVD) detection and classification, and mid-valve plane and apical landmark localization. The model has been evaluated on multiple datasets, including ACDC, M&Ms, M&Ms2, Kaggle, Rescan, and Landmark, etc.

Usage

Installation

You can install the package inside a Conda environment using following commands

git clone https://github.com/mathpluscode/CineMA.git
cd CineMA
conda env update -f envs/environment.yml
conda activate cinema
pip install -e .

Pytorch should be installed separately following the official instructions.

Use fine-tuned models

The fine-tuned models have been released at https://huggingface.co/mathpluscode/CineMA. Example inference scripts are available to test these models.

python examples/inference/segmentation_sax.py
python examples/inference/segmentation_lax_4c.py
python examples/inference/classification_cvd.py
python examples/inference/classification_sex.py
python examples/inference/classification_vendor.py
python examples/inference/regression_ef.py
python examples/inference/regression_bmi.py
python examples/inference/regression_age.py
python examples/inference/landmark_heatmap.py
python examples/inference/landmark_coordinate.py

Training Task	Input View	Input Timeframes	Inference Script
Ventricle and myocardium segmentation	SAX	1	segmentation_sax.py
Ventricle and myocardium segmentation	LAX 4C	1	segmentation_lax_4c.py
CVD classification	SAX or LAX 4C	2 (ED and ES)	classification_cvd.py
Patient sex classification	SAX	2 (ED and ES)	classification_sex.py
CMR machine vendor classification	SAX or LAX 4C	2 (ED and ES)	classification_vendor.py
EF regression	SAX or LAX 4C	2 (ED and ES)	regression_ef.py
Patient BMI regression	SAX	2 (ED and ES)	regression_bmi.py
Patient age regression	SAX	2 (ED and ES)	regression_age.py
Landmark localization by heatmap regression	LAX 2C or LAX 4C	1	landmark_heatmap.py
Landmark localization by coordinates regression	LAX 2C or LAX 4C	1	landmark_coordinate.py

Use pre-trained models

The pre-trained CineMA model backbone is available at https://huggingface.co/mathpluscode/CineMA. Following scripts demonstrated how to fine-tune this backbone using a preprocessed version of ACDC dataset:

python examples/train/segmentation.py
python examples/train/classification.py
python examples/train/regression.py

Task	Fine-tuning Script
Ventricle and myocardium segmentation	segmentation.py
Cardiovascular disease classification	classification.py
Ejection fraction regression	regression.py

Another two scripts demonstrated the masking and prediction process of MAE and the feature extraction from MAE.

python examples/inference/mae.py
python examples/inference/mae_feature_extraction.py

For fine-tuning CineMA on other datasets, pre-process can be performed using the provided scripts following the documentations. Note that it is recommended to download the data under ~/.cache/cinema_datasets as the integration tests uses this path. For instance, the mnms preprocessed data would be ~/.cache/cinema_datasets/mnms/processed. Otherwise define the path using environment variable CINEMA_DATA_DIR.

Training Data	Documentations
ACDC	README.md
M&Ms	README.md
M&Ms2	README.md
Kaggle	README.md
Rescan	README.md
Landmark	README.md
EMIDEC	README.md
Myops2020	README.md

The code for training and evaluating models on these datasets are available.

Task	Data	Documentation
Ventricle and myocardium segmentation	ACDC	cinema/segmentation/acdc/README.md
Ventricle and myocardium segmentation	M&Ms	cinema/segmentation/mnms/README.md
Ventricle and myocardium segmentation	M&Ms2	cinema/segmentation/mnms2/README.md
Ventricle and myocardium segmentation	Kaggle	cinema/segmentation/kaggle/README.md
Ventricle and myocardium segmentation	Rescan	cinema/segmentation/rescan/README.md
Scar segmentation	EMIDEC	cinema/segmentation/emidec/README.md
Scar segmentation	Myops2020	cinema/segmentation/myops2020/README.md
Image classification	ACDC	cinema/classification/acdc/README.md
Image classification	M&Ms	cinema/classification/mnms/README.md
Image classification	M&Ms2	cinema/classification/mnms2/README.md
Image regression	ACDC	cinema/regression/acdc/README.md
Image regression	M&Ms	cinema/regression/mnms/README.md
Image regression	M&Ms2	cinema/regression/mnms2/README.md
Landmark localization by heatmap regression	Landmark	cinema/segmentation/landmark/README.md
Landmark localization by coordinates regression	Landmark	cinema/regression/landmark/README.md

Train your own foundation model

A simplified example script for launch masked autoencoder pretraining is provided: pretrain.py. For DDP supported pretraining, check cinema/mae/pretrain.py. Check examples/dicom_to_nifti.py for UKB data preprocessing.

python examples/train/pretrain.py

References

The code is built upon several open-source projects:

• UK Biobank Cardiac Preprocessing
• Masked Autoencoders
• ConvMAE
• MultiMAE
• MONAI
• PyTorch Vision
• PyTorch Image Models

Acknowledgement

Contact

For any questions or suggestions, please create an issue.

For collaborations, please contact Yunguan Fu ([email protected]).

Citation