FOUND-AF: Benchmarking ECG Foundation Models for Atrial Fibrillation Detection

This repository contains the code and aggregate results needed to reproduce the benchmark described below.

Atrial fibrillation (AF) is the most prevalent sustained cardiac arrhythmia and a leading cause of stroke, motivating reliable automated detection from the electrocardiogram (ECG). This study presents a unified, leakage-free benchmark of nine publicly released ECG foundation models from five model families (HuBERT-ECG, CLEF, ST-MEM, ECG-JEPA, ECGFounder) on four PhysioNet datasets — AFDB, CinC2017, CPSC2021, LTAFDB — totalling 1,954,594 ten-second windows from 7,308 recordings. All models are used as frozen feature extractors under an identical preprocessing pipeline; the resulting embeddings are classified with a fixed-configuration XGBoost classifier evaluated through 5-fold cross-validation grouped by recording. ECGFounder achieved the best performance on every dataset ($F_1$ = 97.87% / 92.77% / 99.50% / 95.68% on AFDB / CinC2017 / CPSC2021 / LTAFDB respectively), with the advantage confirmed statistically significant by bootstrap analysis, and was also the Pareto-preferred model in the accuracy–parameters–latency trade-off.

.
├── code/ Benchmark notebooks (one per dataset) — data loading,
│ model loading, embedding extraction, CV, evaluation
│ ├── 00_environment_diagnostics.ipynb
│ ├── AFDB.ipynb
│ ├── CINC2017.ipynb
│ ├── CPSC2021.ipynb
│ └── LTAFDB.ipynb
├── data/ Dataset documentation (sources, preprocessing,
│ README.md how to regenerate the processed files — raw/processed
│ data itself is not tracked here, see below)
├── results/ Aggregate benchmark outputs per dataset (metrics,
│ README.md efficiency, t-SNE coordinates, summary charts)
├── requirements.txt
├── CITATION.cff
├── LICENSE
└── README.md (this file)

1. Get the data. The four ECG databases are public but not redistributed here (see data/README.md for download links and the exact preprocessing pipeline — 10 s windows, 5 s overlap, per-window z-score normalisation).

2. Install dependencies:

   pip install -r requirements.txt

3. Run a benchmark notebook, e.g. code/AFDB.ipynb. Each notebook:

   • clones the original model repos and downloads pretrained checkpoints (HuBERT-ECG, CLEF, ST-MEM, ECG-JEPA, ECGFounder) from their official sources (GitHub / Hugging Face Hub / Zenodo / Google Drive mirrors),
   • loads the dataset's processed Parquet file,
   • extracts frozen embeddings from all 9 models,
   • runs 5-fold recording-level cross-validation with a fixed-config XGBoost classifier,
   • writes per-model results, efficiency, and t-SNE coordinates to results/<dataset>/.

   The notebooks were originally run on Google Colab (they drive.mount() and read/write under a Drive project folder) — update the PROJECT_DIR / PARQUET_PATH variables in each notebook's first cell to point at a local or different-cloud path if running elsewhere.

See results/README.md for the full breakdown. Headline numbers (best model, ECGFounder, mean F1 across 5 folds):

To keep the repository a reasonable size, the following are excluded (see .gitignore) and documented instead so they can be regenerated:

• Raw and processed ECG data (~3.2 GB of Parquet files) — public sources and preprocessing steps in data/README.md.
• Pretrained model checkpoints — downloaded automatically by the notebooks from the original authors' sources.
• Per-window prediction dumps (up to ~300 MB per dataset) — the aggregate metrics in results/ are kept; raw predictions are reproducible by re-running the notebooks.

See CITATION.cff. Please contact the corresponding author for citation details.

Code is released under the MIT License. The underlying ECG databases and pretrained model checkpoints are third-party resources under their own licenses (see data/README.md).