Arrhythmia Classification with Explainable AI

Arrhythmia Classification with Explainable AI using RESNET50 1d and 1d Gradcam

Comment

Inspiration

Cardiovascular diseases remain one of the leading causes of death worldwide, and many heart rhythm disorders often go undetected until they become serious. While hospitals rely on 12-lead ECGs for diagnosis, wearable devices such as smartwatches typically record only a single ECG lead. This makes automatic arrhythmia classification significantly more challenging due to the limited information available.

The goal of this project was to investigate whether a deep learning model trained on single-lead ECG signals could accurately classify clinically meaningful heart rhythms while remaining suitable for deployment on wearable devices. I also wanted the model to provide visual explanations for its predictions instead of behaving like a complete "black box."

What it does

This project classifies six clinically meaningful heart rhythm categories from a single-lead ECG signal:

  • Atrial Arrhythmia (AF + AFIB + AT)
  • Sinus Arrhythmia
  • Sinus Bradycardia
  • Sinus Rhythm
  • Sinus Tachycardia
  • Supraventricular Tachycardia (SVT + AVRT)

Users can upload a single-lead ECG recording through a simple web interface. The application then displays:

  • Predicted heart rhythm
  • Confidence scores for every class
  • An uncertainty indicator when the prediction confidence falls below a configurable threshold
  • A 1D Grad-CAM visualization highlighting the regions of the ECG that contributed most to the prediction

The project is designed around the workflow of wearable ECG devices, where interpretability is just as important as prediction accuracy.

How I built it

Instead of training a model entirely from scratch, I first used Bootstrap Your Own Latent (BYOL) for self-supervised representation learning to pre-train a 1D ResNet-50 encoder on ECG signals. The pre-trained encoder was then fine-tuned for arrhythmia classification.

During supervised training I:

  • Fine-tuned a BYOL-pretrained 1D ResNet-50 encoder.
  • Introduced a clinically informed label hierarchy by merging electrophysiologically related arrhythmias into six clinically meaningful classes.
  • Used Focal Loss with inverse square-root class weighting to address class imbalance.
  • Applied targeted data augmentation only to minority classes using Gaussian noise and simulated baseline wander.
  • Used cosine learning rate scheduling with warm-up and early stopping.

To improve interpretability, I implemented 1D Grad-CAM, allowing users to visualize which portions of the ECG influenced the model's prediction.

The final model achieved:

  • Accuracy: 93.08%
  • Macro F1-score: 86.45%
  • Weighted F1-score: 92.96%
  • Macro AUROC: 98.32%

Challenges I ran into

One of the biggest challenges was the severe class imbalance in the dataset. Some arrhythmias contained thousands of ECG recordings, while others had only a few dozen examples.

Another challenge was deciding how to represent clinically similar arrhythmias. Instead of treating every rhythm as an independent class, I adopted a clinically informed label hierarchy by grouping electrophysiologically related arrhythmias together. This improved class balance while remaining practical for wearable single-lead ECG applications.

Implementing explainability was another challenge. While Grad-CAM is commonly used for image classification, adapting it for 1D ECG signals required implementing and validating a temporal Grad-CAM visualization.

Finally, balancing model performance with interpretability and deployment considerations required multiple rounds of experimentation with class weighting, data augmentation, and label design.

Accomplishments that I'm proud of

  • Successfully adapted a BYOL-pretrained 1D ResNet-50 for single-lead ECG classification.
  • Improved classification performance by introducing a clinically informed label hierarchy.
  • Built an explainable AI dashboard that combines prediction probabilities with 1D Grad-CAM visualizations.
  • Designed the project specifically for wearable single-lead ECG applications rather than relying on traditional 12-lead ECGs.
  • Improved minority-class performance using targeted ECG data augmentation without sacrificing overall model accuracy.

What I learned

This project taught me that achieving high accuracy alone is not sufficient for healthcare applications. A useful model should also be interpretable, clinically meaningful, and capable of expressing uncertainty when it is not confident.

I also learned that incorporating clinical knowledge into machine learning—such as grouping electrophysiologically related arrhythmias—can significantly improve both model performance and practical usability.

Finally, self-supervised learning proved to be an effective approach for learning robust ECG representations before supervised fine-tuning, especially when labeled medical data are limited.

What's next for Arrhythmia Classification with Explainable AI

The next step is to evaluate the model on external ECG datasets collected from different hospitals, populations, and recording devices to better understand how well it generalizes beyond the training data.

I also plan to evaluate the model on real-world single-lead ECG recordings collected from wearable devices such as smartwatches and portable ECG monitors. Improving uncertainty estimation, expanding the range of detectable arrhythmias, and further validating the Grad-CAM explanations with clinicians are important future directions.

Ultimately, I hope this work contributes toward making explainable and accessible ECG screening more practical outside traditional clinical settings while encouraging timely medical evaluation when abnormal heart rhythms are detected.

Built With

Share this project:

Updates