ECG Arrhythmia classification

The repository contains code for Master's degree dissertation - Diagnosis of Diseases by ECG Using Convolutional Neural Networks. Only CNN neural network models are considered in the paper and the repository. As a part of the work, more than 30 experiments have been run. The table with all experiments and their metrics is available by the link

The best 1D and 2D CNN models are presented in the repository
The repository follows config principle and can be run in the following modes:

• Training - use train.py --config configs/training/<config>.json to train the model
• Validation - use inference.py --config configs/inference/config.json to validate the model
• Pipeline - use pipeline.py --config configs/pipelines/config/json to test the model using ECG data (i.e. data generation, running, visualization the results)

All available models and all necessary information are described below

Python 3.7 and PyTorch are used in the project
GitHub actions are used for installing dependencies and training implemented models

Program - Data Mining
Department - Computer Science

Principal Investigator - Nikolai Yu. Zolotykh
National Research University - Higher School of Economics

Implemented models

1D models:

• Cardiologist-Level Arrhythmia Detection with Convolutional Neural Networks
• ECG Heartbeat Classification Using Convolutional Neural Networks
• Electrocardiogram Generation and Feature Extraction Using a Variational Autoencoder (encoder only)
• Author's EcgResNet34

2D models:

• ECG arrhythmia classification using a 2-D convolutional neural network
• MobileNetV2
• EfficientNetB4

Metrics

name	type	leads	model	accuracy	val loss	epoch	batch
exp-025	1D (1x128) - [PEAK[t] - 64, PEAK[t] + 64]	MLII	https://arxiv.org/pdf/1707.01836.pdf	0,9827	0,0726	19	128
exp-030	1D (1x128) - [PEAK[t] - 64, PEAK[t] + 64]	MLII	https://ieeexplore.ieee.org/stamp/stamp.jsp?arnumber=8952723	0,9864	1,5	509	128
exp-031	1D (1x128) - [PEAK[t] - 64, PEAK[t] + 64]	MLII	https://arxiv.org/pdf/2002.00254.pdf	0,9886	0,15	645	128
exp-018	2D (128x128) - [PEAK[t] - 64, PEAK[t] + 64]	MLII	https://arxiv.org/pdf/1804.06812.pdf	0,9920	0,1	76	64
exp-013	2D (128x128) - [PEAK[t] - 64, PEAK[t] + 64]	MLII	MobileNetV2	0,9934	0,088	251	128
exp-021	2D (128x128) - [PEAK[t-1] + 20, PEAK[t+1] - 20]	MLII	EfficientNetB4	0,9935	0,062	154	128
exp-029	1D (1x128) - [PEAK[t] - 64, PEAK[t] + 64]	MLII	Novel	0,9938	0,0500	634	128

Getting started

Training quick start:

1. Download and unzip files into mit-bih directory
2. Install requirements via pip install -r requirements.txt
3. Generate 1D and 2D data files running cd scripts && python dataset-generation-pool.py
4. Create json annotation files
   • For 1D model - cd scripts && python annotation-generation-1d.py
   • For 2D model - cd scripts && python annotation-generation-2d.py
5. Run training - python train.py --configs/training/<config>.json

See CI examples for each model

Testing and visualization

Using EcgResNet34 model as it shows the best metrics

1. Install requirements via pip install -r requirements.txt
2. Create directory named experiments
3. Download the archive and unzip its content into experiments directory
4. Download WFDB format data
5. Change ecg_data path in configs/pipelines/config.json with no extension

{
  ...
  "ecg_data": "./mit-bih/100",
  ...
}

6. Run pipeline - python pipeline.py --configs/pipelines/config.json

The results will be saved as HTML file in experiments/EcgResNet34/results directory

Experiments

The code of all experiments described in the table is in branches experiments/exp-XXX

Other

The repository contains Jupyter Notebooks (see notebooks folder)

Contributors

• Alexander Lyashuk
• Nikolai Zolotykh

Support

Please give a ⭐️ if this project helped you

License

This project is licensed under the MIT License