Classification of imbalanced ECGs through segmentation models and augmented by conditional diffusion model

PQRST-waves pattern recognition model

The continuous MIT-BIH dataset composed of multiple heartbeats is segmented into individual heartbeats for arrhythmia classification purposes. These segmented waveforms consist of PQRST waves, with the features of each waveform being utilized for diagnosing arrhythmias. In Table 2, Faziludeen & Sabiq (2013) utilized the Pan-Tompkins algorithm (Pan & Tompkins, 1985) to detect the QRS complex. Subsequently, 25 features were extracted from the QRS complex by applying the Daubechies four wavelet, which included the mean, variance, standard deviation, the maximum and minimum of the detail coefficient, as well as approximation coefficient. Additionally, the R-R interval (RRI) was employed as a feature for classifying normal, LBBB (L), and PVC (V) using a support vector machine (SVM), achieving an accuracy of 98.95% (Faziludeen & Sabiq, 2013). Senapati, Senapati & Maka (2014) conducted classification of five arrhythmia classes using the scaled conjugate gradient (SCG) algorithm, employing RRI, QRS interval, and morphological features of QRS and T waves as input. The achieved accuracy was 96.2%, with a sensitivity of 96.5%. However, it was noted that the classification performance for the S class was lower compared to other classes due to the limited amount of train data. In Cho et al. (2015), nine patterns were defined based on the amplitude and phase of the QRS. By employing these nine patterns to classify six arrhythmias, a detection rate of 93.72% was achieved. Li & Zhou (2016) employed wavelet packet entropy (WPE) of the QRS complex and the RRI as features. Random forest (RF) was then utilized for the classification of N, S, V, F, and Q classes, achieving an accuracy of 94.61%. Additionally, Yang & Wei (2020) utilized the Pan-Tompkins algorithm to extract the QRS complex. To identify the P wave, the QRS complex was utilized, and features were extracted from the maximum amplitude. Additionally, the visual morphological pattern of QRS complex (VMP-QRS) algorithm was employed to extract visual features from the data, while the adaptive K-means clustering (AKMC) algorithm was utilized to analyze the features of the data through clustering. Various models including neural networks, radial basis function based support vector machine (RBF-SVM), and K-nearest neighbor (KNN) were employed for the classification of the fifteen arrhythmias. Yang & Wei (2020) noted that the proposed models demonstrated improved performance in automated arrhythmia classification techniques, although their emphasis on visual patterns posed a limitation. It is worth mentioning that most of the research articles on pattern recognition for arrhythmia classification predominantly focus on morphological patterns.

Deep learning model

Classifying arrhythmias through pattern recognition poses challenges due to the substantial differences in ECG signals among patients. Consequently, it has become commonplace to employ deep learning models for feature extraction (Liu et al., 2021). Ochiai, Takahashi & Fukazawa (2018) utilized a convolutional denoising autoencoder (CDAE) model, incorporating a convolution layer into a denoising autoencoder (DAE), for the classification of two arrhythmia classes. However, the outcomes revealed a notably lower sensitivity and positive predictive value for the S class. The reason for this limitation was that the number of training data for S class was relatively small compared to other classification classes. Rana & Kim (2019) proposed a long short term memory (LSTM) model to classify the normal, LBBB (L), RBBB (R), Atrial Premature (A), and PVC (V) classes with 95% accuracy. Xu, Jeong & Li (2020) employed a transfer-learned convolutional neural network+bidirectional long short term memory (CNN+BiLSTM) classification model with the 2017 PhysioNet/CinC Challenge dataset. The model achieved an accuracy of 95.9% in classifying five AAMI arrhythmia classes: N, S, V, F, and Q. Qiu et al. (2021) employed the synthetic minority oversampling technique (SMOTE) to augment the data and showed improved performance in arrhythmia classification using the Faster R-CNN model. Anis & Sharma (2022) demonstrated that the CNN model surpassed previous techniques in terms of classification accuracy and efficiency for N, S, V, F, and Q classes.

Comparing the articles that classified AAMI arrhythmia classes in Table 2, it can see that the accuracy of arrhythmia classification models using deep learning is relatively high compared to the accuracy of pattern recognition. Therefore, the performance of deep learning arrhythmia classification methods is better than that of pattern recognition methods. This study evaluates the performance of the ECG classification model against two models, MobileNetV2 and the transformer, which have demonstrated strong performance among various deep learning architectures.

