deepecgkit 0.1.0__py3-none-any.whl

deepecgkit/models/gru_model.py

@@ -0,0 +1,99 @@
+import torch
+from torch import nn
+
+from deepecgkit.registry import register_model
+
+
+@register_model(
+    name="gru",
+    description="GRU-based model for sequential ECG analysis",
+)
+class GRUECG(nn.Module):
+    """
+    GRU-based model for ECG signal classification.
+
+    A recurrent neural network using GRU layers with concat pooling
+    (adaptive avg + adaptive max + last hidden state) for ECG classification.
+    Based on the RNN architecture from the PTB-XL benchmark.
+
+    Reference:
+        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
+        hidden_size: Size of GRU hidden state (default: 256)
+        num_layers: Number of GRU layers (default: 2)
+        dropout_rate: Dropout probability (default: 0.3)
+        bidirectional: Use bidirectional GRU (default: False)
+
+    Example:
+        >>> model = GRUECG(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, 768) with hidden_size=256, unidirectional
+    """
+
+    def __init__(
+        self,
+        input_channels: int = 12,
+        output_size: int = 5,
+        hidden_size: int = 256,
+        num_layers: int = 2,
+        dropout_rate: float = 0.3,
+        bidirectional: bool = False,
+    ):
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.bidirectional = bidirectional
+        self.num_directions = 2 if bidirectional else 1
+
+        self.gru = nn.GRU(
+            input_size=input_channels,
+            hidden_size=hidden_size,
+            num_layers=num_layers,
+            batch_first=True,
+            dropout=dropout_rate if num_layers > 1 else 0,
+            bidirectional=bidirectional,
+        )
+
+        rnn_out_dim = hidden_size * self.num_directions
+        self._feature_dim = rnn_out_dim * 3
+
+        self.head = nn.Sequential(
+            nn.BatchNorm1d(self._feature_dim),
+            nn.Dropout(dropout_rate / 2),
+            nn.Linear(self._feature_dim, output_size),
+        )
+
+    @property
+    def feature_dim(self) -> int:
+        return self._feature_dim
+
+    def _concat_pool(self, rnn_output: torch.Tensor) -> torch.Tensor:
+        x = rnn_output.transpose(1, 2)
+        avg_pool = torch.mean(x, dim=2)
+        max_pool, _ = torch.max(x, dim=2)
+
+        if self.bidirectional:
+            last = torch.cat(
+                [rnn_output[:, -1, : self.hidden_size], rnn_output[:, 0, self.hidden_size :]], dim=1
+            )
+        else:
+            last = rnn_output[:, -1, :]
+
+        return torch.cat([avg_pool, max_pool, last], dim=1)
+
+    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
+        x = x.transpose(1, 2)
+        output, _ = self.gru(x)
+        return self._concat_pool(output)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.extract_features(x)
+        return self.head(x)

deepecgkit/models/inception_time.py

