deepecgkit 0.1.0__py3-none-any.whl

deepecgkit/models/crnn.py

@@ -0,0 +1,117 @@
+import torch
+from torch import nn
+
+from deepecgkit.registry import register_model
+
+
+@register_model(
+    name="crnn",
+    description="CNN-LSTM hybrid for local feature extraction and temporal aggregation",
+)
+class CRNN(nn.Module):
+    """
+    Convolutional Recurrent Neural Network for ECG classification.
+
+    Uses a CNN front-end for local feature extraction followed by a
+    bidirectional LSTM for temporal aggregation, combining the strengths
+    of both architectures.
+
+    Args:
+        input_channels: Number of input channels (default: 1 for single-lead ECG)
+        output_size: Number of output classes
+        cnn_channels: List of channel sizes for CNN stages (default: [32, 64, 128, 256])
+        lstm_hidden_size: Size of LSTM hidden state (default: 128)
+        lstm_num_layers: Number of LSTM layers (default: 2)
+        bidirectional: Use bidirectional LSTM (default: True)
+        dropout_rate: Dropout probability (default: 0.3)
+
+    Example:
+        >>> model = CRNN(input_channels=1, output_size=4)
+        >>> x = torch.randn(32, 1, 3000)
+        >>> output = model(x)
+        >>> print(output.shape)
+
+        >>> features = model.extract_features(x)
+        >>> print(features.shape)  # (32, 256)
+    """
+
+    def __init__(
+        self,
+        input_channels: int = 1,
+        output_size: int = 4,
+        cnn_channels: list | None = None,
+        lstm_hidden_size: int = 128,
+        lstm_num_layers: int = 2,
+        bidirectional: bool = True,
+        dropout_rate: float = 0.3,
+    ):
+        super().__init__()
+
+        if cnn_channels is None:
+            cnn_channels = [32, 64, 128, 256]
+
+        self.lstm_hidden_size = lstm_hidden_size
+        self.lstm_num_layers = lstm_num_layers
+        self.bidirectional = bidirectional
+        self.num_directions = 2 if bidirectional else 1
+
+        layers = []
+        in_ch = input_channels
+        kernel_sizes = [7, 5, 3, 3]
+        for i, out_ch in enumerate(cnn_channels):
+            ks = kernel_sizes[i] if i < len(kernel_sizes) else 3
+            layers.extend(
+                [
+                    nn.Conv1d(in_ch, out_ch, kernel_size=ks, padding=ks // 2),
+                    nn.BatchNorm1d(out_ch),
+                    nn.ReLU(),
+                    nn.MaxPool1d(kernel_size=2),
+                    nn.Dropout(dropout_rate),
+                ]
+            )
+            in_ch = out_ch
+
+        self.cnn = nn.Sequential(*layers)
+
+        self.lstm = nn.LSTM(
+            input_size=cnn_channels[-1],
+            hidden_size=lstm_hidden_size,
+            num_layers=lstm_num_layers,
+            batch_first=True,
+            dropout=dropout_rate if lstm_num_layers > 1 else 0,
+            bidirectional=bidirectional,
+        )
+
+        self.dropout = nn.Dropout(dropout_rate)
+        self._feature_dim = lstm_hidden_size * self.num_directions
+
+        self.classifier = nn.Sequential(
+            nn.Linear(self._feature_dim, 64),
+            nn.ReLU(),
+            nn.Dropout(dropout_rate),
+            nn.Linear(64, output_size),
+        )
+
+    @property
+    def feature_dim(self) -> int:
+        return self._feature_dim
+
+    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size = x.size(0)
+        x = self.cnn(x)
+        x = x.transpose(1, 2)
+        _, (h_n, _) = self.lstm(x)
+
+        if self.bidirectional:
+            h_n = h_n.view(self.lstm_num_layers, 2, batch_size, self.lstm_hidden_size)
+            h_n = h_n[-1]
+            h_n = torch.cat([h_n[0], h_n[1]], dim=1)
+        else:
+            h_n = h_n[-1]
+
+        return h_n
+
+    def forward(self, x):
+        x = self.extract_features(x)
+        x = self.dropout(x)
+        return self.classifier(x)

deepecgkit/models/deep_res_cnn.py

@@ -0,0 +1,215 @@
+import torch
+import torch.nn.functional as func
+from torch import nn
+
+from deepecgkit.registry import register_model
+
+
+class StemResBlock2D(nn.Module):
+    """
+    First residual block following the stem convolution.
+
+    Unlike subsequent blocks, this starts directly with convolution
+    (no pre-activation BN/ReLU), matching the Keras reference architecture.
+    All convolutions use valid padding (no padding).
+
+    Args:
+        in_channels: Number of input filter channels
+        out_channels: Number of output filter channels
+        dropout_rate: Dropout probability
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        dropout_rate: float = 0.1,
+    ):
+        super().__init__()
+
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=(1, 5))
+        self.bn1 = nn.BatchNorm2d(out_channels)
+        self.dropout1 = nn.Dropout(dropout_rate)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=(1, 5), stride=(1, 2))
+
+        self.shortcut_pool = nn.MaxPool2d(kernel_size=(1, 9), stride=(1, 2))
+        self.shortcut_proj = nn.Conv2d(in_channels, out_channels, kernel_size=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = func.relu(out)
+        out = self.dropout1(out)
+        out = self.conv2(out)
+
+        shortcut = self.shortcut_pool(x)
+        shortcut = self.shortcut_proj(shortcut)
+
+        return out + shortcut
+
+
+class PreActResBlock2D(nn.Module):
+    """
+    Pre-activation 2D residual block for blocks 2-4.
+
+    Follows BN -> ReLU -> Dropout -> Conv -> BN -> ReLU -> Dropout -> Conv(stride=2)
+    pattern matching the Keras reference model. All convolutions use valid padding.
+
+    Args:
+        in_channels: Number of input filter channels
+        out_channels: Number of output filter channels
+        dropout_rate: Dropout probability
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        dropout_rate: float = 0.1,
+    ):
+        super().__init__()
+
+        self.bn1 = nn.BatchNorm2d(in_channels)
+        self.dropout1 = nn.Dropout(dropout_rate)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=(1, 5))
+        self.bn2 = nn.BatchNorm2d(out_channels)
+        self.dropout2 = nn.Dropout(dropout_rate)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=(1, 5), stride=(1, 2))
+
+        self.shortcut_pool = nn.MaxPool2d(kernel_size=(1, 9), stride=(1, 2))
+        self.shortcut_proj = (
+            nn.Conv2d(in_channels, out_channels, kernel_size=1)
+            if in_channels != out_channels
+            else nn.Identity()
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = self.bn1(x)
+        out = func.relu(out)
+        out = self.dropout1(out)
+        out = self.conv1(out)
+
+        out = self.bn2(out)
+        out = func.relu(out)
+        out = self.dropout2(out)
+        out = self.conv2(out)
+
+        shortcut = self.shortcut_pool(x)
+        shortcut = self.shortcut_proj(shortcut)
+
+        return out + shortcut
+
+
+@register_model(
+    name="deep-res-cnn",
+    description="Deep Residual 2D CNN (Elyamani et al. 2022) for multi-lead ECG classification",
+)
+class DeepResCNN(nn.Module):
+    """
+    Deep Residual 2D CNN for ECG classification.
+
+    Faithful implementation of Elyamani et al. (2022). Uses Conv2d with (1, k)
+    kernels and valid padding to process each ECG lead independently along the
+    time axis, then fuses across leads with a (leads, 1) convolution at the end.
+
+    Input convention: (batch, leads, time) -- standard deepecg-kit format.
+    Internally reshaped to (batch, 1, leads, time) for 2D convolution.
+
+    The classifier head includes L2 regularization matching the original Keras
+    model. Call l2_regularization_loss() to obtain the penalty term.
+
+    Reference:
+        https://github.com/HaneenElyamani/ECG-classification
+
+    Args:
+        input_channels: Number of ECG leads (default: 12)
+        output_size: Number of output classes (default: 5)
+        dropout_rate: Dropout probability in residual blocks (default: 0.1)
+
+    Example:
+        >>> model = DeepResCNN(input_channels=12, output_size=5)
+        >>> x = torch.randn(32, 12, 1000)
+        >>> output = model(x)
+        >>> print(output.shape)
+
+        >>> features = model.extract_features(x)
+        >>> print(features.shape)  # (32, 128)
+    """
+
+    def __init__(
+        self,
+        input_channels: int = 12,
+        output_size: int = 5,
+        dropout_rate: float = 0.1,
+    ):
+        super().__init__()
+
+        self.stem_conv = nn.Conv2d(1, 32, kernel_size=(1, 7))
+        self.stem_bn = nn.BatchNorm2d(32)
+
+        self.block1 = StemResBlock2D(32, 64, dropout_rate=dropout_rate)
+        self.block2 = PreActResBlock2D(64, 64, dropout_rate=dropout_rate)
+        self.block3 = PreActResBlock2D(64, 128, dropout_rate=dropout_rate)
+        self.block4 = PreActResBlock2D(128, 128, dropout_rate=dropout_rate)
+
+        self.lead_fusion = nn.Conv2d(128, 128, kernel_size=(input_channels, 1))
+        self.lead_fusion_bn = nn.BatchNorm2d(128)
+        self.pool = nn.AdaptiveAvgPool2d(1)
+
+        self._feature_dim = 128
+
+        self.fc1 = nn.Linear(128, 128)
+        self.fc1_bn = nn.BatchNorm1d(128)
+        self.fc1_drop = nn.Dropout(dropout_rate)
+
+        self.fc2 = nn.Linear(128, 64)
+        self.fc2_bn = nn.BatchNorm1d(64)
+        self.fc2_drop = nn.Dropout(0.15)
+
+        self.fc_out = nn.Linear(64, output_size)
+
+    @property
+    def feature_dim(self) -> int:
+        return self._feature_dim
+
+    def l2_regularization_loss(self) -> torch.Tensor:
+        """Return the L2 penalty for the classifier head weights.
+
+        Matches the Keras kernel_regularizer=L2(lambda) on the two Dense layers.
+        Add this to the training loss for full equivalence with the original model.
+        """
+        return 0.005 * self.fc1.weight.pow(2).sum() + 0.009 * self.fc2.weight.pow(2).sum()
+
+    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
+        x = x.unsqueeze(1)
+
+        x = self.stem_conv(x)
+        x = self.stem_bn(x)
+        x = func.relu(x)
+
+        x = self.block1(x)
+        x = self.block2(x)
+        x = self.block3(x)
+        x = self.block4(x)
+
+        x = self.lead_fusion(x)
+        x = self.lead_fusion_bn(x)
+        x = func.relu(x)
+        x = self.pool(x)
+        return torch.flatten(x, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.extract_features(x)
+
+        x = self.fc1(x)
+        x = self.fc1_bn(x)
+        x = func.relu(x)
+        x = self.fc1_drop(x)
+
+        x = self.fc2(x)
+        x = self.fc2_bn(x)
+        x = func.relu(x)
+        x = self.fc2_drop(x)
+
+        x = self.fc_out(x)
+        return x

deepecgkit/models/ecg_dualnet.py

@@ -0,0 +1,230 @@
+import math
+
+import torch
+from torch import nn
+
+from deepecgkit.registry import register_model
+
+
+class PositionalEncoding(nn.Module):
+    """
+    Sinusoidal positional encoding for transformer branch.
+
+    Args:
+        d_model: Dimension of the model
+        max_len: Maximum sequence length (default: 5000)
+        dropout: Dropout probability (default: 0.1)
+    """
+
+    def __init__(self, d_model: int, max_len: int = 5000, dropout: float = 0.1):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        position = torch.arange(max_len).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
+        pe = torch.zeros(1, max_len, d_model)
+        pe[0, :, 0::2] = torch.sin(position * div_term)
+        pe[0, :, 1::2] = torch.cos(position * div_term)
+        self.register_buffer("pe", pe)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.pe[:, : x.size(1)]
+        return self.dropout(x)
+
+
+class CNNLSTMBranch(nn.Module):
+    """
+    CNN feature extraction followed by LSTM temporal aggregation.
+
+    Args:
+        input_channels: Number of input ECG channels
+        cnn_channels: Output channels for CNN stages (default: 128)
+        lstm_hidden: LSTM hidden size (default: 128)
+        lstm_layers: Number of LSTM layers (default: 1)
+        dropout_rate: Dropout probability (default: 0.3)
+    """
+
+    def __init__(
+        self,
+        input_channels: int,
+        cnn_channels: int = 128,
+        lstm_hidden: int = 128,
+        lstm_layers: int = 1,
+        dropout_rate: float = 0.3,
+    ):
+        super().__init__()
+        self.lstm_hidden = lstm_hidden
+
+        self.cnn = nn.Sequential(
+            nn.Conv1d(input_channels, 64, kernel_size=7, padding=3),
+            nn.BatchNorm1d(64),
+            nn.ReLU(),
+            nn.MaxPool1d(kernel_size=2),
+            nn.Dropout(dropout_rate),
+            nn.Conv1d(64, cnn_channels, kernel_size=5, padding=2),
+            nn.BatchNorm1d(cnn_channels),
+            nn.ReLU(),
+            nn.MaxPool1d(kernel_size=2),
+            nn.Dropout(dropout_rate),
+            nn.Conv1d(cnn_channels, cnn_channels, kernel_size=3, padding=1),
+            nn.BatchNorm1d(cnn_channels),
+            nn.ReLU(),
+            nn.MaxPool1d(kernel_size=2),
+        )
+
+        self.lstm = nn.LSTM(
+            input_size=cnn_channels,
+            hidden_size=lstm_hidden,
+            num_layers=lstm_layers,
+            batch_first=True,
+            dropout=dropout_rate if lstm_layers > 1 else 0,
+            bidirectional=True,
+        )
+
+        self.output_dim = lstm_hidden * 2
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size = x.size(0)
+        x = self.cnn(x)
+        x = x.transpose(1, 2)
+        _, (h_n, _) = self.lstm(x)
+        h_n = h_n.view(-1, 2, batch_size, self.lstm_hidden)
+        h_n = h_n[-1]
+        return torch.cat([h_n[0], h_n[1]], dim=1)
+
+
+class TransformerBranch(nn.Module):
+    """
+    Transformer encoder for global temporal attention.
+
+    Args:
+        input_channels: Number of input ECG channels
+        d_model: Transformer model dimension (default: 128)
+        nhead: Number of attention heads (default: 8)
+        num_layers: Number of transformer encoder layers (default: 2)
+        dim_feedforward: Feedforward dimension (default: 256)
+        dropout_rate: Dropout probability (default: 0.1)
+    """
+
+    def __init__(
+        self,
+        input_channels: int,
+        d_model: int = 128,
+        nhead: int = 8,
+        num_layers: int = 2,
+        dim_feedforward: int = 256,
+        dropout_rate: float = 0.1,
+    ):
+        super().__init__()
+
+        self.input_projection = nn.Linear(input_channels, d_model)
+        self.positional_encoding = PositionalEncoding(d_model, dropout=dropout_rate)
+
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=d_model,
+            nhead=nhead,
+            dim_feedforward=dim_feedforward,
+            dropout=dropout_rate,
+            batch_first=True,
+        )
+
+        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+        self.pool = nn.AdaptiveAvgPool1d(1)
+        self.output_dim = d_model
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x.transpose(1, 2)
+        x = self.input_projection(x)
+        x = self.positional_encoding(x)
+        x = self.encoder(x)
+        x = x.transpose(1, 2)
+        x = self.pool(x)
+        return x.squeeze(-1)
+
+
+@register_model(
+    name="dualnet",
+    description="Dual-path architecture with CNN-LSTM and Transformer branches for ECG",
+)
+class ECGDualNet(nn.Module):
+    """
+    Dual-path ECG classification network.
+
+    Runs a CNN-LSTM branch and a Transformer branch in parallel,
+    then fuses their outputs via concatenation and a fusion head.
+
+    Args:
+        input_channels: Number of input channels (default: 1 for single-lead ECG)
+        output_size: Number of output classes
+        cnn_channels: CNN output channels in CNN-LSTM branch (default: 128)
+        lstm_hidden: LSTM hidden size in CNN-LSTM branch (default: 128)
+        lstm_layers: Number of LSTM layers (default: 1)
+        d_model: Transformer model dimension (default: 128)
+        nhead: Number of attention heads (default: 8)
+        transformer_layers: Number of transformer encoder layers (default: 2)
+        dim_feedforward: Transformer feedforward dimension (default: 256)
+        dropout_rate: Dropout probability (default: 0.3)
+
+    Example:
+        >>> model = ECGDualNet(input_channels=1, output_size=4)
+        >>> x = torch.randn(32, 1, 3000)
+        >>> output = model(x)
+        >>> print(output.shape)
+
+        >>> features = model.extract_features(x)
+        >>> print(features.shape)  # (32, 384)
+    """
+
+    def __init__(
+        self,
+        input_channels: int = 1,
+        output_size: int = 4,
+        cnn_channels: int = 128,
+        lstm_hidden: int = 128,
+        lstm_layers: int = 1,
+        d_model: int = 128,
+        nhead: int = 8,
+        transformer_layers: int = 2,
+        dim_feedforward: int = 256,
+        dropout_rate: float = 0.3,
+    ):
+        super().__init__()
+
+        self.cnn_lstm_branch = CNNLSTMBranch(
+            input_channels=input_channels,
+            cnn_channels=cnn_channels,
+            lstm_hidden=lstm_hidden,
+            lstm_layers=lstm_layers,
+            dropout_rate=dropout_rate,
+        )
+
+        self.transformer_branch = TransformerBranch(
+            input_channels=input_channels,
+            d_model=d_model,
+            nhead=nhead,
+            num_layers=transformer_layers,
+            dim_feedforward=dim_feedforward,
+            dropout_rate=dropout_rate,
+        )
+
+        self._feature_dim = self.cnn_lstm_branch.output_dim + self.transformer_branch.output_dim
+
+        self.classifier = nn.Sequential(
+            nn.Linear(self._feature_dim, 128),
+            nn.ReLU(),
+            nn.Dropout(dropout_rate),
+            nn.Linear(128, output_size),
+        )
+
+    @property
+    def feature_dim(self) -> int:
+        return self._feature_dim
+
+    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
+        cnn_lstm_out = self.cnn_lstm_branch(x)
+        transformer_out = self.transformer_branch(x)
+        return torch.cat([cnn_lstm_out, transformer_out], dim=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.extract_features(x)
+        return self.classifier(x)

deepecgkit/models/fcn_wang.py

@@ -0,0 +1,88 @@
+import torch
+from torch import nn
+
+from deepecgkit.registry import register_model
+
+
+@register_model(
+    name="fcn-wang",
+    description="Fully Convolutional Network (Wang et al.) for time series classification",
+)
+class FCNWang(nn.Module):
+    """
+    Fully Convolutional Network for ECG signal classification.
+
+    Based on Wang et al.'s FCN architecture for time series classification,
+    using three convolutional blocks with batch normalization and global
+    average pooling. No dense layers except the final classifier.
+
+    Reference:
+        Wang Z., Yan W., Oates T. "Time Series Classification from Scratch
+        with Deep Neural Networks: A Strong Baseline" (2017)
+        https://github.com/helme/ecg_ptbxl_benchmarking
+
+    Args:
+        input_channels: Number of input channels (default: 12 for 12-lead ECG)
+        output_size: Number of output classes
+        filters: List of filter counts for each conv block (default: [128, 256, 128])
+        kernel_sizes: List of kernel sizes for each conv block (default: [8, 5, 3])
+        dropout_rate: Dropout probability before classifier (default: 0.3)
+
+    Example:
+        >>> model = FCNWang(input_channels=12, output_size=5)
+        >>> x = torch.randn(32, 12, 1000)
+        >>> output = model(x)
+        >>> print(output.shape)
+
+        >>> features = model.extract_features(x)
+        >>> print(features.shape)  # (32, 128)
+    """
+
+    def __init__(
+        self,
+        input_channels: int = 12,
+        output_size: int = 5,
+        filters: list[int] | None = None,
+        kernel_sizes: list[int] | None = None,
+        dropout_rate: float = 0.3,
+    ):
+        super().__init__()
+
+        if filters is None:
+            filters = [128, 256, 128]
+        if kernel_sizes is None:
+            kernel_sizes = [8, 5, 3]
+
+        assert len(filters) == len(kernel_sizes)
+
+        blocks = []
+        in_ch = input_channels
+        for nf, ks in zip(filters, kernel_sizes):
+            blocks.append(
+                nn.Sequential(
+                    nn.Conv1d(in_ch, nf, kernel_size=ks, padding=(ks - 1) // 2, bias=False),
+                    nn.BatchNorm1d(nf),
+                    nn.ReLU(inplace=True),
+                )
+            )
+            in_ch = nf
+
+        self.conv_blocks = nn.Sequential(*blocks)
+        self.global_pool = nn.AdaptiveAvgPool1d(1)
+        self.dropout = nn.Dropout(dropout_rate)
+        self._feature_dim = filters[-1]
+        self.classifier = nn.Linear(self._feature_dim, output_size)
+
+    @property
+    def feature_dim(self) -> int:
+        return self._feature_dim
+
+    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.conv_blocks(x)
+        x = self.global_pool(x)
+        return torch.flatten(x, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.extract_features(x)
+        x = self.dropout(x)
+        return self.classifier(x)