MobileNetV2 (Sandler et al., 2018), introduced by Google, represents a lightweight residual CNN architecture predominantly utilized in mobile and embedded systems. Its distinction from traditional CNN architectures lies in the adoption of depthwise separable convolutions, enabling the computation of convolutions with a substantially reduced parameter count. This method involves executing depthwise convolutions, which apply a filter to each input channel for capturing spatial features, followed by pointwise convolutions that aggregate the depthwise convolution outputs through 1x1 convolutions. Furthermore, MobileNetV2 innovates with the introduction of the inverted residual structure, which serves to diminish computational load while improving accuracy. This structure is composed of two block types: those with a stride of 1, mirroring residual blocks in preserving feature map dimensions and facilitating learning, and those with a stride of 2, which halve the feature map size, aiding in the model’s efficiency. Consequently, MobileNetV2, despite its efficiency-oriented design, achieves notable performance.

The transformer model, which overcomes the long-term dependency issues inherent in recurrent neural networks (RNNs), is a Sequence-to-Sequence model that has served as the foundation for revolutionary advances in the field of natural language processing. It is composed of encoders and decoders that utilize the attention mechanism (Vaswani et al., 2017). The attention mechanism calculates the importance of different positions from each position in the input, considering their relationships. The calculation process of attention is as given in Eq. (1). (1)${\displaystyle \mathrm{Attention}\left(Q,K,V\right)=\mathrm{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V.}$

The symbols Q, K, and V denote the constructs of query, key, and value, respectively, with each being matrices derived from the multiplication of the input matrix X by the corresponding weight matrices $W_i^Q,W_i^K,W_i^V$. The d_k refers to the dimension of the key vector. Calculation of the attention score is achieved through the dot product between the Q and K matrices, followed by division by $\sqrt{d_k}$ to facilitate stable gradient computation. The attention distribution is then determined using the softmax function, and this distribution is subsequently applied to the V matrix to generate the attention matrix. Through the attention matrix, features from positions of highest importance at each location of the same input can be extracted. The transformer model can be utilized as a classifier by adding a linear layer to the trained encoder. Consequently, the employment of the attention mechanism enables the transformer to efficiently capture dependencies across all positions, with the benefit of parallel computation facilitating a substantial reduction in training.

SMOTE-Tomek

SMOTE-Tomek is an augmentation method that employs both SMOTE (Chawla et al., 2002) oversampling and Tomek-links undersampling techniques. SMOTE generates data for the minority classes by utilizing the nearest neighbor algorithm. On the other hand, Tomek-links are used to remove data from the majority class that exists in pairs at the boundary between the majority and minority classes. Consequently, SMOTE-Tomek adjusts class balance by generating data with SMOTE and removing data using Tomek-links.

SMOTE was applied to the MIT-BIH arrhythmia dataset by Pandey & Janghel (2019) to perform oversampling. Subsequently, a classification of five arrhythmia classes was conducted using an 11-layer deep CNN model. The accuracy of the classification model trained on the augmented data showed an improvement of 2.26% compared to the CNN classifier trained on the original dataset. This indicates that using data augmented by SMOTE can enhance the model’s classification performance. The MIT-BIH arrhythmia dataset was augmented using SMOTE by Xiaolin, Cardiff & John (2020), followed by the classification of N, S, V, F, Q classes using a one-dimension CNN. Compared to the classifier trained on the imbalanced dataset, the performance of the classifier trained on the balanced dataset exhibited a decrease in accuracy by 1%, specificity by 1.54%, precision by 6.89%, and F1 score by 2.48%. However, an improvement of 2.09% was observed in sensitivity.

VAE

VAE (Kingma & Welling, 2022) is a model designed to learn the distribution of data in order to generate similar data. It comprises an encoder and a decoder. The structure of the VAE utilized in this study is depicted in Fig. 1. In this study, the input to the VAE is a segment of a one-dimension ECG signal. This representation has the dimension of (500,1), since the signal has been segmented to represent a single heartbeat. In a VAE, the encoder plays a crucial role in estimating the distribution of the data. It estimates a latent vector containing the mean (µ) and variance (σ²) parameters of the input data distribution. These parameters are assumed to follow a Gaussian distribution. The decoder in a VAE uses latent vectors to learn to approximate the distribution of the data while reconstructing to the input data. The encoder in a VAE employs convolution to train the distribution of the input data. Consequently, the encoder approximates the distribution of the latent vector, denoted as z, given the input data x. This latent vector z represents the mean and variance of the input data. Subsequently, the decoder utilizes the latent vector and convolution to reconstruct the input data. It is important to note that the decoder is symmetric to the encoder in the VAE architecture. The loss function is equal to Eq. (2). (2)${\displaystyle \mathrm{Loss}={\mathbb{E}}_{q_{\varphi }\left(\mathbf{z}|\mathbf{x}\right)}\left[log\left(p\left(\mathbf{x}|\mathbf{z}\right)\right)\right]-\mathrm{KL}\left(q_{\varphi }\left(\mathbf{z}|\mathbf{x}\right)\left|\right|p\left(\mathbf{z}\right)\right).}$

The first term corresponds to the reconstruction loss, while the second term corresponds to the Kullback–Leibler Divergence (KLD). The reconstruction loss measures the difference between the input data and the reconstructed data. KLD is a function that approximates the probability distribution of an idealized latent variable and the probability distribution of an estimated latent variable. Consequently, during training, the model is optimized towards minimizing the overall loss function, which encompasses both the reconstruction loss and the KLD. This optimization process ensures that the model trains to accurately reconstruct the input data while also appropriately modeling the latent space distribution.

To evaluate the performance of the VAE model, the test data, 10% of the total dataset, was used to calculate the average cosine similarity per class. The test data was fed into the trained VAE model, and the decoder was used to reconstruct the data. Subsequently, the cosine similarity between the reconstructed data and the original data was then calculated. The obtained cosine similarity values were 0.99 for the S class, 0.99 for the V class, and 0.99 for the F class. Higher cosine similarity values indicate greater similarity between the original and reconstructed data, suggesting that the VAE model effectively estimated the distribution of the data.

Kuznetsov et al. (2021) used VAE in the Lobachevsky University Electrocardiography Database (LUDB) to generate an ECG for a single heartbeat. Kuznetsov et al. (2021) presented a VAE model that outputs two 25-dimension vectors from the encoder and inputs them to the decoder. The maximum mean discrepancy value used as an evaluation metric for VAE is 3.83 × 10⁻³. Kuznetsov et al. (2021) stated that their proposed model is relatively lightweight, but it cannot generate the entire ECG sequence.

Conditional diffusion

Diffusion (Ho, Jain & Abbeel, 2020) is a generative model that trains the data distribution through a forward process, adding Gaussian noise to the initial data x₀, and a reverse process, eliminating the noise from the Gaussian-noised data x_T, thus reconstructing it to the original data x₀, as illustrated in Fig. 2. Given that the forward and reverse processes are defined as a Markov chain, there exists the capacity to progressively add and remove noise from the data.

Considering the data x_t−1 in the forward process, Eq. (3) illustrates the procedure of adding Gaussian noise and subsequently computing the data x_t at the timestamp. In Eq. (3), β_t is the variance schedule. Using β_t in Eq. (3), defining ${\alpha }_t=1-{\beta }_t,{\bar{\alpha }}_t={\Pi }_{i=1}^t{\beta }_t$, Eq. (3) can be expressed as Eq. (4). Equation (4) signifies the ability to directly compute x_t, by adding Gaussian noise to the original data x₀. Through iterative application of the reparameterization trick to Eq. (4), the noise can be calculated, as demonstrated in Eq. (5). (3)${\displaystyle q\left({\mathbf{x}}_t|{\mathbf{x}}_{t-1}\right):=\mathcal{N}\left({\mathbf{x}}_t;\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1},{\beta }_t\mathbf{I}\right)}$ (4)${\displaystyle q\left({\mathbf{x}}_t|{\mathbf{x}}_0\right)=\mathcal{N}\left({\mathbf{x}}_t;\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0,\left(1-{\bar{\alpha }}_t\right)\mathbf{I}\right)}$ (5)${\displaystyle {\mathbf{x}}_t=\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}ϵϵ\sim \mathcal{N}\left(0,\mathbf{I}\right).}$