@@ -0,0 +1,197 @@
+import torch
+import torch.nn.functional as func
+from torch import nn
+
+from deepecgkit.registry import register_model
+
+
+class InceptionBlock1D(nn.Module):
+    """
+    Single Inception block with multi-scale 1D convolutions.
+
+    Applies parallel convolutions with different kernel sizes to capture
+    temporal patterns at multiple scales, plus a max-pool branch.
+
+    Args:
+        in_channels: Number of input channels
+        n_filters: Number of filters per branch (default: 32)
+        kernel_sizes: Tuple of kernel sizes for parallel branches
+        bottleneck_channels: Channels in bottleneck layer (default: 32)
+        use_bottleneck: Whether to use 1x1 bottleneck before convolutions
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        n_filters: int = 32,
+        kernel_sizes: tuple[int, ...] = (5, 15, 41),
+        bottleneck_channels: int = 32,
+        use_bottleneck: bool = True,
+    ):
+        super().__init__()
+
+        self.use_bottleneck = use_bottleneck
+
+        if use_bottleneck:
+            self.bottleneck = nn.Conv1d(in_channels, bottleneck_channels, kernel_size=1, bias=False)
+            conv_in = bottleneck_channels
+        else:
+            conv_in = in_channels
+
+        self.conv_branches = nn.ModuleList(
+            [
+                nn.Conv1d(
+                    conv_in,
+                    n_filters,
+                    kernel_size=ks,
+                    padding=ks // 2,
+                    bias=False,
+                )
+                for ks in kernel_sizes
+            ]
+        )
+
+        self.maxpool_branch = nn.Sequential(
+            nn.MaxPool1d(kernel_size=3, stride=1, padding=1),
+            nn.Conv1d(in_channels, n_filters, kernel_size=1, bias=False),
+        )
+
+        total_filters = n_filters * (len(kernel_sizes) + 1)
+        self.bn = nn.BatchNorm1d(total_filters)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.use_bottleneck:
+            bottleneck_out = self.bottleneck(x)
+            conv_outputs = [branch(bottleneck_out) for branch in self.conv_branches]
+        else:
+            conv_outputs = [branch(x) for branch in self.conv_branches]
+
+        pool_output = self.maxpool_branch(x)
+        conv_outputs.append(pool_output)
+
+        out = torch.cat(conv_outputs, dim=1)
+        out = self.bn(out)
+        out = func.relu(out)
+        return out
+
+
+class InceptionResidualBlock1D(nn.Module):
+    """
+    Inception block with a residual shortcut connection.
+
+    Args:
+        in_channels: Number of input channels
+        n_filters: Number of filters per Inception branch
+        kernel_sizes: Tuple of kernel sizes for parallel branches
+        bottleneck_channels: Channels in bottleneck layer
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        n_filters: int = 32,
+        kernel_sizes: tuple[int, ...] = (5, 15, 41),
+        bottleneck_channels: int = 32,
+    ):
+        super().__init__()
+
+        self.inception = InceptionBlock1D(
+            in_channels=in_channels,
+            n_filters=n_filters,
+            kernel_sizes=kernel_sizes,
+            bottleneck_channels=bottleneck_channels,
+        )
+
+        out_channels = n_filters * (len(kernel_sizes) + 1)
+
+        self.shortcut = nn.Sequential(
+            nn.Conv1d(in_channels, out_channels, kernel_size=1, bias=False),
+            nn.BatchNorm1d(out_channels),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = self.inception(x)
+        shortcut = self.shortcut(x)
+        return func.relu(out + shortcut)
+
+
+@register_model(
+    name="inception-time",
+    description="Multi-scale temporal CNN inspired by InceptionTime",
+)
+class InceptionTime1D(nn.Module):
+    """
+    InceptionTime model adapted for 1D ECG signal classification.
+
+    Uses parallel convolutions at multiple temporal scales (short/medium/long
+    kernels) with residual connections to capture both rapid arrhythmia spikes
+    and slow rhythm patterns simultaneously.
+
+    Args:
+        input_channels: Number of input channels (default: 1 for single-lead ECG)
+        output_size: Number of output classes
+        n_filters: Number of filters per Inception branch (default: 32)
+        depth: Number of Inception residual blocks (default: 6)
+        kernel_sizes: Tuple of kernel sizes for multi-scale branches
+        bottleneck_channels: Channels in bottleneck layers (default: 32)
+        dropout_rate: Dropout probability (default: 0.3)
+
+    Example:
+        >>> model = InceptionTime1D(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, 128) with n_filters=32
+    """
+
+    def __init__(
+        self,
+        input_channels: int = 1,
+        output_size: int = 4,
+        n_filters: int = 32,
+        depth: int = 6,
+        kernel_sizes: tuple[int, ...] = (5, 15, 41),
+        bottleneck_channels: int = 32,
+        dropout_rate: float = 0.3,
+    ):
+        super().__init__()
+
+        n_branches = len(kernel_sizes) + 1
+        block_out_channels = n_filters * n_branches
+
+        blocks = []
+        for i in range(depth):
+            in_ch = input_channels if i == 0 else block_out_channels
+            blocks.append(
+                InceptionResidualBlock1D(
+                    in_channels=in_ch,
+                    n_filters=n_filters,
+                    kernel_sizes=kernel_sizes,
+                    bottleneck_channels=bottleneck_channels,
+                )
+            )
+
+        self.blocks = nn.Sequential(*blocks)
+        self.global_pool = nn.AdaptiveAvgPool1d(1)
+        self.dropout = nn.Dropout(dropout_rate)
+        self._feature_dim = block_out_channels
+
+        self.classifier = nn.Linear(block_out_channels, output_size)
+
+    @property
+    def feature_dim(self) -> int:
+        return self._feature_dim
+
+    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.blocks(x)
+        x = self.global_pool(x)
+        x = torch.flatten(x, 1)
+        return x
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.extract_features(x)
+        x = self.dropout(x)
+        x = self.classifier(x)
+        return x

deepecgkit/models/kanres_wide_x.py

@@ -0,0 +1,150 @@
+import torch
+from torch import nn
+
+from deepecgkit.registry import register_model
+
+
+class KanResInit(nn.Module):
+    """Initial convolutional block for KanResWideX architecture.
+
+    Args:
+        in_channels: Number of input channels
+        filterno_1: Number of filters in first convolution
+        filterno_2: Number of filters in second convolution
+        filtersize_1: Kernel size for first convolution
+        filtersize_2: Kernel size for second convolution
+        stride: Stride for first convolution
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        filterno_1: int,
+        filterno_2: int,
+        filtersize_1: int,
+        filtersize_2: int,
+        stride: int,
+    ):
+        super().__init__()
+        self.conv1 = nn.Conv1d(in_channels, filterno_1, filtersize_1, stride=stride)
+        self.bn1 = nn.BatchNorm1d(filterno_1)
+        self.relu1 = nn.ReLU()
+        self.conv2 = nn.Conv1d(filterno_1, filterno_2, filtersize_2)
+        self.bn2 = nn.BatchNorm1d(filterno_2)
+        self.relu2 = nn.ReLU()
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu1(x)
+        x = self.conv2(x)
+        x = self.bn2(x)
+        x = self.relu2(x)
+        return x
+
+
+class KanResModule(nn.Module):
+    """Residual module with skip connection for KanResWideX architecture.
+
+    Args:
+        in_channels: Number of input channels
+        filterno_1: Number of filters in first convolution
+        filterno_2: Number of filters in second convolution
+        filtersize_1: Kernel size for first convolution
+        filtersize_2: Kernel size for second convolution
+        stride: Stride for first convolution
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        filterno_1: int,
+        filterno_2: int,
+        filtersize_1: int,
+        filtersize_2: int,
+        stride: int,
+    ):
+        super().__init__()
+        self.conv1 = nn.Conv1d(in_channels, filterno_1, filtersize_1, stride=stride, padding="same")
+        self.bn1 = nn.BatchNorm1d(filterno_1)
+        self.relu1 = nn.ReLU()
+        self.conv2 = nn.Conv1d(filterno_1, filterno_2, filtersize_2, padding="same")
+        self.bn2 = nn.BatchNorm1d(filterno_2)
+        self.relu2 = nn.ReLU()
+
+    def forward(self, x):
+        identity = x
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu1(out)
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu2(out)
+        out = out + identity
+        return out
+
+
+@register_model(
+    name="kanres-deep",
+    description="Deep KAN-ResNet architecture",
+)
+class KanResDeepX(nn.Module):
+    """KanRes-Deep-X model for ECG signal classification.
+
+    A deep residual convolutional neural network architecture designed for ECG
+    signal analysis with 8 residual blocks for improved feature extraction.
+
+    Args:
+        input_channels: Number of input channels (default: 1 for single-lead ECG)
+        output_size: Number of output classes (default: 4)
+        base_channels: Base number of channels for the architecture (default: 32)
+
+    Example:
+        >>> model = KanResDeepX(input_channels=1, output_size=4, base_channels=32)
+        >>> x = torch.randn(32, 1, 3000)
+        >>> output = model(x)
+        >>> print(output.shape)
+
+        >>> features = model.extract_features(x)
+        >>> print(features.shape)  # (32, 32)
+    """
+
+    def __init__(self, input_channels: int = 1, output_size: int = 4, base_channels: int = 32):
+        super().__init__()
+
+        self.base_channels = base_channels
+        init_channels = base_channels * 2
+
+        self.init_block = KanResInit(input_channels, init_channels, base_channels, 8, 3, 1)
+        self.pool = nn.AvgPool1d(kernel_size=2)
+
+        self.res_modules = nn.ModuleList(
+            [
+                KanResModule(base_channels, base_channels * 2, base_channels, 50, 50, 1)
+                for _ in range(8)
+            ]
+        )
+
+        self.global_pool = nn.AdaptiveAvgPool1d(1)
+        self._feature_dim = base_channels
+        self.fc = nn.Linear(base_channels, output_size)
+
+    @property
+    def feature_dim(self) -> int:
+        return self._feature_dim
+
+    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.init_block(x)
+        x = self.pool(x)
+
+        for res_module in self.res_modules:
+            x = res_module(x)
+
+        x = self.global_pool(x)
+        x = x.squeeze(-1)
+        return x
+
+    def forward(self, x):
+        x = self.extract_features(x)
+        x = self.fc(x)
+        return x

