deepecgkit 0.1.0__py3-none-any.whl

deepecgkit/datasets/modules.py

@@ -0,0 +1,222 @@
+"""
+PyTorch Lightning data module for ECG datasets.
+"""
+
+from typing import Dict, Optional, Type, Union
+
+import numpy as np
+import pytorch_lightning as pl
+import torch
+from torch.utils.data import DataLoader, Dataset
+
+from deepecgkit.datasets.base import BaseECGDataset
+from deepecgkit.datasets.splitting import DataSplitter
+
+
+class ECGDataModule(pl.LightningDataModule):
+    """PyTorch Lightning data module for ECG datasets."""
+
+    def __init__(
+        self,
+        dataset: Optional[Union[BaseECGDataset, Dataset]] = None,
+        dataset_class: Optional[Type[BaseECGDataset]] = None,
+        data_dir: Optional[str] = None,
+        sampling_rate: int = 500,
+        leads: Optional[list] = None,
+        transform: Optional[callable] = None,
+        target_transform: Optional[callable] = None,
+        download: bool = False,
+        batch_size: int = 32,
+        val_split: float = 0.2,
+        test_split: float = 0.1,
+        num_workers: int = 4,
+        seed: int = 42,
+        stratify: bool = True,
+        pin_memory: bool = False,
+        persistent_workers: bool = True,
+        prefetch_factor: int = 2,
+        verbose: bool = True,
+        dataset_kwargs: Optional[Dict] = None,
+    ):
+        """Initialize the ECG data module.
+
+        Args:
+            dataset: PyTorch dataset instance (optional)
+            dataset_class: Class of dataset to create (optional)
+            data_dir: Directory containing ECG data files (optional, uses dataset's default if None)
+            sampling_rate: Sampling rate of the ECG signals (Hz)
+            leads: List of leads to use (e.g., ['I', 'II', 'III'] for standard leads)
+            transform: Optional transform to be applied to the ECG signals
+            target_transform: Optional transform to be applied to the labels
+            download: Whether to download the dataset if it doesn't exist
+            batch_size: Batch size for data loaders
+            val_split: Fraction of data to use for validation
+            test_split: Fraction of data to use for testing
+            num_workers: Number of worker processes for data loading
+            seed: Random seed for reproducibility
+            stratify: Whether to use stratified splitting based on labels
+            dataset_kwargs: Additional keyword arguments to pass to the dataset class
+        """
+        super().__init__()
+        self.dataset = dataset
+        self.dataset_class = dataset_class
+        self.data_dir = data_dir
+        self.sampling_rate = sampling_rate
+        self.leads = leads
+        self.transform = transform
+        self.target_transform = target_transform
+        self.download = download
+        self.batch_size = batch_size
+        self.val_split = val_split
+        self.test_split = test_split
+        self.num_workers = num_workers
+        self.seed = seed
+        self.stratify = stratify
+        self.pin_memory = pin_memory
+        self.persistent_workers = persistent_workers
+        self.prefetch_factor = prefetch_factor
+        self.verbose = verbose
+        self.dataset_kwargs = dataset_kwargs or {}
+
+        self.train_dataset = None
+        self.val_dataset = None
+        self.test_dataset = None
+
+    def setup(self, stage: Optional[str] = None):
+        """Set up the datasets.
+
+        Args:
+            stage: Current stage ('fit', 'validate', 'test', or 'predict')
+        """
+        if self.dataset is None:
+            if self.dataset_class is None:
+                raise ValueError("dataset_class must be provided if dataset is None")
+
+            if self.data_dir is None:
+                self.data_dir = self.dataset_class.get_default_data_dir()
+
+            init_kwargs = {
+                "data_dir": self.data_dir,
+                "sampling_rate": self.sampling_rate,
+                "transform": self.transform,
+                "target_transform": self.target_transform,
+                "download": self.download,
+            }
+
+            if self.leads is not None:
+                init_kwargs["leads"] = self.leads
+
+            init_kwargs.update(self.dataset_kwargs)
+
+            self.dataset = self.dataset_class(**init_kwargs)
+
+        stratify_labels = None
+        if (
+            self.stratify and len(self.dataset) >= 8
+        ):  # Need at least 2 samples per class for 4 classes
+            try:
+                all_labels = []
+                for i in range(len(self.dataset)):
+                    _, label = self.dataset[i]
+                    all_labels.append(label)
+                stacked = torch.stack(all_labels).numpy()
+
+                # For multi-label (2D), convert rows to string keys for stratification
+                if stacked.ndim > 1:
+                    stratify_labels = np.array(
+                        ["_".join(str(int(v)) for v in row) for row in stacked]
+                    )
+                else:
+                    stratify_labels = stacked
+
+                # Check if we have enough samples per class for stratification
+                _, counts = np.unique(stratify_labels, return_counts=True)
+                if np.min(counts) < 2:
+                    stratify_labels = None  # Disable stratification
+            except Exception:
+                stratify_labels = None  # Fallback to no stratification
+
+        # Extract groups for patient-level splitting (prevents data leakage)
+        groups = None
+        if hasattr(self.dataset, "record_names") and self.dataset.record_names:
+            groups = np.array(self.dataset.record_names)
+            if self.verbose:
+                n_groups = len(np.unique(groups))
+                print(f"Using patient-level splitting ({n_groups} groups)")
+
+        splitter = DataSplitter(
+            dataset=self.dataset,
+            val_split=self.val_split,
+            test_split=self.test_split,
+            seed=self.seed,
+            stratify=stratify_labels,
+            groups=groups,
+        )
+        self.train_dataset, self.val_dataset, self.test_dataset = splitter.split()
+
+        if self.verbose:
+            print(f"Dataset size: {len(self.train_dataset)}")
+            print(f"Validation set size: {len(self.val_dataset)}")
+            print(f"Test set size: {len(self.test_dataset)}")
+
+    def train_dataloader(self) -> DataLoader:
+        """Get the training data loader."""
+        if self.train_dataset is None:
+            raise RuntimeError("Call setup() before using the data module")
+        return DataLoader(
+            self.train_dataset,
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=self.num_workers,
+            pin_memory=self.pin_memory,
+            persistent_workers=self.persistent_workers and self.num_workers > 0,
+            prefetch_factor=self.prefetch_factor if self.num_workers > 0 else None,
+        )
+
+    def val_dataloader(self) -> DataLoader:
+        """Get the validation data loader."""
+        if self.val_dataset is None:
+            raise RuntimeError("Call setup() before using the data module")
+        return DataLoader(
+            self.val_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=self.pin_memory,
+            persistent_workers=self.persistent_workers and self.num_workers > 0,
+            prefetch_factor=self.prefetch_factor if self.num_workers > 0 else None,
+        )
+
+    def test_dataloader(self) -> DataLoader:
+        """Get the test data loader."""
+        if self.test_dataset is None:
+            raise RuntimeError("Call setup() before using the data module")
+        return DataLoader(
+            self.test_dataset,
+            batch_size=self.batch_size,
+            shuffle=False,
+            num_workers=self.num_workers,
+            pin_memory=self.pin_memory,
+            persistent_workers=self.persistent_workers and self.num_workers > 0,
+            prefetch_factor=self.prefetch_factor if self.num_workers > 0 else None,
+        )
+
+    def get_metadata(self) -> Dict:
+        """Get metadata about the dataset.
+
+        Returns:
+            Dictionary containing dataset metadata
+        """
+        if self.dataset is None:
+            raise RuntimeError("Call setup() before getting metadata")
+        if isinstance(self.dataset, BaseECGDataset):
+            return self.dataset.get_metadata()
+        return {
+            "dataset_size": len(self.dataset),
+            "batch_size": self.batch_size,
+            "num_workers": self.num_workers,
+        }
+
+    def print_metadata(self):
+        """Print metadata about the dataset."""
+        print(self.get_metadata())

