deepecgkit 0.1.0__py3-none-any.whl

deepecgkit/evaluation/visualization.py

@@ -0,0 +1,285 @@
+"""
+Visualization functions for ECG evaluation.
+"""
+
+import warnings
+from pathlib import Path
+from typing import List, Optional
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import seaborn as sns
+from sklearn.calibration import calibration_curve
+from sklearn.metrics import auc, confusion_matrix, roc_curve
+
+
+def plot_confusion_matrix(
+    y_true: np.ndarray,
+    y_pred: np.ndarray,
+    class_names: Optional[List[str]] = None,
+    title: str = "Confusion Matrix",
+    save_path: Optional[str] = None,
+):
+    """Plot confusion matrix."""
+    labels = sorted(set(y_true) | set(y_pred))
+    cm = confusion_matrix(y_true, y_pred, labels=labels)
+
+    if class_names is None or len(class_names) != len(labels):
+        class_names = [str(lbl) for lbl in labels]
+
+    _, ax = plt.subplots(figsize=(8, 6))
+
+    sns.heatmap(
+        cm,
+        annot=True,
+        fmt="d",
+        cmap="Blues",
+        xticklabels=class_names,
+        yticklabels=class_names,
+        ax=ax,
+    )
+
+    ax.set_title(title)
+    ax.set_ylabel("True Label")
+    ax.set_xlabel("Predicted Label")
+
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches="tight")
+        plt.close()
+    else:
+        plt.show()
+
+
+def plot_roc_curve(
+    y_true: np.ndarray,
+    y_scores: np.ndarray,
+    title: str = "ROC Curve",
+    save_path: Optional[str] = None,
+):
+    """Plot ROC curve."""
+    try:
+        fpr, tpr, _ = roc_curve(y_true, y_scores)
+        roc_auc = auc(fpr, tpr)
+
+        plt.figure(figsize=(8, 6))
+        plt.plot(
+            fpr,
+            tpr,
+            color="darkorange",
+            lw=2,
+            label=f"ROC curve (area = {roc_auc:.2f})",
+        )
+        plt.plot([0, 1], [0, 1], color="navy", lw=2, linestyle="--")
+        plt.xlim([0.0, 1.0])
+        plt.ylim([0.0, 1.05])
+        plt.xlabel("False Positive Rate")
+        plt.ylabel("True Positive Rate")
+        plt.title(title)
+        plt.legend(loc="lower right")
+
+        if save_path:
+            plt.savefig(save_path, dpi=300, bbox_inches="tight")
+            plt.close()
+        else:
+            plt.show()
+
+    except ImportError:
+        warnings.warn("Could not plot ROC curve - missing dependencies", stacklevel=2)
+
+
+def plot_predictions(
+    y_true: np.ndarray,
+    y_pred: np.ndarray,
+    title: str = "Predictions vs True Values",
+    save_path: Optional[str] = None,
+):
+    """Plot predictions vs true values."""
+    plt.figure(figsize=(8, 6))
+    plt.scatter(y_true, y_pred, alpha=0.6)
+    plt.plot([y_true.min(), y_true.max()], [y_true.min(), y_true.max()], "r--", lw=2)
+    plt.xlabel("True Values")
+    plt.ylabel("Predictions")
+    plt.title(title)
+
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches="tight")
+        plt.close()
+    else:
+        plt.show()
+
+
+def plot_ecg_signals(
+    data: np.ndarray,
+    sampling_rate: float = 500.0,
+    leads: Optional[List[str]] = None,
+    title: str = "ECG Signals",
+    save_path: Optional[str] = None,
+):
+    """Plot ECG signals."""
+    if data.ndim == 1:
+        data = data.reshape(-1, 1)
+
+    n_samples, n_leads = data.shape
+    time = np.arange(n_samples) / sampling_rate
+
+    if leads is None:
+        leads = [f"Lead {i + 1}" for i in range(n_leads)]
+
+    _fig, axes = plt.subplots(n_leads, 1, figsize=(12, 2 * n_leads), sharex=True)
+    if n_leads == 1:
+        axes = [axes]
+
+    for i, ax in enumerate(axes):
+        ax.plot(time, data[:, i])
+        ax.set_ylabel(f"{leads[i]} (mV)")
+        ax.grid(True, alpha=0.3)
+
+    axes[-1].set_xlabel("Time (s)")
+    plt.suptitle(title)
+    plt.tight_layout()
+
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches="tight")
+        plt.close()
+    else:
+        plt.show()
+
+
+def plot_training_curves(
+    metrics_path: str,
+    save_dir: Optional[str] = None,
+    save_path: Optional[str] = None,
+):
+    """Plot training and validation loss and accuracy as separate figures.
+
+    Args:
+        metrics_path: Path to CSVLogger metrics.csv file.
+        save_dir: Directory to save separate loss.png and accuracy.png files.
+        save_path: Deprecated combined path; if provided and save_dir is not,
+            saves loss plot to this path for backward compatibility.
+    """
+    df = pd.read_csv(metrics_path)
+
+    has_acc = "train_acc" in df.columns
+
+    resolved_dir = None
+    if save_dir:
+        resolved_dir = Path(save_dir)
+    elif save_path:
+        resolved_dir = Path(save_path).parent
+
+    train_loss = df.dropna(subset=["train_loss"]).groupby("epoch")["train_loss"].mean()
+    val_loss = df.dropna(subset=["val_loss"]).groupby("epoch")["val_loss"].mean()
+
+    plt.figure(figsize=(10, 5))
+    plt.plot(train_loss.index, train_loss.values, label="Train Loss")
+    plt.plot(val_loss.index, val_loss.values, label="Val Loss")
+    plt.xlabel("Epoch")
+    plt.ylabel("Loss")
+    plt.title("Training & Validation Loss")
+    plt.legend()
+    plt.grid(True, alpha=0.3)
+    plt.tight_layout()
+
+    if resolved_dir:
+        plt.savefig(str(resolved_dir / "loss.png"), dpi=300, bbox_inches="tight")
+        plt.close()
+    else:
+        plt.show()
+
+    if has_acc:
+        train_acc = df.dropna(subset=["train_acc"]).groupby("epoch")["train_acc"].mean()
+        val_acc = df.dropna(subset=["val_acc"]).groupby("epoch")["val_acc"].mean()
+
+        plt.figure(figsize=(10, 5))
+        plt.plot(train_acc.index, train_acc.values, label="Train Acc")
+        plt.plot(val_acc.index, val_acc.values, label="Val Acc")
+        plt.xlabel("Epoch")
+        plt.ylabel("Accuracy")
+        plt.title("Training & Validation Accuracy")
+        plt.legend()
+        plt.grid(True, alpha=0.3)
+        plt.tight_layout()
+
+        if resolved_dir:
+            plt.savefig(str(resolved_dir / "accuracy.png"), dpi=300, bbox_inches="tight")
+            plt.close()
+        else:
+            plt.show()
+
+
+def plot_calibration_curve(
+    y_true: np.ndarray,
+    y_prob: np.ndarray,
+    class_names: Optional[List[str]] = None,
+    n_bins: int = 10,
+    title: str = "Calibration Plot",
+    save_dir: Optional[str] = None,
+    save_path: Optional[str] = None,
+):
+    """Plot reliability diagram and prediction distribution as separate figures.
+
+    Args:
+        y_true: True labels.
+        y_prob: Predicted probabilities (n_samples, n_classes).
+        class_names: Optional list of class names.
+        n_bins: Number of bins for calibration curve.
+        title: Title for the calibration curve plot.
+        save_dir: Directory to save calibration_curve.png and prediction_distribution.png.
+        save_path: Deprecated combined path; if provided and save_dir is not,
+            uses its parent directory for backward compatibility.
+    """
+    num_classes = y_prob.shape[1]
+    if class_names is None:
+        class_names = [f"Class {i}" for i in range(num_classes)]
+
+    resolved_dir = None
+    if save_dir:
+        resolved_dir = Path(save_dir)
+    elif save_path:
+        resolved_dir = Path(save_path).parent
+
+    plt.figure(figsize=(8, 6))
+    for i in range(num_classes):
+        y_binary = (y_true == i).astype(int)
+        prob_class = y_prob[:, i]
+        if y_binary.sum() == 0:
+            continue
+        try:
+            fraction_of_positives, mean_predicted_value = calibration_curve(
+                y_binary, prob_class, n_bins=n_bins, strategy="uniform"
+            )
+            plt.plot(mean_predicted_value, fraction_of_positives, "s-", label=class_names[i])
+        except ValueError:
+            continue
+
+    plt.plot([0, 1], [0, 1], "k--", label="Perfectly calibrated")
+    plt.xlabel("Mean Predicted Probability")
+    plt.ylabel("Fraction of Positives")
+    plt.title(title)
+    plt.legend(loc="lower right", fontsize="small")
+    plt.grid(True, alpha=0.3)
+    plt.tight_layout()
+
+    if resolved_dir:
+        plt.savefig(str(resolved_dir / "calibration_curve.png"), dpi=300, bbox_inches="tight")
+        plt.close()
+    else:
+        plt.show()
+
+    plt.figure(figsize=(8, 6))
+    for i in range(num_classes):
+        plt.hist(y_prob[:, i], bins=n_bins, alpha=0.5, label=class_names[i])
+    plt.xlabel("Predicted Probability")
+    plt.ylabel("Count")
+    plt.title("Prediction Distribution")
+    plt.legend(fontsize="small")
+    plt.grid(True, alpha=0.3)
+    plt.tight_layout()
+
+    if resolved_dir:
+        plt.savefig(str(resolved_dir / "prediction_distribution.png"), dpi=300, bbox_inches="tight")
+        plt.close()
+    else:
+        plt.show()

