deepecgkit 0.1.0__py3-none-any.whl

deepecgkit/datasets/splitting.py

@@ -0,0 +1,225 @@
+"""
+Data splitting and loader creation utilities.
+
+This module provides classes for splitting datasets and creating data loaders.
+"""
+
+from collections import Counter
+from typing import Optional, Tuple
+
+import numpy as np
+import torch
+from sklearn.model_selection import train_test_split
+from torch.utils.data import Dataset, Subset, random_split
+
+
+class DataSplitter:
+    """Handles dataset splitting into train, validation, and test sets."""
+
+    def __init__(
+        self,
+        dataset: Dataset,
+        val_split: float = 0.2,
+        test_split: float = 0.1,
+        seed: Optional[int] = 42,
+        stratify: Optional[np.ndarray] = None,
+        groups: Optional[np.ndarray] = None,
+    ):
+        """Initialize the data splitter.
+
+        Args:
+            dataset: PyTorch dataset to split
+            val_split: Fraction of data to use for validation
+            test_split: Fraction of data to use for testing
+            seed: Random seed for reproducibility
+            stratify: Array of labels for stratified splitting (optional)
+            groups: Array of group identifiers (e.g. patient/record IDs) to keep
+                all samples from the same group in the same split (optional).
+                Prevents data leakage from correlated samples.
+        """
+        self.dataset = dataset
+        self.val_split = val_split
+        self.test_split = test_split
+        self.seed = seed
+        self.stratify = stratify
+        self.groups = groups
+
+        if val_split + test_split >= 1.0:
+            raise ValueError("val_split + test_split must be less than 1.0")
+
+    def split(self) -> Tuple[Dataset, Dataset, Dataset]:
+        """Split the dataset into train, validation, and test sets.
+
+        Returns:
+            Tuple of (train_dataset, val_dataset, test_dataset)
+        """
+        if self.groups is not None and len(self.groups) > 0:
+            return self._split_by_groups()
+
+        total_size = len(self.dataset)
+
+        # Handle small datasets by ensuring at least 1 sample per split if possible
+        val_size = max(1, int(total_size * self.val_split)) if self.val_split > 0 else 0
+        test_size = max(1, int(total_size * self.test_split)) if self.test_split > 0 else 0
+
+        # Adjust for very small datasets
+        if total_size <= 2:
+            val_size = 0
+            test_size = 0
+        elif total_size == 3:
+            val_size = 1
+            test_size = 0
+        elif val_size + test_size >= total_size:
+            # Reduce splits for small datasets
+            val_size = 1 if self.val_split > 0 else 0
+            test_size = 1 if self.test_split > 0 and total_size > val_size + 1 else 0
+
+        train_size = total_size - val_size - test_size
+
+        if self.stratify is not None and total_size > 2:
+            try:
+                indices = np.arange(total_size)
+                if val_size + test_size > 0:
+                    train_idx, temp_idx = train_test_split(
+                        indices,
+                        test_size=(val_size + test_size) / total_size,
+                        stratify=self.stratify,
+                        random_state=self.seed,
+                    )
+                    if test_size > 0 and val_size > 0:
+                        val_idx, test_idx = train_test_split(
+                            temp_idx,
+                            test_size=test_size / (val_size + test_size),
+                            stratify=self.stratify[temp_idx],
+                            random_state=self.seed,
+                        )
+                    elif test_size > 0:
+                        test_idx = temp_idx
+                        val_idx = []
+                    else:
+                        val_idx = temp_idx
+                        test_idx = []
+                else:
+                    train_idx = indices
+                    val_idx = []
+                    test_idx = []
+
+                train_dataset = Subset(self.dataset, train_idx)
+                val_dataset = (
+                    Subset(self.dataset, val_idx) if len(val_idx) > 0 else Subset(self.dataset, [])
+                )
+                test_dataset = (
+                    Subset(self.dataset, test_idx)
+                    if len(test_idx) > 0
+                    else Subset(self.dataset, [])
+                )
+                return train_dataset, val_dataset, test_dataset
+            except ValueError:
+                pass  # Fall through to random splitting
+
+        if val_size + test_size == 0:
+            train_dataset = self.dataset
+            val_dataset = Subset(self.dataset, [])
+            test_dataset = Subset(self.dataset, [])
+        else:
+            train_dataset, val_dataset, test_dataset = random_split(
+                self.dataset,
+                [train_size, val_size, test_size],
+                generator=torch.Generator().manual_seed(self.seed),
+            )
+
+        return train_dataset, val_dataset, test_dataset
+
+    def _split_by_groups(self) -> Tuple[Dataset, Dataset, Dataset]:
+        """Split by groups so all samples from one group stay in the same split.
+
+        This implements inter-patient splitting for medical datasets where
+        segments from the same recording/patient are highly correlated.
+        """
+        unique_groups = np.unique(self.groups)
+        n_groups = len(unique_groups)
+
+        if n_groups < 3:
+            raise ValueError(f"Need at least 3 groups for train/val/test split, got {n_groups}")
+
+        val_n = max(1, int(n_groups * self.val_split)) if self.val_split > 0 else 0
+        test_n = max(1, int(n_groups * self.test_split)) if self.test_split > 0 else 0
+
+        if val_n + test_n >= n_groups:
+            val_n = 1 if self.val_split > 0 else 0
+            test_n = 1 if self.test_split > 0 and n_groups > val_n + 1 else 0
+
+        # Compute majority label per group for stratified group splitting
+        group_stratify = None
+        if self.stratify is not None:
+            group_to_labels: dict = {}
+            for g, label in zip(self.groups, self.stratify):
+                group_to_labels.setdefault(g, []).append(label)
+            group_stratify = np.array(
+                [Counter(group_to_labels[g]).most_common(1)[0][0] for g in unique_groups]
+            )
+            _, counts = np.unique(group_stratify, return_counts=True)
+            if np.min(counts) < 2:
+                group_stratify = None
+
+        group_indices = np.arange(n_groups)
+
+        try:
+            train_gi, temp_gi = train_test_split(
+                group_indices,
+                test_size=(val_n + test_n) / n_groups,
+                stratify=group_stratify,
+                random_state=self.seed,
+            )
+            if val_n > 0 and test_n > 0:
+                temp_stratify = group_stratify[temp_gi] if group_stratify is not None else None
+                if temp_stratify is not None:
+                    _, tc = np.unique(temp_stratify, return_counts=True)
+                    if np.min(tc) < 2:
+                        temp_stratify = None
+                val_gi, test_gi = train_test_split(
+                    temp_gi,
+                    test_size=test_n / (val_n + test_n),
+                    stratify=temp_stratify,
+                    random_state=self.seed,
+                )
+            elif test_n > 0:
+                test_gi = temp_gi
+                val_gi = np.array([], dtype=int)
+            else:
+                val_gi = temp_gi
+                test_gi = np.array([], dtype=int)
+        except ValueError:
+            # Stratification failed, split without it
+            train_gi, temp_gi = train_test_split(
+                group_indices,
+                test_size=(val_n + test_n) / n_groups,
+                random_state=self.seed,
+            )
+            if val_n > 0 and test_n > 0:
+                val_gi, test_gi = train_test_split(
+                    temp_gi,
+                    test_size=test_n / (val_n + test_n),
+                    random_state=self.seed,
+                )
+            elif test_n > 0:
+                test_gi = temp_gi
+                val_gi = np.array([], dtype=int)
+            else:
+                val_gi = temp_gi
+                test_gi = np.array([], dtype=int)
+
+        # Map group indices back to sample indices
+        train_groups = set(unique_groups[train_gi])
+        val_groups = set(unique_groups[val_gi]) if len(val_gi) > 0 else set()
+        test_groups = set(unique_groups[test_gi]) if len(test_gi) > 0 else set()
+
+        train_idx = [i for i, g in enumerate(self.groups) if g in train_groups]
+        val_idx = [i for i, g in enumerate(self.groups) if g in val_groups]
+        test_idx = [i for i, g in enumerate(self.groups) if g in test_groups]
+
+        return (
+            Subset(self.dataset, train_idx),
+            Subset(self.dataset, val_idx) if val_idx else Subset(self.dataset, []),
+            Subset(self.dataset, test_idx) if test_idx else Subset(self.dataset, []),
+        )

