A Deep Learning Baseline of MIT-BIH ECG Recognition

Details of this project can be found in my blog articles:

使用Python+TensorFlow2构建基于卷积神经网络（CNN）的ECG心电信号识别分类（一）

使用Python+TensorFlow2构建基于卷积神经网络（CNN）的ECG心电信号识别分类（二）

使用Python+TensorFlow2构建基于卷积神经网络（CNN）的ECG心电信号识别分类（三）

使用Python+TensorFlow2构建基于卷积神经网络（CNN）的ECG心电信号识别分类（四）

Requirements

common packages:

pip install wfdb, pywavelets, numpy, matplotlib, seaborn, sklearn, tensorboard

tensorflow packages:

tensorflow>=2.3.0

pytorch packages:

torch>=1.4.0

Usage

for tensorflow:

python main_tf.py

for pytorch:

python main_torch.py

With ecg_model.h5 or ecg_model.pt in the directory, the commands above will load the model and evaluate it.

If you want to train the model, please remove the model file and run the commands.

visualize the log in tensorboard:

tensorboard --logdir=./logs/<log_dir_name>

File Structure

main_tf.py: code implemented by TensorFlow2

main_torch.py: code implemented by PyTorch

utils.py: util functions about data processing and plotting

Note

Be careful! The input of the convolution (1d) layer in tensorflow and pytorch are different.

In tensorflow it is [batch_size, length, channel], while in pytorch is [batch_size, channel, length].

The denoised ECG data shape in numpy format is [batch_size, length].

Therefore, the input of the convolution layer in tensorflow needs to be reshaped to [batch_size, length, channel], while that in pytorch needs to be reshaped to [batch_size, channel, length], where channel equals 1.

Data shape in this version of code may vary with my blogs.

Feel Helpful?

If you find this project helpful, please give me a star. Thanks!

If you could tip me a cup of coke, I would be very grateful!