deepecgkit/models/__init__.py

@@ -0,0 +1,51 @@
+"""
+Deep learning models for ECG analysis.
+
+This module contains various neural network architectures optimized for ECG signal processing.
+"""
+
+from deepecgkit.models.af_classifier import AFModel
+from deepecgkit.models.convnext_v2_1d import ConvNeXtV21D
+from deepecgkit.models.crnn import CRNN
+from deepecgkit.models.deep_res_cnn import DeepResCNN
+from deepecgkit.models.ecg_dualnet import ECGDualNet
+from deepecgkit.models.fcn_wang import FCNWang
+from deepecgkit.models.gru_model import GRUECG
+from deepecgkit.models.inception_time import InceptionTime1D
+from deepecgkit.models.kanres_wide_x import KanResDeepX
+from deepecgkit.models.kanres_x import KanResWideX
+from deepecgkit.models.lstm_model import LSTMECG
+from deepecgkit.models.mamba1d import Mamba1D
+from deepecgkit.models.medformer import Medformer
+from deepecgkit.models.resnet1d import ResNet1D
+from deepecgkit.models.resnet1d_wang import ResNetWang
+from deepecgkit.models.se_resnet1d import SEResNet1D
+from deepecgkit.models.simple_cnn import SimpleCNN
+from deepecgkit.models.tcn import TCN
+from deepecgkit.models.transformer_ecg import TransformerECG
+from deepecgkit.models.xresnet1d import XResNet1D
+from deepecgkit.models.xresnet1d_benchmark import XResNet1dBenchmark
+
+__all__ = [
+    "CRNN",
+    "GRUECG",
+    "LSTMECG",
+    "TCN",
+    "AFModel",
+    "ConvNeXtV21D",
+    "DeepResCNN",
+    "ECGDualNet",
+    "FCNWang",
+    "InceptionTime1D",
+    "KanResDeepX",
+    "KanResWideX",
+    "Mamba1D",
+    "Medformer",
+    "ResNet1D",
+    "ResNetWang",
+    "SEResNet1D",
+    "SimpleCNN",
+    "TransformerECG",
+    "XResNet1D",
+    "XResNet1dBenchmark",
+]