In the reverse process, the process of reconstructing the original data x₀ with noise x_t that follows a Gaussian distribution is shown in Eq. (6). (6)${\displaystyle p_{\theta }\left({\mathbf{x}}_{t-1}|{\mathbf{x}}_t\right):=\mathcal{N}\left({\mathbf{x}}_{t-1};{\mu }_{\theta }\left({\mathbf{x}}_t,t\right),{\Sigma }_{\theta }\left({\mathbf{x}}_t,t\right)\right).}$

As the reverse process requires training, the loss function is equivalent to Eq. (7). (7)${\displaystyle L_{\mathrm{simple}}\left(\theta \right):={\mathbb{E}}_{t,{\mathbf{x}}_0,ϵ}\left[\left|\right|ϵ-{ϵ}_{\theta }\left(\sqrt{{\bar{\alpha }}_t}{\mathbf{x}}_0+\sqrt{1-{\bar{\alpha }}_t}ϵ,t\right)|{|}^2\right].}$

Equation (7) stands for finding the minimizing KLD between the the probability distribution of the actual reverse process and the probability distribution of the denoised x_t−1 given x_t as input to the U-Net model.

In this study, one-dimension ECG sequence data was reshaped into two-dimensions and used as input to the diffusion model. Subsequently, the denoising diffusion probabilistic model (DDPM) was trained as a convolution-based U-Net model. The data, generated in two-dimensions through the reverse process of diffusion, reshaped into one-dimension for classification purposes. Diffusion’s U-Net architecture (Ronneberger, Fischer & Brox, 2015) encompasses a contracting path, an expanding path, and a bridge. The contracting path employs convolution and pooling operations on input data to acquire context. In contrast, the expanding path utilizes features extracted from the contracting path to capture details, incorporating convolution and upsampling operations. Ultimately, the bridge facilitates the connection of features obtained from both the contracting and expanding paths, thereby enhancing overall performance.

