--- a +++ b/README.md @@ -0,0 +1,117 @@ +# 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](https://docs.google.com/spreadsheets/d/1fmLNC1_1xohEUEoNRodkIoFHnJI3O4HfCHN8Ys_b7nE) + +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](https://arxiv.org/abs/1707.01836)  +- [ECG Heartbeat Classification Using Convolutional Neural Networks](https://ieeexplore.ieee.org/stamp/stamp.jsp?arnumber=8952723)  +- [Electrocardiogram Generation and Feature Extraction Using a Variational Autoencoder](https://arxiv.org/pdf/2002.00254.pdf) (encoder only)  +- **Author's EcgResNet34**  + +#### 2D models: + +- [ECG arrhythmia classification using a 2-D convolutional neural network](https://arxiv.org/abs/1804.06812)  +- MobileNetV2  +- EfficientNetB4  + +#### 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](https://physionet.org/static/published-projects/mitdb/mit-bih-arrhythmia-database-1.0.0.zip) + 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](https://github.com/lxdv/ecg-classification/actions) 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](https://drive.google.com/file/d/1wCy9Y4EQmI3gdVTX77U7ZXa5zPaqLQ5S/view?usp=sharing) the archive and unzip its content into `experiments` directory +4. Download [WFDB format](https://www.physionet.org/physiotools/wpg/wpg_35.htm) 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 --config 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](https://docs.google.com/spreadsheets/d/1fmLNC1_1xohEUEoNRodkIoFHnJI3O4HfCHN8Ys_b7nE) +is in branches **experiments/exp-XXX** + +## Other + +The repository contains Jupyter Notebooks (see `notebooks` folder) + + + +## Contributors + +- [Alexander Lyashuk](mailto:lyashuk.me@gmail.com) +- [Nikolai Zolotykh](mailto:nikolai.zolotykh@gmail.com) + +## Support + +Please give a ⭐️ if this project helped you + +## License + +This project is licensed under the [MIT](LICENCE) License