deepecgkit/datasets/preprocessing.py

@@ -0,0 +1,208 @@
+from typing import ClassVar, Dict, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from scipy import signal
+
+
+class ECGStandardizer:
+    def __init__(
+        self,
+        target_sampling_rate: int = 300,
+        target_length: Optional[int] = None,
+        target_duration_seconds: Optional[float] = None,
+        normalization: str = "zscore",
+        clip_method: str = "center",
+    ):
+        self.target_sampling_rate = target_sampling_rate
+        self.normalization = normalization
+        self.clip_method = clip_method
+
+        if target_length is not None and target_duration_seconds is not None:
+            raise ValueError("Cannot specify both target_length and target_duration_seconds")
+
+        if target_duration_seconds is not None:
+            self.target_length = int(target_duration_seconds * target_sampling_rate)
+        else:
+            self.target_length = target_length
+
+    def resample(self, ecg_signal: np.ndarray, original_sampling_rate: int) -> np.ndarray:
+        if original_sampling_rate == self.target_sampling_rate:
+            return ecg_signal
+
+        num_samples = ecg_signal.shape[-1]
+        target_samples = int(num_samples * self.target_sampling_rate / original_sampling_rate)
+
+        if ecg_signal.ndim == 1:
+            return signal.resample(ecg_signal, target_samples)
+        else:
+            resampled = np.zeros((ecg_signal.shape[0], target_samples))
+            for i in range(ecg_signal.shape[0]):
+                resampled[i] = signal.resample(ecg_signal[i], target_samples)
+            return resampled
+
+    def normalize(self, ecg_signal: np.ndarray) -> np.ndarray:
+        if self.normalization == "zscore":
+            mean = np.mean(ecg_signal, axis=-1, keepdims=True)
+            std = np.std(ecg_signal, axis=-1, keepdims=True)
+            return (ecg_signal - mean) / (std + 1e-8)
+
+        elif self.normalization == "minmax":
+            min_val = np.min(ecg_signal, axis=-1, keepdims=True)
+            max_val = np.max(ecg_signal, axis=-1, keepdims=True)
+            range_val = max_val - min_val
+            return np.where(
+                range_val > 1e-8, (ecg_signal - min_val) / range_val, ecg_signal - min_val
+            )
+
+        elif self.normalization == "none":
+            return ecg_signal
+
+        else:
+            raise ValueError(f"Unknown normalization method: {self.normalization}")
+
+    def clip_or_pad(self, ecg_signal: np.ndarray) -> np.ndarray:
+        if self.target_length is None or ecg_signal.shape[-1] == self.target_length:
+            return ecg_signal
+
+        current_length = ecg_signal.shape[-1]
+
+        if ecg_signal.ndim == 1:
+            ecg_signal = ecg_signal[np.newaxis, :]
+
+        if self.clip_method not in ("center", "start", "end"):
+            raise ValueError(f"Unknown clip_method: {self.clip_method}")
+
+        if current_length < self.target_length:
+            return self._pad_signal(ecg_signal, current_length)
+        return self._clip_signal(ecg_signal, current_length)
+
+    def _pad_signal(self, ecg_signal: np.ndarray, current_length: int) -> np.ndarray:
+        pad_width = self.target_length - current_length
+        if self.clip_method == "center":
+            pad_left = pad_width // 2
+            pad_right = pad_width - pad_left
+            return np.pad(ecg_signal, ((0, 0), (pad_left, pad_right)), mode="constant")
+        if self.clip_method == "start":
+            return np.pad(ecg_signal, ((0, 0), (0, pad_width)), mode="constant")
+        return np.pad(ecg_signal, ((0, 0), (pad_width, 0)), mode="constant")
+
+    def _clip_signal(self, ecg_signal: np.ndarray, current_length: int) -> np.ndarray:
+        excess = current_length - self.target_length
+        if self.clip_method == "center":
+            start = excess // 2
+            return ecg_signal[:, start : start + self.target_length]
+        if self.clip_method == "start":
+            return ecg_signal[:, : self.target_length]
+        return ecg_signal[:, -self.target_length :]
+
+    def __call__(self, ecg_signal: np.ndarray, original_sampling_rate: int) -> np.ndarray:
+        ecg_signal = self.resample(ecg_signal, original_sampling_rate)
+        ecg_signal = self.clip_or_pad(ecg_signal)
+        ecg_signal = self.normalize(ecg_signal)
+        return ecg_signal
+
+
+class ECGSegmenter:
+    def __init__(
+        self,
+        segment_duration_seconds: float,
+        sampling_rate: int,
+        overlap: float = 0.0,
+    ):
+        self.segment_duration_seconds = segment_duration_seconds
+        self.sampling_rate = sampling_rate
+        self.overlap = overlap
+        self.segment_length = int(segment_duration_seconds * sampling_rate)
+        self.stride = int(self.segment_length * (1 - overlap))
+
+    def segment(self, ecg_signal: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+        if ecg_signal.ndim == 1:
+            ecg_signal = ecg_signal[np.newaxis, :]
+
+        _, signal_length = ecg_signal.shape
+
+        if signal_length < self.segment_length:
+            return np.array([]), np.array([])
+
+        segments = []
+        start_indices = []
+
+        for start in range(0, signal_length - self.segment_length + 1, self.stride):
+            segment = ecg_signal[:, start : start + self.segment_length]
+            segments.append(segment)
+            start_indices.append(start)
+
+        return np.array(segments), np.array(start_indices)
+
+
+class RhythmAnnotationExtractor:
+    RHYTHM_MAP: ClassVar[Dict[str, int]] = {
+        "(AFIB": 1,
+        "(AFL": 2,
+        "(J": 3,
+        "(N": 0,
+    }
+
+    def __init__(
+        self,
+        sampling_rate: int,
+        binary_classification: bool = False,
+    ):
+        self.sampling_rate = sampling_rate
+        self.binary_classification = binary_classification
+
+    def extract_labels(
+        self,
+        annotation,
+        signal_length: int,
+        original_sampling_rate: Optional[int] = None,
+    ) -> np.ndarray:
+        labels = np.zeros(signal_length, dtype=np.int64)
+
+        if not hasattr(annotation, "aux_note") or not hasattr(annotation, "sample"):
+            return labels
+
+        scale = 1.0
+        if original_sampling_rate is not None and original_sampling_rate != self.sampling_rate:
+            scale = self.sampling_rate / original_sampling_rate
+
+        for i, (sample_idx, aux_note) in enumerate(zip(annotation.sample, annotation.aux_note)):
+            scaled_idx = int(sample_idx * scale)
+            if scaled_idx >= signal_length:
+                break
+
+            rhythm_code = self.RHYTHM_MAP.get(aux_note.strip(), 0)
+
+            if self.binary_classification:
+                rhythm_code = 1 if rhythm_code == 1 else 0
+
+            raw_next = annotation.sample[i + 1] if i + 1 < len(annotation.sample) else signal_length
+            next_sample = min(int(raw_next * scale), signal_length)
+
+            labels[scaled_idx:next_sample] = rhythm_code
+
+        return labels
+
+    def segment_with_labels(
+        self, labels: np.ndarray, segment_start_indices: np.ndarray, segment_length: int
+    ) -> np.ndarray:
+        segment_labels = []
+
+        for start_idx in segment_start_indices:
+            segment_label_region = labels[start_idx : start_idx + segment_length]
+
+            unique, counts = np.unique(segment_label_region, return_counts=True)
+            majority_label = unique[np.argmax(counts)]
+
+            segment_labels.append(majority_label)
+
+        return np.array(segment_labels)
+
+
+def convert_to_tensor(
+    data: Union[np.ndarray, torch.Tensor], dtype: torch.dtype = torch.float32
+) -> torch.Tensor:
+    if isinstance(data, torch.Tensor):
+        return data.to(dtype)
+    return torch.from_numpy(data).to(dtype)