deepecgkit/models/af_classifier.py

@@ -0,0 +1,177 @@
+from pathlib import Path
+from typing import Optional, Union
+
+import torch
+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):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        use_dropout: bool = False,
+        use_pooling: bool = False,
+        dropout_rate: float = 0.1,
+    ):
+        super().__init__()
+        self.conv = nn.Conv1d(
+            in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size
+        )
+        self.batch_norm = nn.BatchNorm1d(out_channels)
+        self.relu = nn.ReLU()
+        self.use_dropout = use_dropout
+        self.use_pooling = use_pooling
+
+        if use_dropout:
+            self.dropout = nn.Dropout(dropout_rate)
+        if use_pooling:
+            self.pool = nn.MaxPool1d(kernel_size=2, stride=2)
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.relu(self.batch_norm(x))
+
+        if self.use_dropout:
+            x = self.dropout(x)
+        if self.use_pooling:
+            x = self.pool(x)
+
+        return x
+
+
+@register_model(
+    name="afmodel",
+    description="Atrial Fibrillation model optimized for 30s ECG segments",
+    default_kwargs={"recording_length": 30},
+)
+class AFModel(nn.Module):
+    """
+    Atrial Fibrillation classification model.
+
+    A convolutional neural network specifically designed for AF detection
+    in ECG signals with configurable recording length support.
+
+    Args:
+        input_channels: Number of input ECG leads (default: 1)
+        output_size: Number of output classes (default: 4)
+        recording_length: Recording length in seconds (must be 6, 10, or 30)
+
+    Example:
+        >>> model = AFModel(recording_length=30)
+        >>> x = torch.randn(32, 1, 9000)
+        >>> output = model(x)
+        >>> print(output.shape)  # (32, 4)
+
+        >>> features = model.extract_features(x)
+        >>> print(features.shape)  # (32, 196)
+    """
+
+    def __init__(
+        self,
+        input_channels: int = 1,
+        output_size: int = 4,
+        recording_length: int = 30,
+    ):
+        super().__init__()
+        assert recording_length in [6, 10, 30], (
+            f"Recording length must be 6, 10, or 30, got {recording_length}"
+        )
+
+        self.recording_length = recording_length
+        kernel_size = 6 if recording_length in [6, 10] else 32
+
+        layer_configs = [
+            (input_channels, 64, True, True),
+            (64, 64, False, False),
+            (64, 64, True, True),
+            (64, 64, False, False),
+            (64, 64, True, True),
+            (64, 64, False, False),
+            (64, 128, True, True),
+            (128, 128, False, False),
+            (128, 128, True, True),
+            (128, 128, False, False),
+            (128, 196, True, True),
+            (196, 196, False, False),
+            (196, 196, True, True),
+        ]
+
+        self.conv_layers = nn.ModuleList(
+            [
+                ConvBlock(in_ch, out_ch, kernel_size, dropout, pooling)
+                for in_ch, out_ch, dropout, pooling in layer_configs
+            ]
+        )
+
+        self.adaptive_pool = nn.AdaptiveAvgPool1d(1)
+        self.flatten = nn.Flatten()
+        self._feature_dim = 196
+        self.dropout = nn.Dropout(0.3)
+        self.output = nn.Linear(in_features=self._feature_dim, out_features=output_size)
+
+    @property
+    def feature_dim(self) -> int:
+        return self._feature_dim
+
+    def extract_features(self, data: torch.Tensor) -> torch.Tensor:
+        for conv_layer in self.conv_layers:
+            data = conv_layer(data)
+        data = self.adaptive_pool(data)
+        data = self.flatten(data)
+        return data
+
+    def forward(self, data):
+        data = self.extract_features(data)
+        data = self.dropout(data)
+        return self.output(data)
+
+    def get_feature_size(self, input_size: int) -> int:
+        with torch.no_grad():
+            dummy_input = torch.randn(1, 1, input_size)
+
+            for conv_layer in self.conv_layers:
+                dummy_input = conv_layer(dummy_input)
+
+            return dummy_input.shape[-1]
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        weights: str,
+        map_location: Optional[Union[str, torch.device]] = None,
+        force_download: bool = False,
+        **kwargs,
+    ) -> "AFModel":
+        """Load a pretrained AFModel.
+
+        Args:
+            weights: Name of pretrained weights (e.g., "afmodel-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 = AFModel.from_pretrained("afmodel-30s")
+            >>> model = AFModel.from_pretrained("afmodel-30s", map_location="cuda")
+            >>> model = AFModel.from_pretrained("/path/to/weights.pt", recording_length=30)
+        """
+        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/convnext_v2_1d.py

