ECG_UNET

class torch_ecg.models.ECG_UNET(classes: Sequence[str], n_leads: int, config: CFG | None = None)

Bases: Module, CkptMixin, SizeMixin, CitationMixin

U-Net for (multi-lead) ECG wave delineation.

The U-Net is a fully convolutional network originally proposed for biomedical image segmentation [1]. This architecture is applied to ECG wave delineation in [2]. This implementation is based on an open-source implementation on GitHub [3].

Parameters:

• classes (Sequence[str]) – List of names of the classes.

• n_leads (int) – Number of input leads (number of input channels).

• config (CFG, optional,) – Other hyper-parameters, including kernel sizes, etc. Refer to the corresponding config file.

References

[1]

Ronneberger, Olaf, Philipp Fischer, and Thomas Brox. “U-net: Convolutional networks for biomedical image segmentation.” International Conference on Medical image computing and computer-assisted intervention. Springer, 2015.

[2]

Moskalenko, Viktor, Nikolai Zolotykh, and Grigory Osipov. “Deep Learning for ECG Segmentation.” International Conference on Neuroinformatics. Springer, Cham, 2019.

[3]

milesial/Pytorch-UNet

compute_output_shape(seq_len: int | None = None, batch_size: int | None = None) → Sequence[int | None]

Compute the output shape of the model.

Parameters:

• seq_len (int, optional) – The length of the input signal tensor.

• batch_size (int, optional) – The batch size of the input signal tensor.

Returns:

output_shape – The output shape of the model.

Return type:

sequence

forward(input: Tensor) → Tensor

Forward pass of the model.

Parameters:

input (torch.Tensor) – Input signal tensor, of shape (batch_size, n_channels, seq_len).

Returns:

output – Output tensor, of shape (batch_size, n_channels, seq_len).

Return type:

torch.Tensor

inference(input: Tensor, bin_pred_thr: float = 0.5) → Tensor

Method for making inference on a single input.