deepecgkit/datasets/unified_af.py

@@ -0,0 +1,316 @@
+import logging
+from pathlib import Path
+from typing import ClassVar, Dict, List, Optional, Tuple, Type, Union
+
+import numpy as np
+import torch
+from torch.utils.data import ConcatDataset, Dataset
+
+from deepecgkit.registry import register_dataset
+
+from .af_classification import AFClassificationDataset
+from .base import BaseECGDataset
+from .ltafdb import LTAFDBDataset
+from .mitbih_afdb import MITBIHAFDBDataset
+
+logger = logging.getLogger(__name__)
+
+# PhysioNet 2017 uses labels: N=0, A=1, O=2, ~=3
+# Unified 4-class scheme:     Normal=0, AF=1, AFL=2, J=3
+# "Other" (O) and "Noisy" (~) have no unified equivalent → dropped
+_PHYSIONET2017_LABEL_REMAP = {
+    0: 0,  # N (Normal) → Normal
+    1: 1,  # A (AF) → AF
+    2: None,  # O (Other) → drop
+    3: None,  # ~ (Noisy) → drop
+}
+
+# Binary scheme: Non-AF=0, AF=1
+# "Other" maps to Non-AF, "Noisy" is still dropped
+_PHYSIONET2017_BINARY_REMAP = {
+    0: 0,  # N (Normal) → Non-AF
+    1: 1,  # A (AF) → AF
+    2: 0,  # O (Other) → Non-AF
+    3: None,  # ~ (Noisy) → drop
+}
+
+
+class _RemappedDataset(Dataset):
+    """Wraps a dataset with label remapping and sample filtering.
+
+    Samples whose labels map to None are excluded from the dataset.
+    """
+
+    def __init__(
+        self,
+        dataset,
+        label_map: Dict[int, Optional[int]],
+        target_transform: Optional[callable] = None,
+    ):
+        self._dataset = dataset
+        self._label_map = label_map
+        self._target_transform = target_transform
+
+        # Filter out samples whose labels map to None
+        if hasattr(dataset, "labels"):
+            self._indices = [
+                i for i, lbl in enumerate(dataset.labels) if label_map.get(int(lbl)) is not None
+            ]
+            self.labels = [label_map[int(dataset.labels[i])] for i in self._indices]
+        else:
+            self._indices = list(range(len(dataset)))
+            self.labels = []
+
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        signal, label = self._dataset[self._indices[idx]]
+        new_label = torch.tensor(self._label_map[label.item()], dtype=torch.long)
+        if self._target_transform is not None:
+            new_label = self._target_transform(new_label)
+        return signal, new_label
+
+    def __len__(self) -> int:
+        return len(self._indices)
+
+
+@register_dataset(
+    name="unified-af",
+    input_channels=1,
+    num_classes=4,
+    description="Unified AF dataset combining PhysioNet 2017, MIT-BIH, LTAFDB",
+)
+class UnifiedAFDataset(BaseECGDataset):
+    """Unified AF Dataset combining multiple PhysioNet AF databases.
+
+    Combines samples from the PhysioNet 2017 Challenge, MIT-BIH AFDB, and LTAFDB
+    into a single dataset for AF classification. Supports both binary (AF vs Non-AF)
+    and 4-class (Normal, AF, AFL, J) classification modes.
+
+    PhysioNet 2017 labels are remapped to the unified scheme:
+      - Normal (N) → Normal, AF (A) → AF
+      - Other (O) and Noisy (~) are dropped in 4-class mode
+      - Other (O) → Non-AF and Noisy (~) is dropped in binary mode
+    """
+
+    CLASS_LABELS: ClassVar[List[str]] = ["Normal", "AF", "AFL", "J"]
+    LEADS: ClassVar[List[str]] = ["ECG"]
+
+    AVAILABLE_DATASETS: ClassVar[Dict[str, Type[BaseECGDataset]]] = {
+        "physionet2017": AFClassificationDataset,
+        "mitbih_afdb": MITBIHAFDBDataset,
+        "ltafdb": LTAFDBDataset,
+    }
+
+    def _resolve_dataset_dir(self, dataset_name: str) -> Path:
+        dataset_class = self.AVAILABLE_DATASETS[dataset_name]
+        return dataset_class.get_default_data_dir()
+
+    def __init__(
+        self,
+        data_dir: Optional[Union[str, Path]] = None,
+        sampling_rate: int = 300,
+        segment_duration_seconds: float = 10.0,
+        datasets: Optional[List[str]] = None,
+        binary_classification: bool = False,
+        normalization: str = "zscore",
+        transform: Optional[callable] = None,
+        target_transform: Optional[callable] = None,
+        download: bool = False,
+        force_download: bool = False,
+        verbose: bool = True,
+        dataset_kwargs: Optional[Dict[str, Dict]] = None,
+    ):
+        self.segment_duration_seconds = segment_duration_seconds
+        self.binary_classification = binary_classification
+        self.verbose = verbose
+        self.dataset_kwargs = dataset_kwargs or {}
+
+        if datasets is None:
+            datasets = ["physionet2017", "mitbih_afdb", "ltafdb"]
+
+        for dataset_name in datasets:
+            if dataset_name not in self.AVAILABLE_DATASETS:
+                raise ValueError(
+                    f"Unknown dataset: {dataset_name}. "
+                    f"Available: {list(self.AVAILABLE_DATASETS.keys())}"
+                )
+
+        self.dataset_names = datasets
+        self.datasets = []
+        self.dataset_sizes = []
+
+        super().__init__(
+            data_dir=data_dir,
+            sampling_rate=sampling_rate,
+            leads=self.LEADS,
+            transform=transform,
+            target_transform=target_transform,
+            download=download,
+            force_download=force_download,
+        )
+
+    def download(self):
+        for dataset_name in self.dataset_names:
+            dataset_class = self.AVAILABLE_DATASETS[dataset_name]
+            dataset_dir = self._resolve_dataset_dir(dataset_name)
+
+            if dataset_dir.exists() and any(dataset_dir.iterdir()):
+                if self.verbose:
+                    print(f"\n{'=' * 60}")
+                    print(f"Skipping {dataset_name}: already exists at {dataset_dir}")
+                    print(f"{'=' * 60}")
+                continue
+
+            if self.verbose:
+                print(f"\n{'=' * 60}")
+                print(f"Downloading dataset: {dataset_name}")
+                print(f"{'=' * 60}")
+
+            dataset_dir.mkdir(parents=True, exist_ok=True)
+
+            # Minimal instance for download only — sub-dataset download()
+            # methods only require self.data_dir and self.verbose.
+            instance = object.__new__(dataset_class)
+            instance.data_dir = dataset_dir
+            instance.verbose = self.verbose
+            instance.download()
+
+    def _get_dataset_kwargs(self, dataset_name: str) -> Dict:
+        base_kwargs = {
+            "sampling_rate": self.sampling_rate,
+            "transform": self.transform,
+            "segment_duration_seconds": self.segment_duration_seconds,
+            "download": False,
+        }
+
+        if dataset_name in ["mitbih_afdb", "ltafdb"]:
+            base_kwargs["target_transform"] = self.target_transform
+            base_kwargs["binary_classification"] = self.binary_classification
+        elif dataset_name == "physionet2017":
+            # target_transform is omitted — applied after label remapping
+            # in _RemappedDataset (see _load_data)
+            pass
+
+        if dataset_name in self.dataset_kwargs:
+            base_kwargs.update(self.dataset_kwargs[dataset_name])
+
+        return base_kwargs
+
+    def _load_data(self):
+        if self.verbose:
+            print("\nLoading unified AF dataset...")
+
+        for dataset_name in self.dataset_names:
+            if self.verbose:
+                print(f"\n{'=' * 60}")
+                print(f"Loading: {dataset_name}")
+                print(f"{'=' * 60}")
+
+            dataset_class = self.AVAILABLE_DATASETS[dataset_name]
+            dataset_dir = self._resolve_dataset_dir(dataset_name)
+
+            kwargs = self._get_dataset_kwargs(dataset_name)
+            kwargs["data_dir"] = dataset_dir
+            kwargs["verbose"] = self.verbose
+
+            try:
+                dataset = dataset_class(**kwargs)
+
+                # Remap PhysioNet 2017 labels to the unified scheme
+                if dataset_name == "physionet2017":
+                    remap = (
+                        _PHYSIONET2017_BINARY_REMAP
+                        if self.binary_classification
+                        else _PHYSIONET2017_LABEL_REMAP
+                    )
+                    original_len = len(dataset)
+                    dataset = _RemappedDataset(
+                        dataset, remap, target_transform=self.target_transform
+                    )
+                    dropped = original_len - len(dataset)
+                    if dropped > 0 and self.verbose:
+                        print(f"  Dropped {dropped} samples with unmappable labels (Other/Noisy)")
+
+                self.datasets.append(dataset)
+                self.dataset_sizes.append(len(dataset))
+
+                if self.verbose:
+                    print(f"Loaded {len(dataset)} samples from {dataset_name}")
+                    if hasattr(dataset, "get_class_distribution"):
+                        dist = dataset.get_class_distribution()
+                        print(f"Class distribution: {dist}")
+
+            except Exception as e:
+                logger.warning("Failed to load %s: %s", dataset_name, e)
+                if self.verbose:
+                    print(f"Failed to load {dataset_name}: {e}")
+                continue
+
+        if len(self.datasets) == 0:
+            raise RuntimeError("No datasets were successfully loaded")
+
+        self.concat_dataset = ConcatDataset(self.datasets)
+
+        if self.verbose:
+            print(f"\n{'=' * 60}")
+            print("Unified dataset created")
+            print(f"Total datasets: {len(self.datasets)}")
+            print(f"Total samples: {len(self.concat_dataset)}")
+            print(f"{'=' * 60}")
+            self._print_overall_distribution()
+
+    def _print_overall_distribution(self):
+        if len(self.datasets) == 0:
+            return
+
+        all_labels = []
+        for dataset in self.datasets:
+            if hasattr(dataset, "labels"):
+                all_labels.extend(dataset.labels)
+
+        if len(all_labels) == 0:
+            return
+
+        unique, counts = np.unique(all_labels, return_counts=True)
+        print("\nOverall class distribution:")
+        for label_idx, count in zip(unique, counts):
+            class_name = self.class_names[label_idx]
+            percentage = (count / len(all_labels)) * 100
+            print(f"  {class_name}: {count} ({percentage:.1f}%)")
+
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        return self.concat_dataset[idx]
+
+    def __len__(self) -> int:
+        return len(self.concat_dataset)
+
+    @property
+    def num_classes(self) -> int:
+        return 2 if self.binary_classification else len(self.CLASS_LABELS)
+
+    @property
+    def class_names(self) -> List[str]:
+        if self.binary_classification:
+            return ["Non-AF", "AF"]
+        return self.CLASS_LABELS
+
+    def get_dataset_info(self) -> Dict:
+        return {
+            "num_datasets": len(self.datasets),
+            "dataset_names": self.dataset_names,
+            "dataset_sizes": self.dataset_sizes,
+            "total_samples": len(self),
+            "num_classes": self.num_classes,
+            "class_names": self.class_names,
+        }
+
+    def get_class_distribution(self) -> Dict[str, int]:
+        all_labels = []
+        for dataset in self.datasets:
+            if hasattr(dataset, "labels"):
+                all_labels.extend(dataset.labels)
+
+        if len(all_labels) == 0:
+            return {}
+
+        unique, counts = np.unique(all_labels, return_counts=True)
+        return {self.class_names[int(label)]: int(count) for label, count in zip(unique, counts)}

