DiSR-ECG

Residual Shifting Conditional Diffusion for Robust ECG Super-Resolution

Diffusion-based super-resolution for 12-lead ECG signals. Reconstructs clean 500Hz high-resolution ECG from noisy low-resolution (50Hz) input, using a ResShift-style conditional diffusion model with Mamba-based temporal modeling.

DiSR-ECG achieves state-of-the-art performance on PTB-XL and generalizes significantly better than prior discriminative methods (SRECG, DCAE-SR, MSECG) under distribution shift, evaluated zero-shot on the Chapman-Shaoxing database.

Method

Fig. 1. Overview of residual shifting diffusion.

Fig. 2. DiSR-ECG model architecture.

Datasets

Dataset	Purpose	Config field	Setup
PTB-XL	Training / in-domain evaluation	train.dataset_dir	Download and point to the extracted folder
MIT-BIH Noise Stress Test Database (NSTDB)	Noise augmentation (bw/ma/em)	train.noise_dir	Download, point to the extracted folder, then move info.csv into that same folder (read as noise_dir/info.csv)
Chapman-Shaoxing (12-lead ECG arrhythmia database)	Cross-dataset zero-shot evaluation	train.ecg_dataset_dir	Download and point to the extracted folder

Data preprocessing is handled in utils/dataloader.py:

• Each 12-lead, 10s record at 500Hz is bandpass filtered (Butterworth, 1–45Hz), forming the HR target
• The HR signal is downsampled by a factor of 10 to 50Hz, then linearly interpolated back to 500Hz to form the LR condition
• For PTB-XL dataset, with probability 0.5 a noise segment (bw/em/ma) from NSTDB is injected into the LR signal at a sampled target SNR; Chapman-Shaoxing is evaluated without noise injection

Environment

Create a dedicated conda env (Python 3.9, CUDA 12.1):

conda create -n disr python=3.9 -y
conda activate disr
pip install torch==2.2.2 torchaudio==2.2.2 --index-url https://download.pytorch.org/whl/cu121
pip install -r requirements.txt

All package versions are pinned in requirements.txt. Note that mamba-ssm compiles CUDA kernels on install, so torch must already be installed first (as above) rather than resolved as part of requirements.txt.

Quickstart

All paths and hyperparameters live in config/base.yaml:

• train.dataset_dir: PTB-XL root
• train.noise_dir: NSTDB noise root
• train.ecg_dataset_dir: additional ECG test set root
• train.pth_path: default checkpoint path, used by train.py (save) and inference.py (load) when --pth_path isn't passed
• train.input_training_file / input_validation_file / input_testing_file / input_ecg_testing_file: index files under dataset_index/
• train.mode: which index file inference.py evaluates on: test for PTB-XL (in-domain), test_ecg for Chapman-Shaoxing (cross-dataset zero-shot)

First run: set train.regenerate_index: True in config/base.yaml, run train.py once to generate dataset_index/*.txt, then set it back to False.

Download the pretrained checkpoint from Releases:

mkdir -p model_pth
wget -O model_pth/DiSR-ECG.pth https://github.com/FLYIH/DiSR-ECG/releases/download/v1.0/DiSR-ECG.pth

Then train from scratch, or run inference/evaluation directly with the pretrained checkpoint:

python train.py --gpu 0 --pth_path ./model_pth/DiSR-ECG.pth

python inference.py --gpu 0 --pth_path ./model_pth/DiSR-ECG.pth --save_path ./sr_signal/

python evaluate.py --gt_index ./dataset_index/test.txt --pd_path ./sr_signal/ --out_csv ./DiSRECG.csv

All flags are optional and fall back to config/base.yaml defaults if omitted.

Or run everything via the provided scripts (from anywhere in the repo):

./scripts/train.sh   # python train.py
./scripts/eval.sh    # python inference.py + python evaluate.py

Results

In-domain (PTB-XL, Fold 10 test set)

Method	MSE (×10⁻³) ↓	CoS (×10⁻²) ↑	SNR (dB) ↑	MAD ↓
LI	7.477	90.211	8.592	0.877
SRECG	0.422	99.367	19.751	0.371
DCAE-SR	4.461	97.560	12.591	0.962
MSECG	0.184	99.745	24.037	0.221
DiSR-ECG	0.181	99.754	24.092	0.216

Cross-dataset, zero-shot (Chapman-Shaoxing)

Method	MSE (×10⁻³) ↓	CoS (×10⁻²) ↑	SNR (dB) ↑	MAD ↓
LI	4.289	95.114	10.549	1.111
SRECG	1.227	98.729	17.535	0.659
DCAE-SR	9.017	96.168	10.829	1.410
MSECG	0.776	99.084	21.332	0.487
DiSR-ECG	0.642	99.376	21.939	0.447