deepecgkit/models/kanres_x.py

@@ -0,0 +1,126 @@
+from pathlib import Path
+from typing import Optional, Union
+
+import torch
+import torch.nn.functional as func
+from torch import nn
+
+from deepecgkit.registry import register_model
+from deepecgkit.utils.weights import get_weight_info, load_pretrained_weights
+
+
+class ConvBlock(nn.Module):
+    """Convolutional block with normalization and activation."""
+
+    def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3):
+        super().__init__()
+        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size, padding=kernel_size // 2)
+        self.bn = nn.BatchNorm1d(out_channels)
+
+    def forward(self, x):
+        return func.relu(self.bn(self.conv(x)))
+
+
+class KanResModule(nn.Module):
+    """Residual module for the KanResWideX architecture."""
+
+    def __init__(self, channels: int):
+        super().__init__()
+        self.conv1 = ConvBlock(channels, channels)
+        self.conv2 = ConvBlock(channels, channels)
+
+    def forward(self, x):
+        residual = x
+        out = self.conv1(x)
+        out = self.conv2(out)
+        return out + residual
+
+
+@register_model(
+    name="kanres",
+    description="KAN-ResNet architecture with wide layers",
+)
+class KanResWideX(nn.Module):
+    """
+    KanRes-Wide-X model for ECG signal classification.
+
+    A convolutional neural network architecture designed for ECG signal analysis
+    with residual connections and wide blocks for improved feature extraction.
+
+    Args:
+        input_channels: Number of input channels (default: 1 for single-lead ECG)
+        output_size: Number of output classes or regression targets
+        base_channels: Base number of channels for the first layer (default: 64)
+
+    Example:
+        >>> model = KanResWideX(input_channels=1, output_size=4)
+        >>> x = torch.randn(32, 1, 3000)
+        >>> output = model(x)
+        >>> print(output.shape)  # [32, 4]
+
+        >>> features = model.extract_features(x)
+        >>> print(features.shape)  # (32, 64)
+    """
+
+    def __init__(self, input_channels: int = 1, output_size: int = 4, base_channels: int = 64):
+        super().__init__()
+
+        self.input_layer = ConvBlock(input_channels, base_channels)
+        self.res_modules = nn.Sequential(
+            KanResModule(base_channels), KanResModule(base_channels), KanResModule(base_channels)
+        )
+        self.global_pool = nn.AdaptiveAvgPool1d(1)
+        self._feature_dim = base_channels
+        self.classifier = nn.Linear(base_channels, output_size)
+
+    @property
+    def feature_dim(self) -> int:
+        return self._feature_dim
+
+    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.input_layer(x)
+        x = self.res_modules(x)
+        x = self.global_pool(x)
+        x = x.squeeze(-1)
+        return x
+
+    def forward(self, x):
+        x = self.extract_features(x)
+        return self.classifier(x)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        weights: str,
+        map_location: Optional[Union[str, torch.device]] = None,
+        force_download: bool = False,
+        **kwargs,
+    ) -> "KanResWideX":
+        """Load a pretrained KanResWideX model.
+
+        Args:
+            weights: Name of pretrained weights (e.g., "kanres-af-30s") or path to weights file
+            map_location: Device to map weights to (e.g., "cpu", "cuda")
+            force_download: If True, re-download weights even if cached
+            **kwargs: Override default model parameters from the weight registry
+
+        Returns:
+            Model with pretrained weights loaded
+
+        Example:
+            >>> model = KanResWideX.from_pretrained("kanres-af-30s")
+            >>> model = KanResWideX.from_pretrained("kanres-af-30s", map_location="cuda")
+            >>> model = KanResWideX.from_pretrained("/path/to/weights.pt", output_size=2)
+        """
+        weight_path = Path(weights)
+        if weight_path.exists():
+            state_dict = torch.load(weight_path, map_location=map_location, weights_only=True)
+            model = cls(**kwargs)
+        else:
+            info = get_weight_info(weights)
+            model_kwargs = {**info["model_kwargs"], **kwargs}
+            model = cls(**model_kwargs)
+            state_dict = load_pretrained_weights(weights, map_location, force_download)
+
+        model.load_state_dict(state_dict)
+        return model