deepecgkit/evaluation/__init__.py

@@ -0,0 +1,40 @@
+"""
+Model evaluation and metrics module for DeepECG-Kit.
+
+This module provides comprehensive evaluation functionality including:
+- Classification metrics (accuracy, precision, recall, F1, AUC)
+- Regression metrics (MSE, MAE, R²)
+- ECG-specific metrics
+- Visualization tools
+- Statistical analysis
+"""
+
+from deepecgkit.evaluation.evaluator import ECGEvaluator
+from deepecgkit.evaluation.metrics import (
+    calculate_classification_metrics,
+    calculate_regression_metrics,
+    confusion_matrix_analysis,
+    roc_auc_score,
+)
+from deepecgkit.evaluation.visualization import (
+    plot_calibration_curve,
+    plot_confusion_matrix,
+    plot_ecg_signals,
+    plot_predictions,
+    plot_roc_curve,
+    plot_training_curves,
+)
+
+__all__ = [
+    "ECGEvaluator",
+    "calculate_classification_metrics",
+    "calculate_regression_metrics",
+    "confusion_matrix_analysis",
+    "plot_calibration_curve",
+    "plot_confusion_matrix",
+    "plot_ecg_signals",
+    "plot_predictions",
+    "plot_roc_curve",
+    "plot_training_curves",
+    "roc_auc_score",
+]