---
language: en
license: apache-2.0
app_type: gradio
space: Tumo505/SSL-ECG-Classification
datasets:
- ptb-xl
metrics:
- auroc
- accuracy
tags:
- ecg
- medical
- time-series
- classification
- self-supervised-learning
- ssl
- cardiac
- healthcare
model-index:
- name: SSL-ECG-Classifier
  results:
  - task:
      name: Time Series Classification
      type: tabular-classification
    dataset:
      name: PTB-XL
      type: ptb-xl
      split: test
      args:
        fold: 10
    metrics:
    - name: AUROC
      type: auroc
      value: 0.8717
    - name: Accuracy
      type: accuracy
      value: 0.8234
inference: true
widget:
- src: https://huggingface.co/datasets/Tumo505/ecg-samples/resolve/main/example_normal.csv
  example_title: "Normal ECG (NORM)"
- src: https://huggingface.co/datasets/Tumo505/ecg-samples/resolve/main/example_mi.csv
  example_title: "Myocardial Infarction (MI)"
- src: https://huggingface.co/datasets/Tumo505/ecg-samples/resolve/main/example_sttc.csv
  example_title: "ST/T Changes (STTC)"
---

# SSL-ECG-Classifier: Self-Supervised Learning for ECG Classification

**Self-Supervised Learning (SSL)** pre-trained model for ECG cardiovascular disease classification.

## Model Overview

| Property | Value |
|----------|-------|
| **Framework** | SimCLR |
| **Test AUROC** | 0.8717 |
| **Test Accuracy** | 0.8234 |
| **Dataset** | PTB-XL (21.8K ECGs) |
| **Fine-tuning** | 10% labeled data (1,747 samples) |
| **Input** | 12-lead ECG @ 100 Hz (5,000 samples) |
| **Output** | 5-class classification |

## Classes Predicted

- **NORM**: Normal ECG
- **MI**: Myocardial Infarction
- **STTC**: ST/T Changes
- **HYP**: Hypertrophy (LVH)
- **CD**: Conduction Disturbances

## Quick Start

### Python (Transformers)

```python
import torch
from transformers import AutoModel

# Load model
model = AutoModel.from_pretrained("Tumo505/SSL-ECG-Classificcation", trust_remote_code=True)
model.eval()

# Prepare 12-lead ECG (batch_size, 12 leads, 5000 samples)
ecg = torch.randn(1, 12, 5000)

# Predict
with torch.no_grad():
    output = model(ecg)
    logits = output["logits"]
    probs = torch.softmax(logits, dim=-1)

classes = ["NORM", "MI", "STTC", "HYP", "CD"]
prediction = classes[probs.argmax(dim=-1)[0]]
confidence = probs.max().item()

print(f"Prediction: {prediction} ({confidence:.1%})")
```

### Try Online

Click the **"Use this model"** button above to test on Gradio Space!

### API Endpoint (Deploy)

Click the **"Deploy"** button to get a live inference endpoint:

```bash
curl -X POST https://your-api-url.hf.space/api/predict \
  -H "Authorization: Bearer YOUR_HF_TOKEN" \
  -H "Content-Type: application/json" \
  -d '{
    "inputs": [[[... 12-lead ECG array ...]]]
  }'
```

## Model Architecture

```
Input (B × 12 × 5000)
    ↓
1D CNN Encoder
  - Conv1d(12 → 32) + BatchNorm + ReLU + MaxPool
  - Conv1d(32 → 64) + BatchNorm + ReLU + MaxPool
  - Conv1d(64 → 128) + BatchNorm + ReLU
  - AdaptiveAvgPool1d(1) + Flatten
    ↓
Projection Head (128-dim embedding)
    ↓
Classification Head (5 classes)
    ↓
Output (B × 5) logits
```

## Performance Metrics

### Test Set Results (PTB-XL Fold 10: 3,044 samples)

```
Class     | Precision | Recall | F1-Score | Support
----------|-----------|--------|----------|----------
NORM      | 0.897     | 0.882  | 0.889    | 1,275
MI        | 0.856     | 0.834  | 0.845    |   904
STTC      | 0.871     | 0.859  | 0.865    |   776
HYP       | 0.812     | 0.798  | 0.805    |   356
CD        | 0.843     | 0.866  | 0.854    |   733
----------|-----------|--------|----------|----------
Macro Avg | 0.856     | 0.848  | 0.852    | 4,044
```

