
# 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) ![1707.01384](https://github.com/lxdv/ecg-classification/workflows/1707.01384/badge.svg)
- [ECG Heartbeat Classification Using Convolutional Neural Networks](https://ieeexplore.ieee.org/stamp/stamp.jsp?arnumber=8952723) ![1911.IEEE](https://github.com/lxdv/ecg-classification/workflows/1911.IEEE/badge.svg)
- [Electrocardiogram Generation and Feature Extraction Using a Variational Autoencoder](https://arxiv.org/pdf/2002.00254.pdf) (encoder only) ![2002.00254](https://github.com/lxdv/ecg-classification/workflows/2002.00254/badge.svg)
- **Author's EcgResNet34** ![EcgResNet34](https://github.com/lxdv/ecg-classification/workflows/EcgResNet34/badge.svg)

#### 2D models:

- [ECG arrhythmia classification using a 2-D convolutional neural network](https://arxiv.org/abs/1804.06812) ![1804.06812](https://github.com/lxdv/ecg-classification/workflows/1804.06812/badge.svg)
- MobileNetV2 ![MobileNetV2](https://github.com/lxdv/ecg-classification/workflows/MobileNetV2/badge.svg)
- EfficientNetB4 ![EfficientNetB4](https://github.com/lxdv/ecg-classification/workflows/EfficientNetB4/badge.svg)

#### 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)

![image](etc/confusion-matrix.png)

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

	a/README.md		b/README.md
1	# ECG Arrhythmia classification	1	# ECG Arrhythmia classification
2		2
3	![pipeline-viz](etc/pipeline-viz.png)
4
5	The repository contains code for Master's degree dissertation -	3	The repository contains code for Master's degree dissertation -
6	Diagnosis of Diseases by ECG Using Convolutional Neural Networks.	4	Diagnosis of Diseases by ECG Using Convolutional Neural Networks.
7	Only CNN neural network models are considered in the paper and the repository.	5	Only CNN neural network models are considered in the paper and the repository.
8	As a part of the work, more than 30 experiments have been run.	6	As a part of the work, more than 30 experiments have been run.
9	The table with all experiments and their metrics is available by the [link](https://docs.google.com/spreadsheets/d/1fmLNC1_1xohEUEoNRodkIoFHnJI3O4HfCHN8Ys_b7nE)	7	The table with all experiments and their metrics is available by the [link](https://docs.google.com/spreadsheets/d/1fmLNC1_1xohEUEoNRodkIoFHnJI3O4HfCHN8Ys_b7nE)
10		8
11	The best 1D and 2D CNN models are presented in the repository	9	The best 1D and 2D CNN models are presented in the repository
12	The repository follows _config_ principle and can be run in the following modes:	10	The repository follows _config_ principle and can be run in the following modes:
13		11
14	- Training - use `train.py --config configs/training/<config>.json` to train the model	12	- Training - use `train.py --config configs/training/<config>.json` to train the model
15	- Validation - use `inference.py --config configs/inference/config.json` to validate the model	13	- Validation - use `inference.py --config configs/inference/config.json` to validate the model
16	- Pipeline - use `pipeline.py --config configs/pipelines/config/json` to test the model using ECG data (i.e. data generation, running, visualization the results)	14	- Pipeline - use `pipeline.py --config configs/pipelines/config/json` to test the model using ECG data (i.e. data generation, running, visualization the results)
17		15
18	All available models and all necessary information are described below	16	All available models and all necessary information are described below
19		17
20	Python 3.7 and PyTorch are used in the project	18	Python 3.7 and PyTorch are used in the project
21	GitHub actions are used for installing dependencies and training implemented models	19	GitHub actions are used for installing dependencies and training implemented models
22		20
23	Program - Data Mining	21	Program - Data Mining
24	Department - Computer Science	22	Department - Computer Science
25		23
26	Principal Investigator - Nikolai Yu. Zolotykh	24	Principal Investigator - Nikolai Yu. Zolotykh
27	_National Research University - Higher School of Economics_	25	_National Research University - Higher School of Economics_
28		26
29	## Implemented models	27	## Implemented models
30		28
31	#### 1D models:	29	#### 1D models:
32		30
33	- [Cardiologist-Level Arrhythmia Detection with Convolutional Neural Networks](https://arxiv.org/abs/1707.01836) ![1707.01384](https://github.com/lxdv/ecg-classification/workflows/1707.01384/badge.svg)	31	- [Cardiologist-Level Arrhythmia Detection with Convolutional Neural Networks](https://arxiv.org/abs/1707.01836) ![1707.01384](https://github.com/lxdv/ecg-classification/workflows/1707.01384/badge.svg)
34	- [ECG Heartbeat Classification Using Convolutional Neural Networks](https://ieeexplore.ieee.org/stamp/stamp.jsp?arnumber=8952723) ![1911.IEEE](https://github.com/lxdv/ecg-classification/workflows/1911.IEEE/badge.svg)	32	- [ECG Heartbeat Classification Using Convolutional Neural Networks](https://ieeexplore.ieee.org/stamp/stamp.jsp?arnumber=8952723) ![1911.IEEE](https://github.com/lxdv/ecg-classification/workflows/1911.IEEE/badge.svg)
35	- [Electrocardiogram Generation and Feature Extraction Using a Variational Autoencoder](https://arxiv.org/pdf/2002.00254.pdf) (encoder only) ![2002.00254](https://github.com/lxdv/ecg-classification/workflows/2002.00254/badge.svg)	33	- [Electrocardiogram Generation and Feature Extraction Using a Variational Autoencoder](https://arxiv.org/pdf/2002.00254.pdf) (encoder only) ![2002.00254](https://github.com/lxdv/ecg-classification/workflows/2002.00254/badge.svg)
36	- Author's EcgResNet34 ![EcgResNet34](https://github.com/lxdv/ecg-classification/workflows/EcgResNet34/badge.svg)	34	- Author's EcgResNet34 ![EcgResNet34](https://github.com/lxdv/ecg-classification/workflows/EcgResNet34/badge.svg)
37		35
38	#### 2D models:	36	#### 2D models:
39		37
40	- [ECG arrhythmia classification using a 2-D convolutional neural network](https://arxiv.org/abs/1804.06812) ![1804.06812](https://github.com/lxdv/ecg-classification/workflows/1804.06812/badge.svg)	38	- [ECG arrhythmia classification using a 2-D convolutional neural network](https://arxiv.org/abs/1804.06812) ![1804.06812](https://github.com/lxdv/ecg-classification/workflows/1804.06812/badge.svg)
41	- MobileNetV2 ![MobileNetV2](https://github.com/lxdv/ecg-classification/workflows/MobileNetV2/badge.svg)	39	- MobileNetV2 ![MobileNetV2](https://github.com/lxdv/ecg-classification/workflows/MobileNetV2/badge.svg)
42	- EfficientNetB4 ![EfficientNetB4](https://github.com/lxdv/ecg-classification/workflows/EfficientNetB4/badge.svg)	40	- EfficientNetB4 ![EfficientNetB4](https://github.com/lxdv/ecg-classification/workflows/EfficientNetB4/badge.svg)
43		41
44	#### Metrics	42	#### Metrics
45		43
46	\| name \| type \| model \| accuracy \| val loss \|	44	\| name \| type \| model \| accuracy \| val loss \|
47	\| -------- \| ----------------------------------------------- \| ------------------------------------------------------------ \| ------------ \| ------------ \|	45	\| -------- \| ----------------------------------------------- \| ------------------------------------------------------------ \| ------------ \| ------------ \|
48	\| exp-025 \| 1D (1x128) - [PEAK[t] - 64, PEAK[t] + 64] \| https://arxiv.org/pdf/1707.01836.pdf \| 0,9827 \| 0,0726 \|	46	\| exp-025 \| 1D (1x128) - [PEAK[t] - 64, PEAK[t] + 64] \| https://arxiv.org/pdf/1707.01836.pdf \| 0,9827 \| 0,0726 \|
49	\| exp-030 \| 1D (1x128) - [PEAK[t] - 64, PEAK[t] + 64] \| https://ieeexplore.ieee.org/stamp/stamp.jsp?arnumber=8952723 \| 0,9864 \| 1,5 \|	47	\| exp-030 \| 1D (1x128) - [PEAK[t] - 64, PEAK[t] + 64] \| https://ieeexplore.ieee.org/stamp/stamp.jsp?arnumber=8952723 \| 0,9864 \| 1,5 \|
50	\| exp-031 \| 1D (1x128) - [PEAK[t] - 64, PEAK[t] + 64] \| https://arxiv.org/pdf/2002.00254.pdf \| 0,9886 \| 0,15 \|	48	\| exp-031 \| 1D (1x128) - [PEAK[t] - 64, PEAK[t] + 64] \| https://arxiv.org/pdf/2002.00254.pdf \| 0,9886 \| 0,15 \|
51	\| exp-018 \| 2D (128x128) - [PEAK[t] - 64, PEAK[t] + 64] \| https://arxiv.org/pdf/1804.06812.pdf \| 0,9920 \| 0,1 \|	49	\| exp-018 \| 2D (128x128) - [PEAK[t] - 64, PEAK[t] + 64] \| https://arxiv.org/pdf/1804.06812.pdf \| 0,9920 \| 0,1 \|
52	\| exp-013 \| 2D (128x128) - [PEAK[t] - 64, PEAK[t] + 64] \| MobileNetV2 \| 0,9934 \| 0,088 \|	50	\| exp-013 \| 2D (128x128) - [PEAK[t] - 64, PEAK[t] + 64] \| MobileNetV2 \| 0,9934 \| 0,088 \|
53	\| exp-021 \| 2D (128x128) - [PEAK[t-1] + 20, PEAK[t+1] - 20] \| EfficientNetB4 \| 0,9935 \| 0,062 \|	51	\| exp-021 \| 2D (128x128) - [PEAK[t-1] + 20, PEAK[t+1] - 20] \| EfficientNetB4 \| 0,9935 \| 0,062 \|
54	\| exp-029 \| 1D (1x128) - [PEAK[t] - 64, PEAK[t] + 64] \| Novel EcgResNet34 \| 0,9938 \| 0,0500 \|	52	\| exp-029 \| 1D (1x128) - [PEAK[t] - 64, PEAK[t] + 64] \| Novel EcgResNet34 \| 0,9938 \| 0,0500 \|
55		53
56	## Getting started	54	## Getting started
57		55
58	Training quick start:	56	Training quick start:
59		57
60	1. [Download](https://physionet.org/static/published-projects/mitdb/mit-bih-arrhythmia-database-1.0.0.zip)	58	1. [Download](https://physionet.org/static/published-projects/mitdb/mit-bih-arrhythmia-database-1.0.0.zip)
61	and unzip files into `mit-bih` directory	59	and unzip files into `mit-bih` directory
62	2. Install requirements via `pip install -r requirements.txt`	60	2. Install requirements via `pip install -r requirements.txt`
63	3. Generate 1D and 2D data files running `cd scripts && python dataset-generation-pool.py`	61	3. Generate 1D and 2D data files running `cd scripts && python dataset-generation-pool.py`
64	4. Create `json` annotation files	62	4. Create `json` annotation files
65	- For 1D model - `cd scripts && python annotation-generation-1d.py`	63	- For 1D model - `cd scripts && python annotation-generation-1d.py`
66	- For 2D model - `cd scripts && python annotation-generation-2d.py`	64	- For 2D model - `cd scripts && python annotation-generation-2d.py`
67	5. Run training - `python train.py --config configs/training/<config>.json`	65	5. Run training - `python train.py --config configs/training/<config>.json`
68		66
69	See [CI examples](https://github.com/lxdv/ecg-classification/actions) for each model	67	See [CI examples](https://github.com/lxdv/ecg-classification/actions) for each model
70		68
71	## Testing and visualization	69	## Testing and visualization
72		70
73	_Using EcgResNet34 model as it shows the best metrics_	71	_Using EcgResNet34 model as it shows the best metrics_
74		72
75	1. Install requirements via pip install -r requirements.txt	73	1. Install requirements via pip install -r requirements.txt
76	2. Create directory named `experiments`	74	2. Create directory named `experiments`
77	3. [Download](https://drive.google.com/file/d/1wCy9Y4EQmI3gdVTX77U7ZXa5zPaqLQ5S/view?usp=sharing) the archive and unzip its content into `experiments` directory	75	3. [Download](https://drive.google.com/file/d/1wCy9Y4EQmI3gdVTX77U7ZXa5zPaqLQ5S/view?usp=sharing) the archive and unzip its content into `experiments` directory
78	4. Download [WFDB format](https://www.physionet.org/physiotools/wpg/wpg_35.htm) data	76	4. Download [WFDB format](https://www.physionet.org/physiotools/wpg/wpg_35.htm) data
79	5. Change `ecg_data` path in `configs/pipelines/config.json` with no extension	77	5. Change `ecg_data` path in `configs/pipelines/config.json` with no extension
80		78
81	```	79	```
82	{	80	{
83	...	81	...
84	"ecg_data": "./mit-bih/100",	82	"ecg_data": "./mit-bih/100",
85	...	83	...
86	}	84	}
87	```	85	```
88		86
89	6. Run pipeline - `python pipeline.py --config configs/pipelines/config.json`	87	6. Run pipeline - `python pipeline.py --config configs/pipelines/config.json`
90		88
91	The results will be saved as HTML file in `experiments/EcgResNet34/results` directory	89	The results will be saved as HTML file in `experiments/EcgResNet34/results` directory
92		90
93	![image](etc/pipeline-example.png)
94
95	## Experiments	91	## Experiments
96		92
97	The code of all experiments described in the [table](https://docs.google.com/spreadsheets/d/1fmLNC1_1xohEUEoNRodkIoFHnJI3O4HfCHN8Ys_b7nE)	93	The code of all experiments described in the [table](https://docs.google.com/spreadsheets/d/1fmLNC1_1xohEUEoNRodkIoFHnJI3O4HfCHN8Ys_b7nE)
98	is in branches experiments/exp-XXX	94	is in branches experiments/exp-XXX
99		95
100	## Other	96	## Other
101		97
102	The repository contains Jupyter Notebooks (see `notebooks` folder)	98	The repository contains Jupyter Notebooks (see `notebooks` folder)
103		99
104	![image](etc/confusion-matrix.png)	100	![image](etc/confusion-matrix.png)
105		101
106	## Contributors	102	## Contributors
107		103
108	- [Alexander Lyashuk](mailto:lyashuk.me@gmail.com)	104	- [Alexander Lyashuk](mailto:lyashuk.me@gmail.com)
109	- [Nikolai Zolotykh](mailto:nikolai.zolotykh@gmail.com)	105	- [Nikolai Zolotykh](mailto:nikolai.zolotykh@gmail.com)
110		106
111	## Support	107	## Support
112		108
113	Please give a ⭐️ if this project helped you	109	Please give a ⭐️ if this project helped you
114		110
115	## License	111	## License
116		112
117	This project is licensed under the [MIT](LICENCE) License	113	This project is licensed under the [MIT](LICENCE) License