PyTorch ECG Augmentations

ECG time-series augmentations library for PyTorch. The focus of this repo is to:

• Provide many transformations in an easy-to-use Python interface.
• Easily control stochastic / sequential transformations.
• Make every transformation differentiable through torch.nn.Module.
• Optimize ECG transformations for CPU and GPU devices.

This repo supports stochastic transformations as used often in self-supervised and semi-supervised learning methods. One can apply a single stochastic augmentation or create as many stochastically transformed ECG examples from a single interface. This package follows the conventions set out by torchaudio-augmentations, with an ECG sample defined as a tensor of [lead, time], or a batched representation [batch, lead, time]. Each individual augmentation can be initialized on its own, or be wrapped around a RandomApply interface which will apply the augmentation with probability p.

Note: Current version runs on beta for multiple-lead ECG, we appreciate any contributions.

Installation

Just clone the repository and install the library:

git clone https://github.com/klean2050/ecg-augmentations
pip install -e ecg-augmentations

Usage

One can define a single or several augmentations, which are applied sequentially to an ECG sample:

from ecg_augmentations import *

# 1 lead, 100 samples
ecg = torch.load("tests/some_ecg.pt")

num_samples = ecg.shape[-1] // 2
transforms = [
    RandomCrop(n_samples=num_samples),
    RandomApply([PRMask()], p=0.4),
    RandomApply([QRSMask()], p=0.4),
    RandomApply([Scale()], p=0.5),
    RandomApply([Permute()], p=0.6),
    GaussianNoise(max_snr=0.005),
    RandomApply([Invert()], p=0.2),
    RandomApply([Reverse()], p=0.2),
]

One can also define a stochastic augmentation on multiple transformations:

transforms = [
    RandomCrop(n_samples=num_samples),
    RandomApply([PRMask(), QRSMask()], p=0.8),
    RandomApply([Invert(), Reverse()], p=0.5)
]

One can return either one ...

transform = Compose(transforms=transforms)
transformed_ecg =  transform(ecg)
>> transformed_ecg.shape = [num_leads, num_samples]

... or many versions of the same example:

transform = ComposeMany(transforms=transforms, num_augmented_samples=2)
transformed_ecg =  transform(ecg)
>> transformed_ecg.shape = [2, num_leads, num_samples]

Citation

@software{avramidis_kleanthis_2023_8053043,
    author = {Avramidis, Kleanthis},
    title = {ecg-augmentations},
    year = 2023,
    publisher = {Zenodo},
    version = {v1.0.0-beta},
    doi = {10.5281/zenodo.8053043},
    url = {https://doi.org/10.5281/zenodo.8053043}
}