Given the morphological similarity observed in certain classes of ECG data, it becomes imperative to generate data that effectively encapsulates the distinctive features of each class. Thus, this study applies classifier-free guidance (Ho & Salimans, 2022) as a means to generate data. Classifier-free guidance facilitates the training process of DDPM by incorporating class label information. The mechanism of classifier-free guidance is represented by Eq. (8). (8)${\displaystyle {\widetilde{ϵ}}_{\theta }\left({\mathbf{z}}_{\lambda },\mathbf{c}\right)=\left(1+\omega \right){ϵ}_{\theta }\left({\mathbf{z}}_{\lambda },\mathbf{c}\right)-\omega {ϵ}_{\theta }\left({\mathbf{z}}_{\lambda }\right)}$

In Eq. (8), ϵ_θ(z_λ, c) represents a conditional diffusion model with class labels as input and ϵ_θ(z_λ) means an unconditional diffusion model that trains without class labels and ω is the weight value. Consequently, the reverse process is acquired through the interpolation of the conditional diffusion model and the unconditional diffusion model. In summary, this study employs class labels as input and employs U-Net to learn the data distribution. Subsequently, data from classes S, V, and F are sampled based on the number of N class data. To evaluate the performance of the conditional diffusion model, the test data, which is 10% of the total dataset, was used to calculate the average cosine similarity per class. In this case, the cosine similarity was calculated using the actual noise and the noise predicted by the model. The obtained cosine similarity values are as follows: 0.989 for the S class, 0.992 for the V class, and 0.997 for the F class. These results indicate that the conditional diffusion model successfully learned the data distribution.

Annotation-based segmentation model

The structure of the model that performs annotation-based segmentation and classifies arrhythmias is shown in Fig. 4. In the MIT-BIH dataset, the R-peak label of each heartbeat is mapped. The study was divided into separate segments between 0.25 s before and 0.83 s after the mapped R-peak. The MIT-BIH dataset already provides arrhythmia labels for each R-peak, facilitating our use of this data for augmentation. To address imbalance, this research generated additional segment data for the imbalanced classes S, V, and F to equalize them with the N class using augmentation techniques such as SMOTE-Tomek, VAE, and conditional diffusion. This augmented data was then merged with the original dataset to create a comprehensive training dataset. Subsequently, the dataset was then trained to classify four arrhythmia classes using MobileNetV2 and the transformer models. The specifics of the arrhythmia classification process are further elaborated in classification section.

