This repository contains code, results and dataset links for our IEEE:SMC 2021 oral paper titled IMLE-Net: An Interpretable Multi-level Multi-channel Model for ECG Classification. 📝

Authors: Likith Reddy, Vivek Talwar, Shanmukh Alle, Raju. S. Bapi, U. Deva Priyakumar.

More details on the paper can be found here. Also available on arxiv.

Raise an issue for any query regarding the code, paper or for any support.

Tested and works with the latest tensorflow and torch versions at the time of this commit.

• Introduction
• Highlights
• Results
• Dataset
• Getting started
• Getting the weights
• License and Citation

Early detection of cardiovascular diseases is crucial for effective treatment and an electrocardiogram (ECG) is pivotal for diagnosis. The accuracy of Deep Learning based methods for ECG signal classification has progressed in recent years to reach cardiologist-level performance. In clinical settings, a cardiologist makes a diagnosis based on the standard 12-channel ECG recording. Automatic analysis of ECG recordings from a multiple-channel perspective has not been given enough attention, so it is essential to analyze an ECG recording from a multiple-channel perspective. We propose a model that leverages the multiple-channel information available in the standard 12-channel ECG recordings and learns patterns at the beat, rhythm, and channel level. The experimental results show that our model achieved a macro-averaged ROC-AUC score of 0.9216, mean accuracy of 88.85% and a maximum F1 score of 0.8057 on the PTB-XL dataset. The attention visualization results from the interpretable model are compared against the cardiologist’s guidelines to validate the correctness and usability.

• A model that learns patterns at the beat, rhythm, and channel level with high accuracy💯.
• An interpretable model that gives an explainability at the beat, rhythm and channel level💥.
• Complete preprocessing pipeline, training and inference codes are provided.
• Training weights are available to try out the model.

Performance metrics

Visualization of normalized attention scores with red having a higher attention score and yellow having a lower attention score for a 12-lead ECG signal.

Channel Importance scores for the same 12-lead ECG signal.

The PTB-XL dataset can be downloaded from the Physionet website.

• To prepare the dataset, cd IMLE-Net
• Download the dataset using the terminal wget -r -N -c -np -nH --cut-dirs 4 -O data/ptb.zip https://physionet.org/static/published-projects/ptb-xl/ptb-xl-a-large-publicly-available-electrocardiography-dataset-1.0.2.zip
• Unzip and rename the dataset, unzip data/ptb.zip -d data/ && mv data/ptb-xl-a-large-publicly-available-electrocardiography-dataset-1.0.2 data/ptb && rm data/ptb.zip

More details on the dataset can be found here.

• All the development work is done on Python 3.7
• Install necessary dependencies using requirements.txt file. Run pip install -r requirements.txt in terminal
• Alternatively, set up environment and train the model using Dockerfile. Run docker build -f Dockerfile -t <image_name> .

• train.py trains a particular model from scratch
• preprocessing contains the preprocessing scripts
• models contains scripts for each model
• utils contains utilities for dataloader, callbacks and metrics

• The models are implemented in either tensorflow or torch
• Models implemented in tensorflow are imle_net, mousavi and rajpurkar
• Models implemented in torch are ecgnet and resnet101
• To log the training and validation metrics using wandb tool, set --loggr to True
• To train a particular model from scratch, cd IMLE-Net
• To train a model on sub-diseases of MI, set sub to True
• To run tensorflow models, python train.py --model imle_net --batchsize 32 --epochs 60 --loggr False
• To run torch models, python torch_train.py --model ecgnet --batchsize 32 --epochs 60 --loggr False

• To test a model on test dataset after training the model, cd IMLE-Net
• For tensorflow models, python test.py --model imle_net --batchsize 32
• For torch models, python torch_test.py --model ecgnet --batchsize 32

• It's a three step process, first train the model on main dataset with sub set to False.
• Second, train the model on sub-diseases of MI with sub set to True.
• Third, for inference and visualization run, python inference.py --dir filepath

• The logs are saved in logs/ directory.
• The model checkpoints are saved in checkpoints/ directory.
• The visualizations are saved in results/ directory.

Download the weights for several models trained on the PTB-XL dataset.

The software is licensed under the Apache License 2.0. Please cite the following paper if you have used this code:

@INPROCEEDINGS{9658706,  
author={Reddy, Likith and Talwar, Vivek and Alle, Shanmukh and Bapi, Raju. S. and Priyakumar, U. Deva},  
booktitle={2021 IEEE International Conference on Systems, Man, and Cybernetics (SMC)},   
title={IMLE-Net: An Interpretable Multi-level Multi-channel Model for ECG Classification},   
year={2021},  
pages={1068-1074}, 
doi={10.1109/SMC52423.2021.9658706}}