/ecg-classification

ECG Arrhythmia classification using CNN

Primary LanguageJupyter NotebookMIT LicenseMIT

ECG Arrhythmia classification

pipeline-viz

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:

2D models:

Metrics

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

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 --config 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",
  ...
}
  1. Run pipeline - python pipeline.py --config configs/pipelines/config.json

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

image

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)

image

Contributors

Support

Please give a ⭐️ if this project helped you

License

This project is licensed under the MIT License