ECG Arrhythmia Classification Using STFT-Based Spectrogram and Convolutional Neural Network

The classification of electrocardiogram (ECG) signals is very important for the automatic diagnosis of heart disease. Traditionally, it is divided into two steps, including the step of feature extraction and the step of pattern classification. Owing to recent advances in artificial intelligence, it has been demonstrated that deep neural network, which trained on a huge amount of data, can carry out the task of feature extraction directly from the data and recognize cardiac arrhythmias better than professional cardiologists. This paper proposes an ECG arrhythmia classification method using two-dimensional (2D) deep convolutional neural network (CNN). The time domain signals of ECG, belonging to five heart beat types including normal beat (NOR), left bundle branch block beat (LBB), right bundle branch block beat (RBB), premature ventricular contraction beat (PVC), and atrial premature contraction beat (APC), were first transformed into time-frequency spectrograms by short-time Fourier transform. Subsequently, the spectrograms of the five arrhythmia types were utilized as input to the 2D-CNN such that the ECG arrhythmia types were identified and classified finally. Using ECG recordings from the MIT-BIH arrhythmia database as the training and testing data, the classification results show that the proposed 2D-CNN model can reach an averaged accuracy of 99.00%. On the other hand, in order to achieve optimal classification performances, the model parameter optimization was investigated. It was found when the learning rate is 0.001 and the batch size parameter is 2500, the classifier achieved the highest accuracy and the lowest loss. We also compared the proposed 2D-CNN model with a conventional one-dimensional CNN model. Comparison results show that the 1D-CNN classifier can achieve an averaged accuracy of 90.93%. Therefore, it is validated that the proposed CNN classifier using ECG spectrograms as input can achieve improved classification accuracy without additional manual pre-processing of the ECG signals.