Zero-Shot ECG Classification with Multimodal Learning and Test-time Clinical Knowledge Enhancement, ICML 2024.

To clone this repository:

git clone https://github.com/cheliu-computation/MERL.git

Datasets we used are as follows:

• MIMIC-IV-ECG: We downloaded the MIMIC-IV-ECG dataset as the ECG signals and paired ECG reports.

• PTB-XL: We downloaded the PTB-XL dataset which consisting four subsets, Superclass, Subclass, Form, Rhythm.

• CPSC2018: We downloaded the CPSC2018 dataset which consisting three training sets.

• CSN(Chapman-Shaoxing-Ningbo): We downloaded the CSN dataset.

We preprocessed pretraining datasets and split the dataset into train/val set using the code in pretrain/preprocess.ipynb.
We preprocessed downstream datasets and split the dataset into train/val/test set using the code in finetune/preprocess.ipynb.
We also provide the train/val/test split csv file in finetune/data_split

We pre-trained MERL on MIMIC-IV-ECG using this command:

bash MERL/pretrain/launch.sh

Pre-trained models can be found here.
We uploaded the pretrained models with resenet and vit.
xxx_ckpt.pth is the whole pretrained model for zeroshot classification.
xxx_encoder.pth is the ecg encoder only for linear probing.

We evlauate the performance of MERL on three scenarios: zero-shot classification, linear probing, and domain transferring.

We evaluate linear classification performance of our model using this command:

cd MERL/zeroshot
bash zeroshot.sh

We also release the CKEPE prompt in zeroshot/CKEPE_prompt.json.
Due to the copyright, we are unable to release the original SCP-code database, but you can find all information in: https://www.iso.org/standard/84664.html.

We provide bash script for evaluating linear probing performance of MERL:

cd MERL/finetune/sub_script
bash run_all_linear.sh

You can use --dataset to set specific dataset for finetuning. Here, 3 datsets are available: chexpert, rsna and covidx. You can use --ratio to set the fraction of training data for finetuning.

For domain trasnfering scenario, you do not reimplement any new experiments. You can only compute the metric across the overlapped categories.

If you found our work useful in your research, please consider citing our works(s) at:

@inproceedings{liuzero,
  title={Zero-Shot ECG Classification with Multimodal Learning and Test-time Clinical Knowledge Enhancement},
  author={Liu, Che and Wan, Zhongwei and Ouyang, Cheng and Shah, Anand and Bai, Wenjia and Arcucci, Rossella},
  booktitle={Forty-first International Conference on Machine Learning}
}