An Improved VIT Neural Network for Arrhythmia Classification Based on Wavelet Spectral Feature Extraction

The extraction of morphological features from electrocardiograms (ECGs) is a crucial means of detecting different types of cardiac arrhythmias. Effective methods to enhance the efficiency of feature extraction have been lacking, which in turn impacts the accuracy of detection and classification. In response to this issue, an improved Vision Transformer (VIT)[1] neural network based on wavelet spectral feature extraction is proposed for the automatic classification of ECG features. In the processing of ECG signals, heartbeats are separated from the original signal, and then a time-series data augmentation method called TTS-GAN (Transformer Time-Series GAN)[2] is employed to expand the ECG signal. Next, the signal is transformed into the wavelet domain, and the spectrogram is input into the VIT neural network for feature learning. This study utilizes dual-lead ECG medical sample data from the MIT-BIH database and classifies them according to the American Association of Medical Instrumentation (AAMI) ECG five-class standard. Experimental results indicate that compared to traditional VIT network classification methods, the improved VIT neural network model, based on wavelet spectral feature extraction, enhances attention to local details and specific regions. Accuracy, sensitivity, and the comprehensive evaluation F1 score all show significant improvement; therefore, it is better suited for clinical diagnosis, potentially replacing the need for ECG doctors in certain cases.