Automated segmentation model

Figure 5 presents the configuration of the model designed for automated segmentation employing deep learning frameworks. This model was implemented the Faster R-CNN architecture, similar to that described in Qiu et al. (2021), although certain parameters were adjusted to meet the constraints of the experimental setup. The preprocessing procedure adopted for this automated segmentation model is consistent with that of the annotation-based segmentation model.

The first step of the automated segmentation model is “Feature Extraction”, which is a step to train the features of the classes to extract the feature weights of the model. The Faster R-CNN model facilitates object detection by leveraging a detailed feature map of the data. Consequently, for effective ECG classification, it is essential to derive a feature map that adequately represents the data’s attributes. In this study, the “Feature Extraction” process, indicated by the green dotted line in Fig. 5, is designed to extract the model’s feature weights for the purpose of training the characteristics of the arrhythmia classes N, S, V, and F. The segmentation framework, forming the backbone of our approach, utilizes R-peaks on the preprocessed data to produce segments for each class. Since class imbalance is prevalent, Backbone Augmentation is performed to generate additional segments, using methods such as SMOTE-Tomek, VAE, and conditional diffusion to augment the representation of the minority classes. Fundamentally, as illustrated in Fig. 5, Backbone Augmentation is designed to increase the number of representations of features for imbalanced classes, thus promoting a more equitable training of features. After this augmentation, the backbone process uses MobileNetV2 to deepen the learning and refinement of each class’s features. As a result, the “Feature Extraction” step is a collection of trained weight values, which is the basis for the next segmentation step.

The second step of the automated segmentation model is “Segmentation”. The “Segmentation” process plays a crucial role in Faster R-CNN by predicting the regions where objects are located. This process leverages data from consecutive beats and labeled information to calculate a feature map through the application of a pretrained backbone. This procedure generates anchor boxes of different predefined sizes for predicting the regions containing objects. The next step is to execute a region proposal network (RPN), which is responsible for calculating the probability of object presence for each anchor box and adjusting the coordinates of the box using a bounding box regressor. The process can extract candidate regions using anchor boxes and object presence probabilities and bounding box regressors. This is the process of learning anchor boxes to better detect the location of objects. Since arrhythmias are caused by the heart beating irregularly and the length of each heartbeat is not constant, Faster R-CNN can effectively segment candidate regions. Using the feature map extracted by the Backbone network and the predicted candidate regions, select the anchor boxes with a high probability of object presence. This process is called region of interest (ROI). However, since the ROIs are not uniform in size, training requires the next step: Region Pooling.

The third step is “Region Pooling”, which equalizes the size of the data for training. Owing to the variability in segment sizes across different types of arrhythmias, the dimensions of the predicted ROI are not constant and it is important to standardize the size of the input data for effective classification of arrhythmia classes with deep learning. Consequently, “Region Pooling” is necessary to normalize the segments of varying sizes. “Region pooling” is the technique of transforming the feature map of an ROI into a fixed-size vector through max pooling. In our study, max pooling was employed on the ROIs within the feature map generated by the backbone process, ensuring uniformity in ROI sizes. In other words, the feature map of the ROI was segmented into small, fixed-size windows, with max pooling applied to each window. This approach effectively standardized the feature map dimensions while preserving only the most significant features from the ROIs of disparate sizes, thus maintaining a consistent feature map size.

The fourth step is the “Classification” step, which is the model that classifies arrhythmias into the red areas in Fig. 5. The outcome of “Region Pooling” is just dividing the preprocessed data into segments of uniform size. The MIT-BIH dataset is inherently imbalanced, exhibiting significant class imbalances. To address this issue, Feature Augmentation is conducted utilizing techniques such as SMOTE-Tomek, VAE, and conditional diffusion. Subsequently, a train dataset that integrates both the augmented and original data is employed in conjunction with CNN and the transformer models to classify the arrhythmia classes N, S, V, and F. Further details on the classification of arrhythmias are provided in classification section.