ECGEXPFormer — Explainable ECG-based Atrial Fibrillation Detection

A deep-learning pipeline for detecting Atrial Fibrillation (AF) from ECG recordings, with SHAP explanations that reveal which parts of the signal drive each prediction. Two architectures are compared — a Transformer on raw 1-D ECG patches and a CNN on ECG spectrograms — over a shared, reproducible preprocessing → training → explanation workflow.

Developed as a Bachelor's thesis project. The original dataset is private clinical data and is not distributed here; see docs/DATA_FORMAT.md to plug in your own ECG source.

---

Table of contents

• Highlights
• Repository structure
• Installation
• The pipeline
• Configuration
• Models
• Explainability
• Reproducibility notes
• Citation
• References
• License

---

Highlights

• End-to-end pipeline: raw WFDB records → filtered/segmented/labeled patches → memory-bounded global shuffle → model training → SHAP explanations.
• Two models, one trainer: the Transformer and the CNN share a single training loop; architecture and data are the only differences.
• Config-driven & reproducible: every run is described by a YAML file and a fixed seed; the resolved config is saved next to the checkpoints.
• Explainable: per-time-sample SHAP attributions are rendered directly on the ECG waveform (red = pushes toward AF, blue = pushes away).

---

Repository structure

ECGEXPFormer/
├── configs/                       # YAML experiment configs
│   ├── transformer.yaml
│   └── cnn.yaml
├── docs/
│   └── DATA_FORMAT.md             # expected input & intermediate formats
├── src/ecgexpformer/
│   ├── config.py                  # single Config dataclass (+ YAML I/O)
│   ├── utils.py                   # seeding & device selection
│   ├── data/
│   │   └── dataset.py             # NpzArrayDataset, ECGDataset, ECGDatasetSpec
│   ├── models/
│   │   ├── transformer.py         # ECGFormerForClassification (+ spectrogram variant)
│   │   └── cnn.py                 # ECGCNNForClassification, ConvBlock
│   ├── preprocessing/
│   │   ├── preprocess.py          # WFDB -> labeled .npz patches
│   │   └── shuffle.py             # two-pass global shuffle
│   ├── training/
│   │   ├── metrics.py             # running sensitivity/specificity
│   │   ├── trainer.py             # shared train/eval loop
│   │   └── train.py               # CLI entry point (transformer | cnn)
│   └── explain/
│       └── shap_explain.py        # SHAP attributions on the ECG waveform
├── pyproject.toml                 # packaging + console scripts
├── requirements.txt
├── LICENSE
└── CITATION.cff

---

Installation

git clone https://github.com/manudella/ECGEXPFormer.git
cd ECGEXPFormer

python -m venv .venv
source .venv/bin/activate          # Windows: .venv\Scripts\activate

# Editable install with all dependencies and the ecg-* console commands
pip install -e .
# (or: pip install -r requirements.txt)

This registers four commands: ecg-preprocess, ecg-shuffle, ecg-train, ecg-explain. Each is also runnable as python -m ecgexpformer.<module>.

---

The pipeline

 raw WFDB (.dat/.hea/.atr)
        │  ecg-preprocess
        ▼
 preprocessed_data_*/  (per-record .npz: patches/spectrograms + labels)
        │  ecg-shuffle
        ▼
 shuffled_preprocessed_data_*/   ──►  ecg-train  ──►  checkpoints_*/  + confusion matrices
                                                          │  ecg-explain
                                                          ▼
                                                     shap_images/

1. Preprocess raw recordings

ecg-preprocess --data-dir ./data/ATA/ \
               --train-out ./preprocessed_data_ata_ecg_1c/ \
               --test-out  ./preprocessed_data_ata_ecg_test_1c/

Filters the ECG (baseline + 0.5–40 Hz bandpass), merges the two leads, segments into 10 s patches, and labels each patch (0 = normal, 1 = AF, 2 = other) from the WFDB annotations. See docs/DATA_FORMAT.md for details.

2. Shuffle globally (memory-bounded)

