A vision transformer framework for arrhythmia classification with spectrogram augmented time-frequency transformed ECG signals

Abstract Models using raw electrocardiogram (ECG) signals for the classification of cardiac arrhythmias do not extract significant spectral information essential to detect complex arrhythmias. Time-frequency (TF) transforms can efficiently describe the spectral-temporal information. However, current methods, which place greater emphasis on local dependencies, are not effective in extracting complex spectral-temporal features. In addition, imbalanced databases result in biased models, decreasing the detection rate of rare arrhythmias of clinical importance. In this work, we propose a novel multi-head self-attention (MHSA) Vision Transformer (ViT) framework (MHSA-ViT-TF) that uses TF representations of raw ECG beats and TF masking augmentation for cardiac arrhythmia classification. The ViT model uses MHSA to retrieve rich contextual and spatial features from time-frequency images. We employed continuous wavelet transform (CWT), short-time Fourier transform (STFT), and Mel-scaled spectrogram TF representations and compared the performance of the ViT model with imbalanced data. Furthermore, TF masking was applied to STFT spectrograms to address class imbalance, a comparatively less investigated approach in ECG signal analysis. We tested the performance of the proposed MHSA-ViT-TF model using the MIT-BIH arrhythmia database and the St. Petersburg INCART 12-lead arrhythmia database. The experimental findings proved that the integration of STFT spectrograms with TF masking augmentation and the ViT classification model produced a prominent performance compared to the state-of-the-art techniques with 99.56 ± 0.9% accuracy, precision of 99.56 ± 0.9%, recall of 99.56 ± 0.9%, and an F-score of 99.56 ± 0.9% using the MIT-BIH arrhythmia database for the seven classes of arrhythmias considered. The relative improvement with balanced data enhanced with spectrograms is 0.66%. We presented the attention maps, presenting the input image regions considered in the final model predictions. Furthermore, the proposed MHSA-ViT-TF model also demonstrated consistent performance using balanced data, achieving an accuracy of 99.50% on the St. Petersburg INCART 12-lead arrhythmia database.