@@ -0,0 +1,183 @@
+import torch
+from torch import nn
+
+from deepecgkit.registry import register_model
+
+
+class GRN(nn.Module):
+    """
+    Global Response Normalization layer for 1D signals.
+
+    Normalizes feature responses globally across the temporal dimension,
+    the key differentiator of ConvNeXtV2 over ConvNeXtV1.
+
+    Args:
+        dim: Number of channels
+    """
+
+    def __init__(self, dim: int):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, dim, 1))
+        self.beta = nn.Parameter(torch.zeros(1, dim, 1))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        gx = torch.norm(x, p=2, dim=2, keepdim=True)
+        nx = gx / (gx.mean(dim=1, keepdim=True) + 1e-6)
+        return self.gamma * (x * nx) + self.beta + x
+
+
+class LayerNorm1d(nn.Module):
+    """
+    LayerNorm for (B, C, T) tensors, normalizing over the channel dimension.
+
+    Args:
+        dim: Number of channels
+    """
+
+    def __init__(self, dim: int):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x.transpose(1, 2)
+        x = self.norm(x)
+        x = x.transpose(1, 2)
+        return x
+
+
+class ConvNeXtV2Block1D(nn.Module):
+    """
+    Single ConvNeXtV2 block with inverted bottleneck and GRN.
+
+    depthwise conv -> LayerNorm -> pointwise expand -> GELU -> GRN -> pointwise project
+
+    Args:
+        dim: Number of input/output channels
+        expansion_factor: Expansion ratio for inverted bottleneck (default: 4)
+        kernel_size: Kernel size for depthwise convolution (default: 7)
+    """
+
+    def __init__(self, dim: int, expansion_factor: int = 4, kernel_size: int = 7):
+        super().__init__()
+        hidden_dim = dim * expansion_factor
+        self.dwconv = nn.Conv1d(
+            dim,
+            dim,
+            kernel_size=kernel_size,
+            padding=kernel_size // 2,
+            groups=dim,
+        )
+        self.norm = LayerNorm1d(dim)
+        self.pwconv_expand = nn.Conv1d(dim, hidden_dim, kernel_size=1)
+        self.act = nn.GELU()
+        self.grn = GRN(hidden_dim)
+        self.pwconv_project = nn.Conv1d(hidden_dim, dim, kernel_size=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        x = self.dwconv(x)
+        x = self.norm(x)
+        x = self.pwconv_expand(x)
+        x = self.act(x)
+        x = self.grn(x)
+        x = self.pwconv_project(x)
+        return x + residual
+
+
+@register_model(
+    name="convnext-v2",
+    description="Modern ConvNet with depthwise convolutions and Global Response Normalization for ECG",
+)
+class ConvNeXtV21D(nn.Module):
+    """
+    ConvNeXtV2 adapted for 1D ECG signals.
+
+    A modern convolutional architecture using depthwise separable convolutions,
+    LayerNorm, GELU activation, and Global Response Normalization (GRN).
+
+    Args:
+        input_channels: Number of input channels (default: 1 for single-lead ECG)
+        output_size: Number of output classes
+        dims: Channel dimensions for each stage (default: [64, 128, 256, 512])
+        depths: Number of blocks per stage (default: [2, 2, 6, 2])
+        kernel_size: Kernel size for depthwise convolutions (default: 7)
+        expansion_factor: Expansion ratio for inverted bottleneck (default: 4)
+        dropout_rate: Dropout probability (default: 0.3)
+
+    Example:
+        >>> model = ConvNeXtV21D(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, 512)
+    """
+
+    def __init__(
+        self,
+        input_channels: int = 1,
+        output_size: int = 4,
+        dims: list | None = None,
+        depths: list | None = None,
+        kernel_size: int = 7,
+        expansion_factor: int = 4,
+        dropout_rate: float = 0.3,
+    ):
+        super().__init__()
+
+        if dims is None:
+            dims = [64, 128, 256, 512]
+        if depths is None:
+            depths = [2, 2, 6, 2]
+
+        self.stem = nn.Sequential(
+            nn.Conv1d(input_channels, dims[0], kernel_size=7, stride=4, padding=3),
+            LayerNorm1d(dims[0]),
+        )
+
+        stages = []
+        for i in range(len(dims)):
+            stage_blocks = []
+            for _ in range(depths[i]):
+                stage_blocks.append(ConvNeXtV2Block1D(dims[i], expansion_factor, kernel_size))
+            stages.append(nn.Sequential(*stage_blocks))
+
+        self.stages = nn.ModuleList(stages)
+
+        self.downsamples = nn.ModuleList()
+        for i in range(len(dims) - 1):
+            self.downsamples.append(
+                nn.Sequential(
+                    LayerNorm1d(dims[i]),
+                    nn.Conv1d(dims[i], dims[i + 1], kernel_size=2, stride=2),
+                )
+            )
+
+        self.pool = nn.AdaptiveAvgPool1d(1)
+        self.final_norm = nn.LayerNorm(dims[-1])
+        self.dropout = nn.Dropout(dropout_rate)
+        self._feature_dim = dims[-1]
+
+        self.classifier = nn.Linear(dims[-1], output_size)
+
+    @property
+    def feature_dim(self) -> int:
+        return self._feature_dim
+
+    def extract_features(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.stem(x)
+        for i, stage in enumerate(self.stages):
+            x = stage(x)
+            if i < len(self.downsamples):
+                x = self.downsamples[i](x)
+        x = self.pool(x)
+        x = torch.flatten(x, 1)
+        x = self.final_norm(x)
+        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