StretchCompress#

class torch_ecg.augmenters.StretchCompress(ratio: Real = 6, prob: float = 0.5, inplace: bool = True, **kwargs: Any)[source]#

Bases: Augmenter

Stretch-or-compress augmenter on ECG tensors.

Rescaling the ECGs by a factor sampled from a normal distribution along the time axis.

Parameters:

  • ratio (numbers.Real, default 6) – Mean ratio of the stretch or compress. If it is in the interval[1, 100], then it will be transformed to [0, 1]. The ratio of one batch element is sampled from a normal distribution.

  • prob (float, default 0.5) – Probability of the augmenter to be applied.

  • inplace (bool, default True) – If True, the input ECGs will be modified inplace.

  • kwargs (dict, optional) – Additional keyword arguments.

Example

sc = StretchCompress()
sig = torch.randn((32, 12, 5000))
labels = torch.randint(0, 2, (32, 5000, 26))
label = torch.randint(0, 2, (32, 26), dtype=torch.float32)
mask = torch.randint(0, 2, (32, 5000, 3), dtype=torch.float32)
sig, label, mask = sc(sig, label, mask)
extra_repr_keys() → List[str][source]#

Extra keys for __repr__() and __str__().

forward(sig: Tensor, *labels: Sequence[Tensor] | None, **kwargs: Any) → Tuple[Tensor, ...][source]#

Forward method of the augmenter.

Parameters:

  • sig (torch.Tensor) – Batched ECGs to be stretched or compressed, of shape (batch, lead, siglen).

  • labels (Sequence[torch.Tensor], optional) – Label tensors of the ECGs, If set, labels of ndim = 3, of shape (batch, label_len, channels) will be stretched or compressed. siglen should be divisible by label_len.

  • kwargs (dict, optional) – Not used, but kept for consistency with other augmenters.

Returns:

  • sig (torch.Tensor) – The stretched or compressed ECG tensors.

  • labels (Sequence[torch.Tensor], optional) – The stretched or compressed label tensors.