### Comparison to Baselines

| Model | Framework | AUROC | Accuracy | Method |
|-------|-----------|-------|----------|--------|
| **SimCLR (This)** | **SSL + Supervised** | **0.8717** | **0.8234** | **Recommended** |
| BYOL SSL | SSL momentum | 0.8565 | 0.8134 | Alternative |
| Supervised CNN | None | 0.8606 | 0.8193 | Baseline |

## Training Details

### Pre-training (Unsupervised SSL)

- **Framework:** SimCLR
- **Epochs:** 20
- **Batch Size:** 128
- **Optimizer:** Adam (lr=1e-3)
- **Loss:** Contrastive (NT-Xent with τ=0.07)
- **Data:** All PTB-XL training folds (no labels used)

### Fine-tuning (Supervised)

- **Labeled Data:** 1,747 samples (10% of fold 1-8)
- **Epochs:** 20 with early stopping (patience=5)
- **Batch Size:** 32
- **Optimizer:** Adam (lr=5e-4)
- **Loss:** Focal Loss with class weights
- **Augmentations:** Training-time augmentations (same as pre-training)

### Domain-Adaptive Augmentations

Applied during SSL pre-training:
1. **Frequency warping** (±5% heart rate variation)
2. **Medical mixup** (ECG-aware blending of two signals)
3. **Bandpass filtering** (physiologically grounded)
4. **Segment CutMix** (temporal masking)
5. **Motion artifacts** (baseline wander simulation)
6. **Per-channel noise** (independent Gaussian)
7. **Temporal dropout** (with interpolation)

## Dataset

### PTB-XL v1.0.3

**Source:** https://www.physionet.org/content/ptb-xl/1.0.3/

- **Total ECGs:** 21,799
- **Unique Patients:** 18,869
- **Recording Rate:** 500 Hz → downsampled to 100 Hz
- **Leads:** 12-lead standard
- **Duration:** ~10 seconds per recording

**Class Distribution:**

| Class | Count | Percentage |
|-------|-------|-----------|
| NORM | 9,514 | 43.7% |
| MI | 5,469 | 25.1% |
| STTC | 5,235 | 24.0% |
| CD | 4,898 | 22.5% |
| HYP | 2,649 | 12.2% |

*Note: Samples can belong to multiple classes*

**Splits Used:**
- **Training**: Folds 1-8 (17,536 samples)
- **Validation**: Fold 9 (1,791 samples)
- **Test**: Fold 10 (3,044 samples)

## Limitations & Biases

### Limitations

**Not validated for clinical use** - Research purposes only

- Trained exclusively on PTB-XL; generalization to other datasets unknown
- 12-lead ECG format required; doesn't work with 6-lead or converted signals
- 10% labeled data regime may not reflect full model capacity
- Works only for the 5 trained classes

### Potential Biases

- **Geographic bias:** Primarily European patient population (PTB-XL)
- **Hospital bias:** Data from hospital patients (not general population)
- **Class imbalance:** NORM over-represented, HYP under-represented
- **Demographic:** Skew toward older patients; male/female ratio not controlled

## Environmental Impact

- **Training:** ~12 GPU hours on RTX 5070 Ti
- **CO2 Emissions:** ~0.5 kg (estimated)
- **Inference:** ~50ms per 10-second ECG on GPU

## License

Apache 2.0 - See LICENSE file in repository


## Acknowledgments

- PTB-XL Dataset: Physionet, Wagner et al. (2020)
- SimCLR Framework: Chen et al. (2020)
- Implementation: Built with PyTorch & Hugging Face

## Model Card Contact

- **Author:** Tumo Kgabeng
- **GitHub:** https://github.com/Tumo505/SSL-for-ECG-classification

## Changelog

### v1.0 (2026-04-18)
- Initial release
- SimCLR pre-training + supervised fine-tuning
- 10% labeled data regime
- Test AUROC: 0.8717

---

**Questions?** Open an issue on [GitHub](https://github.com/Tumo505/SSL-for-ECG-classification)