ecg-shuffle --input-dir  ./preprocessed_data_ata_ecg_1c/ \
            --temp-dir   ./temp_dir_1c/ \
            --output-dir ./shuffled_preprocessed_data_ata_ecg_1c/ \
            --data-key   patches

3. Train a model

# Transformer on 1-D ECG patches
ecg-train --model transformer --config configs/transformer.yaml

# CNN on ECG spectrograms
ecg-train --model cnn --config configs/cnn.yaml

Checkpoints and confusion matrices are written to the directories named in the config. The resolved config is saved to <checkpoint_dir>/resolved_config.yaml.

4. Explain predictions with SHAP

ecg-explain --config configs/transformer.yaml \
            --checkpoint ./checkpoints_ata/ecgformer_ata_model_epoch_30.pth \
            --data-dir   ./shuffled_preprocessed_data_ata_ecg_1c/ \
            --output-dir ./shap_images/

Saves ECG plots grouped by (predicted, true) class, with each waveform colored by its per-sample SHAP attribution.

---

Configuration

All knobs live in a single Config dataclass, populated from a YAML file. A few common overrides are exposed on the CLI (--epochs, --train-dir, --test-dir); anything else is edited in the YAML. Because the run is fully described by its config plus a fixed seed, results are reproducible and the exact settings travel with the checkpoints.

---

Models

Model	Input	Idea
ECGFormerForClassification	1-D ECG patch (B, 2000)	A 1-D Vision-Transformer: a strided Conv1d patch embedding splits the signal into a sequence of tokens (e.g. 50 tokens of 40 samples each), a learnable [CLS] token and positional embeddings are added, pre-norm Transformer encoder layers model temporal relationships, and classification reads from the [CLS] token (or mean/max pooling).
ECGCNNForClassification	2-D spectrogram (B, H, W)	Stack of residual Conv-BN-ReLU blocks with strided downsampling, global average/max pooling, and a linear head.

A spectrogram Transformer variant (ECGFormerForSpecClassification) is also included; it tokenizes each spectrogram time-frame.

Why the tokenizer matters

An earlier version projected the entire 2000-sample patch into a single token, which left self-attention with a sequence of length 1 — effectively an MLP with no temporal modeling. The Conv1d patch embedding fixes this "single- token collapse": each token now summarizes a short local window, and attention operates over the resulting sequence. Tokenization is controlled by patch_size / patch_stride, and pooling by pooling (cls / mean / max). Training adds a OneCycle LR schedule and gradient clipping for stability (both configurable).

Note: because the architecture changed, checkpoints trained with the old single-token model are not loadable here — retrain from the current code.

---

Explainability

Explanations use SHAP's DeepExplainer on the trained model. For each patch, the attribution vector for the predicted class is mapped onto the ECG waveform with a diverging colormap: red segments push the prediction toward AF, blue segments push it away. This makes it possible to see whether the model attends to physiologically meaningful regions of the signal.

---

Reproducibility notes

• Set once via --seed / config.seed; seeds Python, NumPy and PyTorch.
• No data is committed — raw, preprocessed, shuffled, checkpoints and images are all git-ignored.
• The Transformer tokenizes each patch into a sequence (see Why the tokenizer matters); patch_size, pooling, scheduler and max_grad_norm are all set in the YAML config.

---

Citation

If you use this code, please cite it (see CITATION.cff):

@software{dellabona_ecgexpformer,
  author = {Dellabona, Manuel},
  title  = {ECGEXPFormer: Explainable ECG-based Atrial Fibrillation Detection},
  url    = {https://github.com/manudella/ECGEXPFormer},
  license = {MIT}
}

---

References

• AF Detection & ECG Classification
  • Rohr et al., 2022
  • Zihlmann et al., 2017
  • Che et al., 2021 (BMC Medical Informatics and Decision Making)
• Transformers in Medical Signals
  • Hong et al., 2019 (IJCAI)
  • Zhang et al., 2021
• SHAP
  • Lundberg and Lee, 2017

---

License

Released under the MIT License. Contributions and issues are welcome.