deepecgkit/models/lstm_model.py

@@ -0,0 +1,96 @@
+import torch
+from torch import nn
+
+from deepecgkit.registry import register_model
+
+
+@register_model(
+    name="lstm",
+    description="LSTM-based model for sequential ECG analysis",
+)
+class LSTMECG(nn.Module):
+    """
+    LSTM-based model for ECG signal classification.
+
+    A recurrent neural network using bidirectional LSTM layers
+    for temporal pattern recognition in ECG signals.
+
+    Args:
+        input_channels: Number of input channels (default: 1 for single-lead ECG)
+        output_size: Number of output classes
+        hidden_size: Size of LSTM hidden state (default: 128)
+        num_layers: Number of LSTM layers (default: 2)
+        dropout_rate: Dropout probability (default: 0.3)
+        bidirectional: Use bidirectional LSTM (default: True)
+
+    Example:
+        >>> model = LSTMECG(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) with bidirectional=True
+    """
+
+    def __init__(
+        self,
+        input_channels: int = 1,
+        output_size: int = 4,
+        hidden_size: int = 128,
+        num_layers: int = 2,
+        dropout_rate: float = 0.3,
+        bidirectional: bool = True,
+    ):
+        super().__init__()
+
+        self.input_channels = input_channels
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.bidirectional = bidirectional
+        self.num_directions = 2 if bidirectional else 1
+
+        self.lstm = nn.LSTM(
+            input_size=input_channels,
+            hidden_size=hidden_size,
+            num_layers=num_layers,
+            batch_first=True,
+            dropout=dropout_rate if num_layers > 1 else 0,
+            bidirectional=bidirectional,
+        )
+
+        self.dropout = nn.Dropout(dropout_rate)
+        self._feature_dim = 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_lstm_features(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size = x.size(0)
+        x = x.transpose(1, 2)
+        _, (h_n, _) = self.lstm(x)
+
+        if self.bidirectional:
+            h_n = h_n.view(self.num_layers, 2, batch_size, self.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 extract_features(self, x: torch.Tensor) -> torch.Tensor:
+        return self._extract_lstm_features(x)
+
+    def forward(self, x):
+        x = self.extract_features(x)
+        x = self.dropout(x)
+        return self.classifier(x)