Aims:
We propose a novel ECG classification approach that leverages YOLOv8- based ECG waveform image classification and tree-model-based superclass prediction exploiting metadata and sub-classes, providing a comprehensive method for classifying ECG subject conditions.
Methods:
Image Preprocessing: We maximize information usage of the PTB-XL dataset by creating a unique image dataset for YOLOv8, plotting ECG signals on blank background. This allows the YOLOv8 model to focus solely on learning features from the ECG waveforms.
Superclass Prediction: To utilize the previously untouched metadata, we propose a correlational analysis based on decision trees (potentially random forest or C50 classifiers) to impute missing superclass labels. This analysis leverages existing subclass labels and subject metadata to predict superclasses, improving training data quality.
Ensemble Classification: While YOLOv8 performs purely image-based ECG classification, the machine learning model predicts superclasses based on metadata. This combined ensemble approach aims for a comprehensive classification into both subclasses and superclasses.
Results:
Initial testing utilized a training set (18865 images) and a validation set (2500 images). The YOLOv8n model achieved a top 1% classification accuracy of 64.7% after 40 epochs and 66.7% after 80 epochs across 5 super-classes. Evaluation of a larger YOLOv8x model with validation loss decrease based on early stopping is ongoing. Insights from the decision tree analysis will be used to re-annotate images with missing superclass labels, leveraging subclass predictions for improved YOLOv8 training.
Conclusion:
This work demonstrates the effectiveness of combining YOLOv8 for image-based ECG classification with tree-based machine learning models for metadata analysis. This ensemble approach has the potential to improve overall ECG subject classification accuracy. Future work will explore the use of decision trees and random forest/C50 classifiers for superclass prediction through metadata and potentially ensemble the two models for comprehensive prediction of the ECG superclass of the subject.