Initial release

README.md

@@ -0,0 +1,189 @@
+---
+library_name: pytorch
+tags:
+- ecg
+- arrhythmia
+- rhythm-classification
+- ltaf
+- physionet
+- 1d-resnet
+license: mit
+datasets:
+- physionet/ltafdb
+---
+
+# LTAF ECG Rhythm Classifier — RhythmResNet1D + TTA
+
+A from-scratch 1D-ResNet trained on PhysioNet's
+[Long-Term Atrial Fibrillation (LTAF)](https://physionet.org/content/ltafdb/)
+database for **6-class rhythm classification** on two-lead 128 Hz ECG.
+
+| Metric | Single-window | **+ 7-view TTA (recommended)** |
+|---|---:|---:|
+| Test accuracy | 0.636 | **0.684** |
+| Test balanced accuracy | 0.740 | **0.778** |
+| **Test macro F1** | **0.614** | **0.656** |
+
+vs. frozen Chronos-2 + MLP baseline on the same 6-class subset:
+test macro F1 = 0.299 — i.e. **+36 pp / 2.2× the F1**.
+
+Per-class F1 (TTA-7): NSR 0.76, AFIB 0.62, SBR 0.82, AB 0.77, SVTA 0.15, B 0.82.
+
+## Classes
+
+| Code | Expansion |
+|------|-----------|
+| NSR  | Normal sinus rhythm |
+| AFIB | Atrial fibrillation |
+| SBR  | Sinus bradycardia (<60 bpm, sinus origin) |
+| AB   | Atrial bigeminy (every other beat is an APC) |
+| SVTA | Supraventricular tachyarrhythmia (≥3 consec SV ectopics @ >100 bpm) |
+| B    | Ventricular bigeminy (every other beat is a PVC) |
+
+`VT`, `T`, and `IVR` are excluded — their LTAF test supports (31, 26, 1) are too small for stable F1 estimation.
+
+## Quickstart
+
+```bash
+pip install torch huggingface_hub numpy
+```
+
+```python
+import numpy as np
+import torch
+from huggingface_hub import hf_hub_download
+from model import RhythmResNet1D, RHYTHM_CLASS_NAMES
+
+# Download checkpoint + model code from HF
+ckpt = hf_hub_download("rmxjck/ltaf-ecg-rhythm-classifier", "best_classifier.pt")
+model = RhythmResNet1D.load(ckpt, device="cuda")
+model.eval()
+
+# Input: (B, 2, 1280) — 10 s @ 128 Hz, 2 leads, per-channel z-scored.
+x = torch.randn(1, 2, 1280).cuda()  # replace with real ECG
+with torch.no_grad():
+    logits = model(x)
+    pred_idx = logits.argmax(-1).item()
+print(model.class_names[pred_idx])
+```
+
+For best results, use the **7-view TTA** wrapper in `inference.py`
+(averages softmax across 7 random window-start offsets — adds ~4 pp F1
+at the cost of 7× inference compute).
+
+```bash
+python inference.py
+```
+
+## Architecture
+
+`RhythmResNet1D(num_classes=6, n_channels=2, base_channels=64,
+blocks_per_stage=2)`:
+
+- **Stem:** Conv1d(2, 64, k=15, stride=2) → BN → ReLU → MaxPool(2).
+- **4 ResNet stages × 2 basic blocks** (Conv1d k=7, BN, ReLU, Dropout, +skip).
+  Channels: 64 → 128 → 256 → 512. Time downsamples 2× at the start of each
+  stage past the first.
+- **Head:** AdaptiveAvgPool1d → Linear(512 → 128) → ReLU → Dropout(0.2)
+  → Linear(128 → 6).
+- **Total parameters:** 8,794,246.
+
+## Input format
+
+- `(B, 2, 1280)` float32
+- 2-lead ECG at **128 Hz** (LTAF leads `ECG1`, `ECG2`)
+- 10 s window
+- Per-channel z-scored: `(x - x.mean(axis=-1)) / x.std(axis=-1)`
+
+## Test-time augmentation (TTA)
+
+Pass a longer signal slice (≥1280 samples) to `predict_tta()` and it
+samples 7 random 10 s windows, averages the softmax outputs, then
+argmaxes. Why it helps: training uses random window-start sampling
+within each rhythm bout, so the model learns to be invariant to that
+shift. At eval time, taking multiple shifts and averaging cancels the
+position-specific noise. **+4.2 pp test macro F1, no retraining.**
+
+```python
+# (2, 30*128) signal, 30 s long
+cls, prob, full_probs = predict_tta(model, long_signal, n_views=7, device="cuda")
+```
+
+## Training recipe
+
+```bash
+.venv/bin/python scripts/train_ecg_rhythm_scratch.py \
+    --arch resnet1d --window-sizes 10 \
+    --epochs 30 --batch-size 64 --lr 5e-4 \
+    --base-channels 64 \
+    --use-val-as-train \
+    --classes NSR AFIB SBR AB SVTA B \
+    --output-dir results/ecg_classifier/sweep/c6_resnet1d_w10_e30_wide
+```
+
+- Dataset: LTAF train+val combined (75 records). 8 records held out for
+  early stopping. Test (9 records, 3,716 windows) untouched.
+- Loss: weighted cross-entropy with sqrt-dampened inverse-frequency
+  class weights (cap 10), label smoothing 0.1.
+- Cosine LR schedule from 5e-4 → 0 over 30 epochs. AdamW (wd 1e-4).
+- Best checkpoint by held-out macro F1.
+- Training time on a single H100 80GB: **~6 minutes**.
+
+Source repo: `scripts/train_ecg_rhythm_scratch.py` and
+`src/models/ts_llm/ecg_rhythm_scratch.py` in
+[rmxjck/TSLM-Arena](https://github.com/rmxjck/TSLM-Arena).
+
+## Test set details
+
+LTAF held-out split (deterministic seed 42, record-level): 9 records
+(`100, 104, 105, 11, 200, 32, 48, 49, 68`), 3,716 windows.
+
+Confusion matrix (rows = true, cols = pred), with TTA:
+
+|       | NSR | AFIB | SBR | AB  | SVTA | B   |
+|-------|----:|-----:|----:|----:|-----:|----:|
+| NSR   | 1109 | 286 | 95  | 114 | 185  | 35  |
+| AFIB  | 189  | 628 | 25  | 29  | 294  | 14  |
+| SBR   | 26   | 0   | 279 | 0   | 0    | 0   |
+| AB    | 9    | 13  | 0   | 225 | 3    | 0   |
+| SVTA  | 9    | 14  | 0   | 3   | 34   | 0   |
+| B     | 4    | 0   | 0   | 1   | 3    | 90  |
+
+Per-class supports: NSR 1824, AFIB 1179, SBR 305, AB 250, SVTA 60, B 98.
+
+## What was tried and didn't help
+
+This model was the best of 30+ experiments. What did *not* improve over
+this baseline:
+
+- HRV side-channel input (8-dim RR-derived features fused with CNN trunk):
+  hurts F1 by 3-8 pp because the CNN already extracts equivalent
+  information from raw QRS timing.
+- Cross-corpus augmentation (MIT-BIH AFDB added to training): hurts
+  AFIB F1 by 14 pp because AFDB's clean AFIB blocks bias the model
+  toward over-calling AFIB on LTAF's paroxysmal transitions.
+- Wider models (96-channel, 12 M params): overfits.
+- Longer training (50 epochs): overfits.
+- Multi-model soft-voting ensembles: members make correlated errors.
+- Focal loss: matches CE within noise.
+- Multi-scale training (5 / 10 / 30 s windows): underperforms 10 s alone.
+- Bigger external models (torchecg ResNet-50 51.9 M, Stanford 27 M):
+  underperform a 2.2 M home-rolled ResNet1D at 12 epochs.
+
+## Not for clinical use
+
+Research artifact only. **Not FDA-cleared.** Not suitable for triage,
+diagnosis, or any patient-facing application. Uses the LTAF benchmark
+which has known label noise from its original PhysioNet curation.
+
+## Citation
+
+```bibtex
+@misc{petrutiu2008ltafdb,
+  title         = {Abrupt Changes in Fibrillatory Wave Characteristics at the Termination of Paroxysmal Atrial Fibrillation in Humans},
+  author        = {Petrutiu, Simona and Sahakian, Alan V. and Swiryn, Steven},
+  year          = {2008},
+  howpublished  = {PhysioNet},
+  url           = {https://physionet.org/content/ltafdb/}
+}
+```

best_classifier.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e52d7453256a051cbcc5516a1462d9869a5baf8e2e5caeb2d2a0e5b69fa3e961
+size 35256343

inference.py

@@ -0,0 +1,119 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: MIT
+"""Inference example for the LTAF ECG rhythm classifier.
+
+Two modes:
+- Single-window: pass a (B, 2, 1280) z-scored 10 s @ 128 Hz tensor.
+- TTA-7 (recommended, +4 pp F1): pass a longer signal slice and the
+  function will pull 7 random 10 s windows from it and soft-vote.
+
+Usage:
+    .venv/bin/python inference.py
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Tuple
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from huggingface_hub import hf_hub_download
+
+from model import RHYTHM_CLASS_NAMES, RhythmResNet1D
+
+
+WINDOW_SECONDS = 10
+SOURCE_HZ = 128
+WINDOW_SAMPLES = WINDOW_SECONDS * SOURCE_HZ  # 1280
+
+
+def load_model(device: str = "cpu") -> RhythmResNet1D:
+    """Download the checkpoint from HF and load it."""
+    ckpt_path = hf_hub_download(
+        "rmxjck/ltaf-ecg-rhythm-classifier",
+        "best_classifier.pt",
+    )
+    return RhythmResNet1D.load(ckpt_path, device=device)
+
+
+def zscore(window: np.ndarray) -> np.ndarray:
+    """Per-channel z-score a (C, L) array."""
+    mean = window.mean(axis=-1, keepdims=True)
+    std = window.std(axis=-1, keepdims=True)
+    return ((window - mean) / (std + 1e-6)).astype(np.float32, copy=False)
+
+
+def predict_single(
+    model: RhythmResNet1D,
+    window: np.ndarray,
+    device: str = "cpu",
+) -> Tuple[str, float]:
+    """Predict on one (2, 1280) z-scored window. Returns (class_name, prob)."""
+    if window.shape != (2, WINDOW_SAMPLES):
+        raise ValueError(f"Expected (2, {WINDOW_SAMPLES}), got {window.shape}")
+    x = torch.from_numpy(window).float().unsqueeze(0).to(device)
+    with torch.no_grad():
+        probs = F.softmax(model(x), dim=-1)[0]
+    idx = int(probs.argmax().item())
+    return model.class_names[idx], float(probs[idx].item())
+
+
+def predict_tta(
+    model: RhythmResNet1D,
+    long_signal: np.ndarray,
+    n_views: int = 7,
+    device: str = "cpu",
+    seed: int = 42,
+) -> Tuple[str, float, np.ndarray]:
+    """TTA-soft-voting prediction over a longer (2, L) signal.
+
+    Samples ``n_views`` random 10 s windows from ``long_signal`` (L >= 1280),
+    z-scores each independently, runs them through the model, and averages
+    the softmax probabilities.
+
+    Returns (class_name, prob, full_probs) where full_probs is shape (6,).
+    """
+    n_ch, n_samples = long_signal.shape
+    if n_ch != 2:
+        raise ValueError(f"Expected 2-channel signal, got {n_ch}")
+    if n_samples < WINDOW_SAMPLES:
+        raise ValueError(f"Need at least {WINDOW_SAMPLES} samples, got {n_samples}")
+    rng = np.random.default_rng(seed)
+    starts = rng.integers(0, n_samples - WINDOW_SAMPLES + 1, size=n_views)
+    accum = torch.zeros(model.num_classes, device=device)
+    for s in starts:
+        window = zscore(long_signal[:, s:s + WINDOW_SAMPLES])
+        x = torch.from_numpy(window).float().unsqueeze(0).to(device)
+        with torch.no_grad():
+            probs = F.softmax(model(x), dim=-1)[0]
+        accum += probs
+    probs_avg = accum / n_views
+    idx = int(probs_avg.argmax().item())
+    return model.class_names[idx], float(probs_avg[idx].item()), probs_avg.cpu().numpy()
+
+
+def demo():
+    print("Loading model from HF...")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = load_model(device)
+    print(f"Loaded {model.__class__.__name__} on {device}")
+    print(f"Classes: {model.class_names}")
+    print(f"Params: {sum(p.numel() for p in model.parameters()):,}")
+
+    # Synthetic example: random noise (will get garbage prediction).
+    print("\n--- single-window demo (random input) ---")
+    fake_window = zscore(np.random.randn(2, WINDOW_SAMPLES).astype(np.float32))
+    cls, prob = predict_single(model, fake_window, device=device)
+    print(f"prediction: {cls} ({prob:.1%})")
+
+    print("\n--- TTA-7 demo (random 30 s input) ---")
+    fake_long = np.random.randn(2, 30 * SOURCE_HZ).astype(np.float32)
+    cls, prob, full = predict_tta(model, fake_long, n_views=7, device=device)
+    print(f"prediction: {cls} ({prob:.1%})")
+    print(f"all class probs: {dict(zip(model.class_names, [round(p, 3) for p in full.tolist()]))}")
+
+
+if __name__ == "__main__":
+    demo()

model.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: MIT
+"""Self-contained RhythmResNet1D for LTAF rhythm classification.
+
+Vendored from rmxjck/TSLM-Arena (src/models/ts_llm/ecg_rhythm_scratch.py)
+so the model can be loaded with no external project imports.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+RHYTHM_CLASS_NAMES = ["NSR", "AFIB", "SBR", "AB", "SVTA", "B"]
+
+
+class _BasicBlock1D(nn.Module):
+    """Two-conv residual block with optional stride-2 downsample."""
+
+    def __init__(self, in_c: int, out_c: int, kernel: int = 7, stride: int = 1,
+                 dropout: float = 0.1):
+        super().__init__()
+        pad = kernel // 2
+        self.conv1 = nn.Conv1d(in_c, out_c, kernel_size=kernel, stride=stride,
+                               padding=pad, bias=False)
+        self.bn1 = nn.BatchNorm1d(out_c)
+        self.conv2 = nn.Conv1d(out_c, out_c, kernel_size=kernel, stride=1,
+                               padding=pad, bias=False)
+        self.bn2 = nn.BatchNorm1d(out_c)
+        self.drop = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+        if stride != 1 or in_c != out_c:
+            self.proj = nn.Sequential(
+                nn.Conv1d(in_c, out_c, kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm1d(out_c),
+            )
+        else:
+            self.proj = nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        identity = self.proj(x)
+        h = F.relu(self.bn1(self.conv1(x)), inplace=True)
+        h = self.drop(h)
+        h = self.bn2(self.conv2(h))
+        return F.relu(h + identity, inplace=True)
+
+
+class RhythmResNet1D(nn.Module):
+    """1D ResNet — stem + 4 stages, each stage halves time and doubles channels."""
+
+    def __init__(
+        self,
+        num_classes: int = 6,
+        class_names: List[str] = RHYTHM_CLASS_NAMES,
+        n_channels: int = 2,
+        base_channels: int = 64,
+        blocks_per_stage: int = 2,
+        stem_kernel: int = 15,
+        block_kernel: int = 7,
+        dropout: float = 0.2,
+    ):
+        super().__init__()
+        assert len(class_names) == num_classes
+        self.num_classes = num_classes
+        self.class_names = list(class_names)
+        self.n_channels = n_channels
+        self.base_channels = base_channels
+        self.blocks_per_stage = blocks_per_stage
+
+        self.stem = nn.Sequential(
+            nn.Conv1d(n_channels, base_channels, kernel_size=stem_kernel,
+                      stride=2, padding=stem_kernel // 2, bias=False),
+            nn.BatchNorm1d(base_channels),
+            nn.ReLU(inplace=True),
+            nn.MaxPool1d(2),
+        )
+        stages = []
+        in_c = base_channels
+        out_c = base_channels
+        for s in range(4):
+            for b in range(blocks_per_stage):
+                stride = 2 if (b == 0 and s > 0) else 1
+                stages.append(_BasicBlock1D(in_c, out_c, kernel=block_kernel,
+                                            stride=stride, dropout=dropout))
+                in_c = out_c
+            out_c = min(out_c * 2, 512)
+        self.stages = nn.Sequential(*stages)
+        self.pool = nn.AdaptiveAvgPool1d(1)
+        self.head = nn.Sequential(
+            nn.Linear(in_c, 128),
+            nn.ReLU(inplace=True),
+            nn.Dropout(dropout),
+            nn.Linear(128, num_classes),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = self.stem(x)
+        h = self.stages(h)
+        feat = self.pool(h).squeeze(-1)
+        return self.head(feat)
+
+    @classmethod
+    def load(cls, path: str | Path, device: str = "cpu") -> "RhythmResNet1D":
+        ckpt = torch.load(path, map_location=device, weights_only=False)
+        model = cls(
+            num_classes=ckpt["num_classes"], class_names=ckpt["class_names"],
+            n_channels=ckpt["n_channels"],
+            base_channels=ckpt.get("base_channels", 64),
+            blocks_per_stage=ckpt.get("blocks_per_stage", 2),
+        )
+        model.load_state_dict(ckpt["state_dict"])
+        model.to(device).eval()
+        return model

requirements.txt

@@ -0,0 +1,3 @@
+torch>=2.0
+huggingface_hub>=0.20
+numpy>=1.24