IMLE-Net: An Interpretable Multi-level Multi-channel Model for ECG Classification

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.

Table of contents

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

Introduction 🔥

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.

Highlights ✨

• 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.

Results 🕺

Performance metrics

	Macro ROC-AUC	Mean Accuracy	Max. F1-score
Resnet101	0.8952	86.78	0.7558
Mousavi et al.	0.8654	84.19	0.7315
ECGNet	0.9101	87.35	0.7712
Rajpurkar et al.	0.9155	87.91	0.7895
IMLE-Net	0.9216	88.85	0.8057

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.

Dataset ⚡

Download ⌛

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

Getting started 🥷

• 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.

Getting started 🥷

Setting up the environment

• 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> .

What each file does

• 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

Training the model

• 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

Testing the model

• 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

Reimplementation and visualization

• 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

Logs and checkpoints

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

Getting the weights 🏋️

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

Name	Author's	Community
Mousavi et al.	link	link
ECGNet	link	link
Rajpurkar et al.	link	link
IMLE-Net	link	link

License and Citation 📰

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}}