deepecgkit.datasets

Dataset classes, data modules, and preprocessing utilities for ECG analysis.

Base Classes

BaseECGDataset

Bases: Dataset, ABC

Base class for all ECG datasets in deepecg-kit.

Subclasses implement _download and _load_data. The accessor layer (__getitem__, __len__, num_classes, class_names, get_record_info, get_class_distribution) is provided here for any dataset that fills the standard storage attributes during _load_data:

• self.signals: sequence of per-sample arrays (list or ndarray)
• self.labels: sequence of integer class labels
• self.record_names: one source-record name per sample
• CLASS_LABELS: class-level list of class names
• binary_classification: optional bool collapsing labels to Non-AF/AF

Datasets with a different storage model (multi-label PTB-XL, the composed unified-af) override the relevant accessors.

Source code in deepecgkit/datasets/base.py

class BaseECGDataset(Dataset, ABC):
    """Base class for all ECG datasets in deepecg-kit.

    Subclasses implement ``_download`` and ``_load_data``. The accessor layer
    (``__getitem__``, ``__len__``, ``num_classes``, ``class_names``,
    ``get_record_info``, ``get_class_distribution``) is provided here for any
    dataset that fills the standard storage attributes during ``_load_data``:

    - ``self.signals``: sequence of per-sample arrays (list or ndarray)
    - ``self.labels``: sequence of integer class labels
    - ``self.record_names``: one source-record name per sample
    - ``CLASS_LABELS``: class-level list of class names
    - ``binary_classification``: optional bool collapsing labels to Non-AF/AF

    Datasets with a different storage model (multi-label PTB-XL, the composed
    unified-af) override the relevant accessors.
    """

    #: Class names for the standard multi-class labelling.
    CLASS_LABELS: ClassVar[list[str]] = []
    #: Class names used when binary_classification is enabled.
    BINARY_CLASS_LABELS: ClassVar[list[str]] = ["Non-AF", "AF"]
    #: On-disk directory name under ~/.deepecgkit/datasets. Defaults to the
    #: lowercased class name; pin it before renaming a class so existing
    #: downloads keep being found.
    DIR_NAME: ClassVar[str | None] = None
    #: Approximate size of what download() fetches, for `deepecg list-datasets`.
    DOWNLOAD_SIZE: ClassVar[str | None] = None

    binary_classification: bool = False
    #: Re-raise per-record load errors instead of skipping the record.
    strict: bool = False

    #: Exception types treated as a corrupt/unreadable source record.
    RECORD_ERRORS: ClassVar[tuple[type[Exception], ...]] = (
        OSError,
        ValueError,
        KeyError,
        IndexError,
        NotImplementedError,
    )

    @classmethod
    def get_default_data_dir(cls) -> Path:
        """Get the default data directory for this dataset.

        Returns:
            ``<data home>/datasets/<name>``, where the data home is
            ``~/.deepecgkit`` unless the DEEPECGKIT_HOME environment variable
            points elsewhere.
        """
        return get_data_home() / "datasets" / (cls.DIR_NAME or cls.__name__.lower())

    @classmethod
    def is_downloaded(cls, data_dir: str | Path | None = None) -> bool:
        """Report whether a usable copy of the dataset is present.

        The default treats any non-empty directory as the dataset. Datasets
        that download archives override this to look for the files their
        loader needs, so an interrupted download is detected and repeated
        rather than mistaken for the dataset.

        Args:
            data_dir: Directory to inspect; the default data directory when None

        Returns:
            True when the dataset can be loaded from data_dir as it stands.
        """
        target = Path(data_dir) if data_dir is not None else cls.get_default_data_dir()
        return target.is_dir() and any(target.iterdir())

    @classmethod
    def _download_records(
        cls,
        data_dir: Path,
        names: list[str] | Callable[[], list[str]],
        *,
        mirror: str,
        desc: str,
        archive_url: str,
        archive_name: str,
        workers: int = 1,
        keep_alive: bool = True,
        verbose: bool = True,
        **archive_kwargs: Any,
    ) -> None:
        """Fetch a dataset's files from a mirror, or its archive when that fails.

        PhysioNet's own servers hand out whole-project archives slowly; the
        open-data mirror serves the same files individually and much faster,
        so it is tried first. The archive stays as the fallback because it
        needs nothing but physionet.org to be reachable.

        Args:
            data_dir: Directory to download into
            names: File paths relative to the mirror's base URL, or a function
                returning them once it has fetched whatever index lists them
            mirror: Base URL of the mirrored project directory
            desc: Description for the progress bars
            archive_url: URL of the whole-project zip archive
            archive_name: File name to save the archive under
            workers: Files fetched at a time from the mirror
            keep_alive: Reuse connections across files; right for many small
                files, wrong for large ones that should resume individually
            verbose: Print progress messages
            **archive_kwargs: Forwarded to download_and_extract_zip
        """
        archive = data_dir / archive_name
        with marking_download(data_dir):
            try:
                resolved = names() if callable(names) else names
                download_files(
                    [mirror],
                    resolved,
                    data_dir,
                    desc=desc,
                    workers=workers,
                    keep_alive=keep_alive,
                )
                # An archive download abandoned on an earlier run would
                # otherwise keep the directory looking unfinished.
                for leftover in (archive, partial_path(archive)):
                    leftover.unlink(missing_ok=True)
                return
            except (RuntimeError, OSError, ValueError, KeyError) as error:
                logger.warning(
                    f"The mirror at {mirror} could not be used ({error}). "
                    f"Falling back to the archive on physionet.org, which is slower."
                )
            download_and_extract_zip(
                archive_url,
                data_dir,
                zip_name=archive_name,
                desc=desc,
                verbose=verbose,
                **archive_kwargs,
            )

    @classmethod
    def _not_found_message(cls, data_dir: Path, expected: str) -> str:
        """Build the error text for a dataset that is absent or incomplete.

        Args:
            data_dir: Directory that was searched
            expected: What the loader looked for there

        Returns:
            A message naming the directory and the ways to fix it.
        """
        return (
            f"{cls.__name__} data was not found in {data_dir} (expected {expected}). "
            f"Pass download=True to fetch it, or force_download=True to replace an "
            f"incomplete copy (CLI: --force-download). To keep datasets elsewhere, pass "
            f"data_dir or set the {DATA_HOME_ENV} environment variable."
        )

    def __init__(
        self,
        data_dir: str | Path | None = None,
        sampling_rate: int = 500,
        leads: list[str] | None = None,
        transform: Callable[[torch.Tensor], torch.Tensor] | None = None,
        target_transform: Callable[[torch.Tensor], torch.Tensor] | None = None,
        download: bool = False,
        force_download: bool = False,
    ):
        """Initialize the base ECG dataset.

        Args:
            data_dir: Directory where the dataset is stored or will be downloaded.
                     If None, uses the default data directory.
            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 locally.
            force_download: Whether to force re-download even if the dataset exists
        """
        self.data_dir = Path(data_dir) if data_dir is not None else self.get_default_data_dir()
        self.sampling_rate = sampling_rate
        self.leads = leads
        self.transform = transform
        self.target_transform = target_transform

        if force_download:
            with download_lock(self.data_dir):
                self._clear_dataset()
                self.download()
        elif download and not self.is_downloaded(self.data_dir):
            with download_lock(self.data_dir):
                # Another process may have fetched it while this one waited.
                if not self.is_downloaded(self.data_dir):
                    self.download()

        self._load_data()

    def _data_dir_is_empty(self) -> bool:
        """True when the data directory is missing or holds no entries."""
        return not self.data_dir.exists() or not any(self.data_dir.iterdir())

    @classmethod
    def _require_clearable(cls, directory: Path) -> None:
        """Refuse to delete a directory that is not recognisably this dataset.

        A forced download removes the whole directory first. The default
        directory is the package's own, but a data_dir of the caller's
        choosing may hold anything, so it is only cleared when it is empty, a
        complete copy of the dataset, or a download this package began.

        Raises:
            RuntimeError: If the directory holds something else.
        """
        recognised = (
            directory.resolve() == cls.get_default_data_dir().resolve()
            or not any(directory.iterdir())
            or (directory / IN_PROGRESS_MARKER).exists()
            or cls.is_downloaded(directory)
        )
        if not recognised:
            raise RuntimeError(
                f"Refusing to clear {directory} for a forced download: it is not empty and "
                f"does not look like a copy of {cls.__name__}, and everything in it would be "
                f"deleted. Remove it yourself, or choose an empty data directory."
            )

    def _clear_dataset(self) -> None:
        """Clear the dataset directory for re-download."""
        if self.data_dir.exists():
            self._require_clearable(self.data_dir)
            logger.info(f"Clearing existing dataset at {self.data_dir}...")
            shutil.rmtree(self.data_dir)
        self.data_dir.mkdir(parents=True, exist_ok=True)

    @classmethod
    def _download(cls, data_dir: Path, verbose: bool = True) -> None:
        """Fetch the dataset archive into data_dir. Implemented per dataset."""
        raise NotImplementedError(
            f"{cls.__name__} does not implement _download(), so it cannot be fetched "
            f"automatically. Place the data in {data_dir} and pass download=False."
        )

    @classmethod
    def download_to(cls, data_dir: str | Path | None = None, verbose: bool = True) -> Path:
        """Download the dataset without constructing (and loading) an instance.

        Args:
            data_dir: Target directory; the default data directory when None
            verbose: Print progress messages

        Returns:
            The directory the dataset was downloaded into.
        """
        target = Path(data_dir) if data_dir is not None else cls.get_default_data_dir()
        with download_lock(target):
            cls._download(target, verbose=verbose)
        return target

    def download(self) -> None:
        """Download the dataset into this instance's data directory."""
        self._download(self.data_dir, verbose=getattr(self, "verbose", True))

    @abstractmethod
    def _load_data(self) -> None:
        """Load the dataset into memory."""

    def _record_load_error(self, record_name: str, error: Exception) -> None:
        """Note a source record that failed to load; re-raise in strict mode.

        Args:
            record_name: Name of the failing record
            error: The exception it raised
        """
        if self.strict:
            raise error
        if not hasattr(self, "_load_errors"):
            self._load_errors: list[tuple[str, Exception]] = []
        self._load_errors.append((record_name, error))
        if getattr(self, "verbose", True):
            logger.warning(f"Error loading {record_name}: {error}")

    def _finish_load(self, loaded_records: int, total_records: int) -> None:
        """Surface skipped records loudly once loading completes.

        A partially loaded dataset must never look like a fully loaded one:
        skips produce an unconditional warning (independent of verbose), and
        a load in which every record failed raises.

        Args:
            loaded_records: Number of records that loaded successfully
            total_records: Number of records that were attempted
        """
        errors = getattr(self, "_load_errors", [])
        if errors:
            first_name, first_error = errors[0]
            warnings.warn(
                f"Skipped {len(errors)}/{total_records} records while loading "
                f"{type(self).__name__}; first error ({first_name}): {first_error}. "
                f"Pass strict=True to raise instead.",
                RuntimeWarning,
                stacklevel=3,
            )
        if total_records and loaded_records == 0 and errors:
            raise RuntimeError(
                f"All {total_records} records failed to load for {type(self).__name__}; "
                f"first error ({errors[0][0]}): {errors[0][1]}"
            )

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
        """Get a sample from the dataset.

        Args:
            idx: Index of the sample to get

        Returns:
            Tuple of (ecg_signal, label) where ecg_signal has shape
            (num_leads, signal_length).
        """
        signal = convert_to_tensor(self.signals[idx], dtype=torch.float32)
        label = torch.tensor(self.labels[idx], dtype=torch.long)

        if self.transform is not None:
            signal = self.transform(signal)
        if self.target_transform is not None:
            label = self.target_transform(label)

        return signal, label

    def __len__(self) -> int:
        """Get the number of samples in the dataset."""
        return len(self.labels)

    @property
    def num_classes(self) -> int:
        """Get the number of classes in the dataset."""
        return len(self.class_names)

    @property
    def class_names(self) -> list[str]:
        """Get the names of the classes in the dataset."""
        if self.binary_classification:
            return self.BINARY_CLASS_LABELS
        return self.CLASS_LABELS

    @property
    def group_ids(self) -> list[str]:
        """Identifiers used to keep correlated samples inside a single split.

        Defaults to record_names, which is correct whenever one record corresponds
        to one subject. Datasets where a single subject contributes several records
        must override this to return subject-level identifiers, otherwise samples
        from the same subject can be divided across train and test.

        Returns:
            One identifier per sample, or an empty list if grouping is unavailable.
        """
        return list(getattr(self, "record_names", []) or [])

    def get_record_info(self, idx: int) -> dict:
        """Describe one sample: source record, label, class name, and shape."""
        label = int(self.labels[idx])
        return {
            "record_name": self.record_names[idx],
            "label": label,
            "class_name": self.class_names[label],
            "signal_shape": self.signals[idx].shape,
        }

    def stratification_labels(self) -> np.ndarray | None:
        """Return every sample's label as __getitem__ gives it, without reading a signal.

        ECGDataModule stratifies its split on these. Collecting them through
        __getitem__ reads every signal on the way, which for a memory-mapped
        dataset of several gigabytes means paging all of it in.

        Returns:
            Integer labels, one per sample, or None when they cannot be had
            without reading the samples: a target_transform is set, or the
            subclass builds its samples in a __getitem__ of its own.
        """
        labels = getattr(self, "labels", None)
        if (
            labels is None
            or self.target_transform is not None
            or type(self).__getitem__ is not BaseECGDataset.__getitem__
        ):
            return None
        return np.asarray(labels).astype(np.int64)

    def get_class_distribution(self) -> dict[str, int]:
        """Count samples per class name."""
        labels = getattr(self, "labels", None)
        if labels is None or len(labels) == 0:
            return {}
        unique, counts = np.unique(labels, return_counts=True)
        return {
            self.class_names[int(label)]: int(count)
            for label, count in zip(unique, counts, strict=False)
        }

    def _print_class_distribution(self) -> None:
        """Print class distribution statistics."""
        distribution = self.get_class_distribution()
        if not distribution:
            return
        logger.info("\nClass distribution:")
        for class_name, count in distribution.items():
            logger.info(f"  {class_name}: {count}")

    def get_preprocessing(self) -> dict[str, Any]:
        """Describe how a raw recording becomes one of this dataset's samples.

        Checkpoints record this (see ECGTrainer.data_config) so that a new
        recording can be put through the same steps before it reaches the
        model: resample to sampling_rate, cut into segments of
        segment_duration_seconds, then normalise each segment.

        Returns:
            Mapping with sampling_rate, segment_duration_seconds (None when
            recordings are not segmented), normalization, and leads.
        """
        standardizer = getattr(self, "standardizer", None)
        return {
            "sampling_rate": self.sampling_rate,
            "segment_duration_seconds": getattr(self, "segment_duration_seconds", None),
            "normalization": getattr(standardizer, "normalization", None),
            "leads": list(self.leads) if self.leads else None,
        }

    def get_metadata(self) -> dict:
        """Get metadata about the dataset.

        Returns:
            Dictionary containing metadata such as:
            - sampling_rate: Sampling rate of the signals
            - num_leads: Number of leads
            - lead_names: Names of the leads
            - num_classes: Number of classes
            - class_names: Names of the classes
            - dataset_size: Number of samples
        """
        return {
            "sampling_rate": self.sampling_rate,
            "num_leads": len(self.leads) if self.leads else None,
            "lead_names": self.leads,
            "num_classes": self.num_classes,
            "class_names": self.class_names,
            "dataset_size": len(self),
        }

num_classes property

num_classes: int

Get the number of classes in the dataset.

class_names property

class_names: list[str]

Get the names of the classes in the dataset.

group_ids property

group_ids: list[str]

Identifiers used to keep correlated samples inside a single split.

Defaults to record_names, which is correct whenever one record corresponds to one subject. Datasets where a single subject contributes several records must override this to return subject-level identifiers, otherwise samples from the same subject can be divided across train and test.

Returns:

Type	Description
list[str]	One identifier per sample, or an empty list if grouping is unavailable.

get_default_data_dir classmethod

get_default_data_dir() -> Path

Get the default data directory for this dataset.

Returns:

Type	Description
Path	<data home>/datasets/<name>, where the data home is
Path	~/.deepecgkit unless the DEEPECGKIT_HOME environment variable
Path	points elsewhere.

Source code in deepecgkit/datasets/base.py

@classmethod
def get_default_data_dir(cls) -> Path:
    """Get the default data directory for this dataset.

    Returns:
        ``<data home>/datasets/<name>``, where the data home is
        ``~/.deepecgkit`` unless the DEEPECGKIT_HOME environment variable
        points elsewhere.
    """
    return get_data_home() / "datasets" / (cls.DIR_NAME or cls.__name__.lower())

is_downloaded classmethod

is_downloaded(data_dir: str | Path | None = None) -> bool

Report whether a usable copy of the dataset is present.

The default treats any non-empty directory as the dataset. Datasets that download archives override this to look for the files their loader needs, so an interrupted download is detected and repeated rather than mistaken for the dataset.

Parameters:

Name	Type	Description	Default
data_dir	str | Path | None	Directory to inspect; the default data directory when None	None

Returns:

Type	Description
bool	True when the dataset can be loaded from data_dir as it stands.

Source code in deepecgkit/datasets/base.py

@classmethod
def is_downloaded(cls, data_dir: str | Path | None = None) -> bool:
    """Report whether a usable copy of the dataset is present.

    The default treats any non-empty directory as the dataset. Datasets
    that download archives override this to look for the files their
    loader needs, so an interrupted download is detected and repeated
    rather than mistaken for the dataset.

    Args:
        data_dir: Directory to inspect; the default data directory when None

    Returns:
        True when the dataset can be loaded from data_dir as it stands.
    """
    target = Path(data_dir) if data_dir is not None else cls.get_default_data_dir()
    return target.is_dir() and any(target.iterdir())

download_to classmethod

download_to(
    data_dir: str | Path | None = None, verbose: bool = True
) -> Path

Download the dataset without constructing (and loading) an instance.

Parameters:

Name	Type	Description	Default
data_dir	str | Path | None	Target directory; the default data directory when None	None
verbose	bool	Print progress messages	True

Returns:

Type	Description
Path	The directory the dataset was downloaded into.

Source code in deepecgkit/datasets/base.py

@classmethod
def download_to(cls, data_dir: str | Path | None = None, verbose: bool = True) -> Path:
    """Download the dataset without constructing (and loading) an instance.

    Args:
        data_dir: Target directory; the default data directory when None
        verbose: Print progress messages

    Returns:
        The directory the dataset was downloaded into.
    """
    target = Path(data_dir) if data_dir is not None else cls.get_default_data_dir()
    with download_lock(target):
        cls._download(target, verbose=verbose)
    return target

download

download() -> None

Download the dataset into this instance's data directory.

Source code in deepecgkit/datasets/base.py

def download(self) -> None:
    """Download the dataset into this instance's data directory."""
    self._download(self.data_dir, verbose=getattr(self, "verbose", True))

get_record_info

get_record_info(idx: int) -> dict

Describe one sample: source record, label, class name, and shape.

Source code in deepecgkit/datasets/base.py

def get_record_info(self, idx: int) -> dict:
    """Describe one sample: source record, label, class name, and shape."""
    label = int(self.labels[idx])
    return {
        "record_name": self.record_names[idx],
        "label": label,
        "class_name": self.class_names[label],
        "signal_shape": self.signals[idx].shape,
    }

stratification_labels

stratification_labels() -> np.ndarray | None

Return every sample's label as getitem gives it, without reading a signal.

ECGDataModule stratifies its split on these. Collecting them through getitem reads every signal on the way, which for a memory-mapped dataset of several gigabytes means paging all of it in.

Returns:

Type	Description
ndarray | None	Integer labels, one per sample, or None when they cannot be had
ndarray | None	without reading the samples: a target_transform is set, or the
ndarray | None	subclass builds its samples in a getitem of its own.

Source code in deepecgkit/datasets/base.py

def stratification_labels(self) -> np.ndarray | None:
    """Return every sample's label as __getitem__ gives it, without reading a signal.

    ECGDataModule stratifies its split on these. Collecting them through
    __getitem__ reads every signal on the way, which for a memory-mapped
    dataset of several gigabytes means paging all of it in.

    Returns:
        Integer labels, one per sample, or None when they cannot be had
        without reading the samples: a target_transform is set, or the
        subclass builds its samples in a __getitem__ of its own.
    """
    labels = getattr(self, "labels", None)
    if (
        labels is None
        or self.target_transform is not None
        or type(self).__getitem__ is not BaseECGDataset.__getitem__
    ):
        return None
    return np.asarray(labels).astype(np.int64)

get_class_distribution

get_class_distribution() -> dict[str, int]

Count samples per class name.

Source code in deepecgkit/datasets/base.py

def get_class_distribution(self) -> dict[str, int]:
    """Count samples per class name."""
    labels = getattr(self, "labels", None)
    if labels is None or len(labels) == 0:
        return {}
    unique, counts = np.unique(labels, return_counts=True)
    return {
        self.class_names[int(label)]: int(count)
        for label, count in zip(unique, counts, strict=False)
    }

get_preprocessing

get_preprocessing() -> dict[str, Any]

Describe how a raw recording becomes one of this dataset's samples.

Checkpoints record this (see ECGTrainer.data_config) so that a new recording can be put through the same steps before it reaches the model: resample to sampling_rate, cut into segments of segment_duration_seconds, then normalise each segment.

Returns:

Type	Description
dict[str, Any]	Mapping with sampling_rate, segment_duration_seconds (None when
dict[str, Any]	recordings are not segmented), normalization, and leads.

Source code in deepecgkit/datasets/base.py

def get_preprocessing(self) -> dict[str, Any]:
    """Describe how a raw recording becomes one of this dataset's samples.

    Checkpoints record this (see ECGTrainer.data_config) so that a new
    recording can be put through the same steps before it reaches the
    model: resample to sampling_rate, cut into segments of
    segment_duration_seconds, then normalise each segment.

    Returns:
        Mapping with sampling_rate, segment_duration_seconds (None when
        recordings are not segmented), normalization, and leads.
    """
    standardizer = getattr(self, "standardizer", None)
    return {
        "sampling_rate": self.sampling_rate,
        "segment_duration_seconds": getattr(self, "segment_duration_seconds", None),
        "normalization": getattr(standardizer, "normalization", None),
        "leads": list(self.leads) if self.leads else None,
    }

get_metadata

get_metadata() -> dict

Get metadata about the dataset.

Returns:

Type	Description
dict	Dictionary containing metadata such as:
dict	sampling_rate: Sampling rate of the signals
dict	num_leads: Number of leads
dict	lead_names: Names of the leads
dict	num_classes: Number of classes
dict	class_names: Names of the classes
dict	dataset_size: Number of samples

Source code in deepecgkit/datasets/base.py

def get_metadata(self) -> dict:
    """Get metadata about the dataset.

    Returns:
        Dictionary containing metadata such as:
        - sampling_rate: Sampling rate of the signals
        - num_leads: Number of leads
        - lead_names: Names of the leads
        - num_classes: Number of classes
        - class_names: Names of the classes
        - dataset_size: Number of samples
    """
    return {
        "sampling_rate": self.sampling_rate,
        "num_leads": len(self.leads) if self.leads else None,
        "lead_names": self.leads,
        "num_classes": self.num_classes,
        "class_names": self.class_names,
        "dataset_size": len(self),
    }

Data Module

ECGDataModule

Data module for ECG datasets.

Handles dataset creation, train/val/test splitting, and DataLoader construction.

Source code in deepecgkit/datasets/modules.py

class ECGDataModule:
    """Data module for ECG datasets.

    Handles dataset creation, train/val/test splitting, and DataLoader construction.
    """

    def __init__(
        self,
        dataset: BaseECGDataset | Dataset | None = None,
        dataset_class: type[BaseECGDataset] | None = None,
        data_dir: str | None = None,
        sampling_rate: int | None = None,
        leads: list | None = None,
        transform: Callable[[torch.Tensor], torch.Tensor] | None = None,
        target_transform: Callable[[torch.Tensor], torch.Tensor] | None = None,
        download: bool = False,
        batch_size: int = 32,
        val_split: float = 0.2,
        test_split: float = 0.1,
        num_workers: int = 0,
        seed: int = 42,
        stratify: bool = True,
        pin_memory: bool = False,
        persistent_workers: bool = True,
        prefetch_factor: int = 2,
        verbose: bool = True,
        dataset_kwargs: dict | None = 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: Rate the dataset resamples its signals to (Hz). None keeps the
                dataset class's own default, which is its native or recommended rate;
                after setup() the attribute holds the rate actually in effect
            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. The built-in
                datasets hold their signals in memory, so the default of 0 is the
                fastest setting; with more, guard the calling script with
                ``if __name__ == "__main__":`` on macOS and Windows
            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
        """
        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
        self._split_settings: tuple | None = None

    def setup(self, stage: str | None = None) -> None:
        """Set up the datasets.

        The dataset is loaded and split on the first call. Later calls, which
        the trainer makes for every stage, return at once unless the dataset or
        a split setting has changed since.

        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,
                "transform": self.transform,
                "target_transform": self.target_transform,
                "download": self.download,
            }

            # Only an explicit rate overrides the dataset's own default: a
            # blanket value here would silently resample every dataset.
            if self.sampling_rate is not None:
                init_kwargs["sampling_rate"] = self.sampling_rate

            if self.leads is not None:
                init_kwargs["leads"] = self.leads

            init_kwargs.update(self.dataset_kwargs)

            self.dataset = self.dataset_class(**init_kwargs)

        effective_rate = getattr(self.dataset, "sampling_rate", None)
        if effective_rate is not None:
            self.sampling_rate = effective_rate

        split_settings = (
            id(self.dataset),
            self.val_split,
            self.test_split,
            self.seed,
            self.stratify,
        )
        if self.train_dataset is not None and split_settings == self._split_settings:
            return
        self._split_settings = split_settings

        stratify_labels = None
        if (
            self.stratify and len(self.dataset) >= 8
        ):  # Need at least 2 samples per class for 4 classes
            try:
                stacked = self._collect_labels()

                # 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:
                    warnings.warn(
                        "Stratified splitting disabled: at least one class has fewer "
                        "than 2 samples.",
                        RuntimeWarning,
                        stacklevel=2,
                    )
                    stratify_labels = None
            except (RuntimeError, ValueError, TypeError, IndexError) as error:
                warnings.warn(
                    f"Stratified splitting disabled: could not collect labels ({error}).",
                    RuntimeWarning,
                    stacklevel=2,
                )
                stratify_labels = None

        # Extract groups for patient-level splitting (prevents data leakage)
        groups = None
        group_ids = getattr(self.dataset, "group_ids", None) or getattr(
            self.dataset, "record_names", None
        )
        if group_ids:
            if len(group_ids) != len(self.dataset):
                raise ValueError(
                    f"{type(self.dataset).__name__} exposes {len(group_ids)} group ids "
                    f"for {len(self.dataset)} samples. Grouped splitting cannot be "
                    "performed safely."
                )
            groups = np.array(group_ids)
            if self.verbose:
                n_groups = len(np.unique(groups))
                logger.info(f"Using grouped splitting ({n_groups} groups)")
        else:
            warnings.warn(
                f"{type(self.dataset).__name__} exposes no group ids, so the split is "
                "performed at sample level. If samples are segments of longer recordings, "
                "segments of the same recording will appear in both train and test and the "
                "reported scores will be optimistically biased.",
                RuntimeWarning,
                stacklevel=2,
            )

        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:
            logger.info(f"Training set size: {len(self.train_dataset)}")
            logger.info(f"Validation set size: {len(self.val_dataset)}")
            logger.info(f"Test set size: {len(self.test_dataset)}")

    def _collect_labels(self) -> np.ndarray:
        """Collect every sample's label, for stratifying the split.

        A dataset that can hand its labels over directly is asked first;
        otherwise every sample is read for its label.
        """
        size = len(self.dataset)
        direct = getattr(self.dataset, "stratification_labels", None)
        labels: np.ndarray | None = direct() if callable(direct) else None
        if labels is None or len(labels) != size:
            labels = torch.stack([self.dataset[i][1] for i in range(size)]).numpy()
        return labels

    def _loader(self, dataset: Any, shuffle: bool) -> DataLoader:
        """Build a loader over one split.

        Raises:
            RuntimeError: If setup() has not been called.
        """
        if dataset is None:
            raise RuntimeError("Call setup() before using the data module")
        options: dict[str, Any] = {}
        if self.num_workers > 0:
            # These two only mean something with worker processes, and PyTorch
            # before 2.0 rejects a prefetch_factor given without any.
            options["persistent_workers"] = self.persistent_workers
            options["prefetch_factor"] = self.prefetch_factor
        return DataLoader(
            dataset,
            batch_size=self.batch_size,
            shuffle=shuffle,
            num_workers=self.num_workers,
            pin_memory=self.pin_memory,
            **options,
        )

    def train_dataloader(self) -> DataLoader:
        """Get the training data loader."""
        return self._loader(self.train_dataset, shuffle=True)

    def val_dataloader(self) -> DataLoader:
        """Get the validation data loader."""
        return self._loader(self.val_dataset, shuffle=False)

    def test_dataloader(self) -> DataLoader:
        """Get the test data loader."""
        return self._loader(self.test_dataset, shuffle=False)

    def describe(self) -> dict[str, Any]:
        """Describe the data a model is trained on, in the form checkpoints record.

        ECGTrainer stores this as its data_config, which is what lets a
        checkpoint be used on new recordings: it names the preprocessing the
        inputs need and the classes the outputs stand for.

        Returns:
            Mapping with the registered dataset name, its preprocessing
            (sampling rate, segment length, normalisation), the class names,
            the sample shape, the labelling mode, and the split protocol.
        """
        if self.dataset is None:
            raise RuntimeError("Call setup() before describing the data module")
        dataset = self.dataset

        input_channels = input_length = None
        if len(dataset) > 0:
            sample = dataset[0][0]
            if getattr(sample, "ndim", 0) >= 2:
                input_channels, input_length = int(sample.shape[0]), int(sample.shape[-1])

        get_preprocessing = getattr(dataset, "get_preprocessing", None)
        class_names = getattr(dataset, "class_names", None)
        return {
            "dataset": getattr(type(dataset), "_registry_name", None),
            "preprocessing": get_preprocessing() if callable(get_preprocessing) else None,
            "class_names": list(class_names) if class_names else None,
            "input_channels": input_channels,
            "input_length": input_length,
            "multi_label": bool(getattr(dataset, "multi_label", False)),
            "binary_classification": bool(getattr(dataset, "binary_classification", False)),
            "split": {
                "seed": self.seed,
                "val_split": self.val_split,
                "test_split": self.test_split,
            },
        }

    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) -> None:
        """Print metadata about the dataset."""
        print(self.get_metadata())

setup

setup(stage: str | None = None) -> None

Set up the datasets.

The dataset is loaded and split on the first call. Later calls, which the trainer makes for every stage, return at once unless the dataset or a split setting has changed since.

Parameters:

Name	Type	Description	Default
stage	str | None	Current stage ('fit', 'validate', 'test', or 'predict')	None

Source code in deepecgkit/datasets/modules.py

def setup(self, stage: str | None = None) -> None:
    """Set up the datasets.

    The dataset is loaded and split on the first call. Later calls, which
    the trainer makes for every stage, return at once unless the dataset or
    a split setting has changed since.

    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,
            "transform": self.transform,
            "target_transform": self.target_transform,
            "download": self.download,
        }

        # Only an explicit rate overrides the dataset's own default: a
        # blanket value here would silently resample every dataset.
        if self.sampling_rate is not None:
            init_kwargs["sampling_rate"] = self.sampling_rate

        if self.leads is not None:
            init_kwargs["leads"] = self.leads

        init_kwargs.update(self.dataset_kwargs)

        self.dataset = self.dataset_class(**init_kwargs)

    effective_rate = getattr(self.dataset, "sampling_rate", None)
    if effective_rate is not None:
        self.sampling_rate = effective_rate

    split_settings = (
        id(self.dataset),
        self.val_split,
        self.test_split,
        self.seed,
        self.stratify,
    )
    if self.train_dataset is not None and split_settings == self._split_settings:
        return
    self._split_settings = split_settings

    stratify_labels = None
    if (
        self.stratify and len(self.dataset) >= 8
    ):  # Need at least 2 samples per class for 4 classes
        try:
            stacked = self._collect_labels()

            # 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:
                warnings.warn(
                    "Stratified splitting disabled: at least one class has fewer "
                    "than 2 samples.",
                    RuntimeWarning,
                    stacklevel=2,
                )
                stratify_labels = None
        except (RuntimeError, ValueError, TypeError, IndexError) as error:
            warnings.warn(
                f"Stratified splitting disabled: could not collect labels ({error}).",
                RuntimeWarning,
                stacklevel=2,
            )
            stratify_labels = None

    # Extract groups for patient-level splitting (prevents data leakage)
    groups = None
    group_ids = getattr(self.dataset, "group_ids", None) or getattr(
        self.dataset, "record_names", None
    )
    if group_ids:
        if len(group_ids) != len(self.dataset):
            raise ValueError(
                f"{type(self.dataset).__name__} exposes {len(group_ids)} group ids "
                f"for {len(self.dataset)} samples. Grouped splitting cannot be "
                "performed safely."
            )
        groups = np.array(group_ids)
        if self.verbose:
            n_groups = len(np.unique(groups))
            logger.info(f"Using grouped splitting ({n_groups} groups)")
    else:
        warnings.warn(
            f"{type(self.dataset).__name__} exposes no group ids, so the split is "
            "performed at sample level. If samples are segments of longer recordings, "
            "segments of the same recording will appear in both train and test and the "
            "reported scores will be optimistically biased.",
            RuntimeWarning,
            stacklevel=2,
        )

    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:
        logger.info(f"Training set size: {len(self.train_dataset)}")
        logger.info(f"Validation set size: {len(self.val_dataset)}")
        logger.info(f"Test set size: {len(self.test_dataset)}")

train_dataloader

train_dataloader() -> DataLoader

Get the training data loader.

Source code in deepecgkit/datasets/modules.py

def train_dataloader(self) -> DataLoader:
    """Get the training data loader."""
    return self._loader(self.train_dataset, shuffle=True)

val_dataloader

val_dataloader() -> DataLoader

Get the validation data loader.

Source code in deepecgkit/datasets/modules.py

def val_dataloader(self) -> DataLoader:
    """Get the validation data loader."""
    return self._loader(self.val_dataset, shuffle=False)

test_dataloader

test_dataloader() -> DataLoader

Get the test data loader.

Source code in deepecgkit/datasets/modules.py

def test_dataloader(self) -> DataLoader:
    """Get the test data loader."""
    return self._loader(self.test_dataset, shuffle=False)

describe

describe() -> dict[str, Any]

Describe the data a model is trained on, in the form checkpoints record.

ECGTrainer stores this as its data_config, which is what lets a checkpoint be used on new recordings: it names the preprocessing the inputs need and the classes the outputs stand for.

Returns:

Type	Description
dict[str, Any]	Mapping with the registered dataset name, its preprocessing
dict[str, Any]	(sampling rate, segment length, normalisation), the class names,
dict[str, Any]	the sample shape, the labelling mode, and the split protocol.

Source code in deepecgkit/datasets/modules.py

def describe(self) -> dict[str, Any]:
    """Describe the data a model is trained on, in the form checkpoints record.

    ECGTrainer stores this as its data_config, which is what lets a
    checkpoint be used on new recordings: it names the preprocessing the
    inputs need and the classes the outputs stand for.

    Returns:
        Mapping with the registered dataset name, its preprocessing
        (sampling rate, segment length, normalisation), the class names,
        the sample shape, the labelling mode, and the split protocol.
    """
    if self.dataset is None:
        raise RuntimeError("Call setup() before describing the data module")
    dataset = self.dataset

    input_channels = input_length = None
    if len(dataset) > 0:
        sample = dataset[0][0]
        if getattr(sample, "ndim", 0) >= 2:
            input_channels, input_length = int(sample.shape[0]), int(sample.shape[-1])

    get_preprocessing = getattr(dataset, "get_preprocessing", None)
    class_names = getattr(dataset, "class_names", None)
    return {
        "dataset": getattr(type(dataset), "_registry_name", None),
        "preprocessing": get_preprocessing() if callable(get_preprocessing) else None,
        "class_names": list(class_names) if class_names else None,
        "input_channels": input_channels,
        "input_length": input_length,
        "multi_label": bool(getattr(dataset, "multi_label", False)),
        "binary_classification": bool(getattr(dataset, "binary_classification", False)),
        "split": {
            "seed": self.seed,
            "val_split": self.val_split,
            "test_split": self.test_split,
        },
    }

get_metadata

get_metadata() -> dict

Get metadata about the dataset.

Returns:

Type	Description
dict	Dictionary containing dataset metadata

Source code in deepecgkit/datasets/modules.py

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,
    }

print_metadata

print_metadata() -> None

Print metadata about the dataset.

Source code in deepecgkit/datasets/modules.py

def print_metadata(self) -> None:
    """Print metadata about the dataset."""
    print(self.get_metadata())

Data Splitting

DataSplitter

Handles dataset splitting into train, validation, and test sets.

Source code in deepecgkit/datasets/splitting.py

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: int | None = 42,
        stratify: np.ndarray | None = None,
        groups: np.ndarray | None = 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 as error:
                warnings.warn(
                    f"Stratified splitting failed ({error}); falling back to random splitting.",
                    RuntimeWarning,
                    stacklevel=2,
                )

        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, strict=False):
                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:
                warnings.warn(
                    "Grouped split is not stratified: a label is the majority label of fewer "
                    "than two groups.",
                    RuntimeWarning,
                    stacklevel=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:
                        warnings.warn(
                            "Validation/test split of the held-out groups is not stratified: "
                            "a label is the majority label of fewer than two held-out groups.",
                            RuntimeWarning,
                            stacklevel=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 as error:
            warnings.warn(
                f"Grouped stratified split failed ({error}); splitting without stratification.",
                RuntimeWarning,
                stacklevel=2,
            )
            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, []),
        )

split

split() -> tuple[Dataset, Dataset, Dataset]

Split the dataset into train, validation, and test sets.

Returns:

Type	Description
tuple[Dataset, Dataset, Dataset]	Tuple of (train_dataset, val_dataset, test_dataset)

Source code in deepecgkit/datasets/splitting.py

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 as error:
            warnings.warn(
                f"Stratified splitting failed ({error}); falling back to random splitting.",
                RuntimeWarning,
                stacklevel=2,
            )

    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

Preprocessing

ECGStandardizer

Resamples, length-standardises, and normalises raw ECG signals.

Parameters:

Name	Type	Description	Default
target_sampling_rate	int	Sampling rate signals are resampled to (Hz)	300
target_length	int | None	Fixed output length in samples; mutually exclusive with target_duration_seconds, None leaves lengths unchanged	None
target_duration_seconds	float | None	Fixed output duration; converted to samples at the target rate	None
normalization	str	"zscore", "minmax", or "none"	'zscore'
clip_method	str	Where clipping/padding takes place: "center", "start", or "end"	'center'

Source code in deepecgkit/datasets/preprocessing.py

class ECGStandardizer:
    """Resamples, length-standardises, and normalises raw ECG signals.

    Args:
        target_sampling_rate: Sampling rate signals are resampled to (Hz)
        target_length: Fixed output length in samples; mutually exclusive
            with target_duration_seconds, None leaves lengths unchanged
        target_duration_seconds: Fixed output duration; converted to samples
            at the target rate
        normalization: "zscore", "minmax", or "none"
        clip_method: Where clipping/padding takes place: "center", "start",
            or "end"
    """

    def __init__(
        self,
        target_sampling_rate: int = 300,
        target_length: int | None = None,
        target_duration_seconds: float | None = 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:
        """Resample a (leads, samples) or 1-D signal to the target rate."""
        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:
        """Normalise per lead using the configured method."""
        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:
        """Bring the signal to target_length by clipping or zero-padding."""
        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:
        """Apply resample, clip_or_pad, and normalize in order."""
        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

resample

resample(
    ecg_signal: ndarray, original_sampling_rate: int
) -> np.ndarray

Resample a (leads, samples) or 1-D signal to the target rate.

Source code in deepecgkit/datasets/preprocessing.py

def resample(self, ecg_signal: np.ndarray, original_sampling_rate: int) -> np.ndarray:
    """Resample a (leads, samples) or 1-D signal to the target rate."""
    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

normalize

normalize(ecg_signal: ndarray) -> np.ndarray

Normalise per lead using the configured method.

Source code in deepecgkit/datasets/preprocessing.py

def normalize(self, ecg_signal: np.ndarray) -> np.ndarray:
    """Normalise per lead using the configured method."""
    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}")

clip_or_pad

clip_or_pad(ecg_signal: ndarray) -> np.ndarray

Bring the signal to target_length by clipping or zero-padding.

Source code in deepecgkit/datasets/preprocessing.py

def clip_or_pad(self, ecg_signal: np.ndarray) -> np.ndarray:
    """Bring the signal to target_length by clipping or zero-padding."""
    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)

ECGSegmenter

Cuts a continuous recording into fixed-length, optionally overlapping segments.

Parameters:

Name	Type	Description	Default
segment_duration_seconds	float	Length of each segment in seconds	required
sampling_rate	int	Sampling rate of the incoming signal (Hz)	required
overlap	float	Fraction of a segment shared with its successor, in [0, 1)	0.0

Source code in deepecgkit/datasets/preprocessing.py

class ECGSegmenter:
    """Cuts a continuous recording into fixed-length, optionally overlapping segments.

    Args:
        segment_duration_seconds: Length of each segment in seconds
        sampling_rate: Sampling rate of the incoming signal (Hz)
        overlap: Fraction of a segment shared with its successor, in [0, 1)
    """

    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]:
        """Cut the signal into segments.

        Args:
            ecg_signal: Signal of shape (leads, samples) or (samples,)

        Returns:
            Tuple of (segments, start indices); both empty when the signal is
            shorter than one segment.
        """
        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)

segment

segment(
    ecg_signal: ndarray,
) -> tuple[np.ndarray, np.ndarray]

Cut the signal into segments.

Parameters:

Name	Type	Description	Default
ecg_signal	ndarray	Signal of shape (leads, samples) or (samples,)	required

Returns:

Type	Description
ndarray	Tuple of (segments, start indices); both empty when the signal is
ndarray	shorter than one segment.

Source code in deepecgkit/datasets/preprocessing.py

def segment(self, ecg_signal: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Cut the signal into segments.

    Args:
        ecg_signal: Signal of shape (leads, samples) or (samples,)

    Returns:
        Tuple of (segments, start indices); both empty when the signal is
        shorter than one segment.
    """
    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)

RhythmAnnotationExtractor

Expands wfdb rhythm annotations into per-sample label timelines.

Parameters:

Name	Type	Description	Default
sampling_rate	int	Rate of the (possibly resampled) signal the labels must align with (Hz)	required
binary_classification	bool	Collapse all rhythms to AF (1) vs Non-AF (0)	False

Source code in deepecgkit/datasets/preprocessing.py

class RhythmAnnotationExtractor:
    """Expands wfdb rhythm annotations into per-sample label timelines.

    Args:
        sampling_rate: Rate of the (possibly resampled) signal the labels
            must align with (Hz)
        binary_classification: Collapse all rhythms to AF (1) vs Non-AF (0)
    """

    RHYTHM_MAP: ClassVar[dict[str, int]] = {
        "(AFIB": 1,
        "(AFL": 2,
        "(J": 3,
        "(N": 0,
    }
    #: What the aux_note of a rhythm-change annotation starts with.
    RHYTHM_PREFIX: ClassVar[str] = "("

    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: Any,
        signal_length: int,
        original_sampling_rate: int | None = None,
    ) -> np.ndarray:
        """Expand a wfdb annotation into one label per signal sample.

        Only rhythm-change annotations, whose aux_note opens a rhythm such as
        ``(AFIB`` or ``(N``, move the label: each one sets the rhythm from its
        own position until the next rhythm change. Some annotation files
        (LTAFDB's among them) also hold an annotation for every beat, with an
        empty aux_note. Those say nothing about the rhythm and are skipped;
        treating them as changes would end every episode at its first beat.

        Args:
            annotation: wfdb annotation with sample and aux_note attributes
            signal_length: Length of the (resampled) signal in samples
            original_sampling_rate: Rate the annotation indices refer to,
                used to rescale them onto the resampled signal

        Returns:
            Integer label per sample, 0 where no rhythm is annotated.
        """
        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

        changes: list[tuple[int, int]] = []
        for sample_idx, aux_note in zip(annotation.sample, annotation.aux_note, strict=False):
            rhythm = aux_note.strip("\x00 \t\r\n")
            if not rhythm.startswith(self.RHYTHM_PREFIX):
                continue

            scaled_idx = int(sample_idx * scale)
            if scaled_idx >= signal_length:
                break

            rhythm_code = self.RHYTHM_MAP.get(rhythm, 0)
            if self.binary_classification:
                rhythm_code = 1 if rhythm_code == 1 else 0
            changes.append((scaled_idx, rhythm_code))

        for position, (start, rhythm_code) in enumerate(changes):
            end = changes[position + 1][0] if position + 1 < len(changes) else signal_length
            labels[start:end] = rhythm_code

        return labels

    def segment_with_labels(
        self, labels: np.ndarray, segment_start_indices: np.ndarray, segment_length: int
    ) -> np.ndarray:
        """Assign each segment the majority label of its samples.

        Args:
            labels: Per-sample label timeline from extract_labels
            segment_start_indices: Segment start positions from ECGSegmenter
            segment_length: Segment length in samples

        Returns:
            One label per segment.
        """
        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)

extract_labels

extract_labels(
    annotation: Any,
    signal_length: int,
    original_sampling_rate: int | None = None,
) -> np.ndarray

Expand a wfdb annotation into one label per signal sample.

Only rhythm-change annotations, whose aux_note opens a rhythm such as (AFIB or (N, move the label: each one sets the rhythm from its own position until the next rhythm change. Some annotation files (LTAFDB's among them) also hold an annotation for every beat, with an empty aux_note. Those say nothing about the rhythm and are skipped; treating them as changes would end every episode at its first beat.

Parameters:

Name	Type	Description	Default
annotation	Any	wfdb annotation with sample and aux_note attributes	required
signal_length	int	Length of the (resampled) signal in samples	required
original_sampling_rate	int | None	Rate the annotation indices refer to, used to rescale them onto the resampled signal	None

Returns:

Type	Description
ndarray	Integer label per sample, 0 where no rhythm is annotated.

Source code in deepecgkit/datasets/preprocessing.py

def extract_labels(
    self,
    annotation: Any,
    signal_length: int,
    original_sampling_rate: int | None = None,
) -> np.ndarray:
    """Expand a wfdb annotation into one label per signal sample.

    Only rhythm-change annotations, whose aux_note opens a rhythm such as
    ``(AFIB`` or ``(N``, move the label: each one sets the rhythm from its
    own position until the next rhythm change. Some annotation files
    (LTAFDB's among them) also hold an annotation for every beat, with an
    empty aux_note. Those say nothing about the rhythm and are skipped;
    treating them as changes would end every episode at its first beat.

    Args:
        annotation: wfdb annotation with sample and aux_note attributes
        signal_length: Length of the (resampled) signal in samples
        original_sampling_rate: Rate the annotation indices refer to,
            used to rescale them onto the resampled signal

    Returns:
        Integer label per sample, 0 where no rhythm is annotated.
    """
    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

    changes: list[tuple[int, int]] = []
    for sample_idx, aux_note in zip(annotation.sample, annotation.aux_note, strict=False):
        rhythm = aux_note.strip("\x00 \t\r\n")
        if not rhythm.startswith(self.RHYTHM_PREFIX):
            continue

        scaled_idx = int(sample_idx * scale)
        if scaled_idx >= signal_length:
            break

        rhythm_code = self.RHYTHM_MAP.get(rhythm, 0)
        if self.binary_classification:
            rhythm_code = 1 if rhythm_code == 1 else 0
        changes.append((scaled_idx, rhythm_code))

    for position, (start, rhythm_code) in enumerate(changes):
        end = changes[position + 1][0] if position + 1 < len(changes) else signal_length
        labels[start:end] = rhythm_code

    return labels

segment_with_labels

segment_with_labels(
    labels: ndarray,
    segment_start_indices: ndarray,
    segment_length: int,
) -> np.ndarray

Assign each segment the majority label of its samples.

Parameters:

Name	Type	Description	Default
labels	ndarray	Per-sample label timeline from extract_labels	required
segment_start_indices	ndarray	Segment start positions from ECGSegmenter	required
segment_length	int	Segment length in samples	required

Returns:

Type	Description
ndarray	One label per segment.

Source code in deepecgkit/datasets/preprocessing.py

def segment_with_labels(
    self, labels: np.ndarray, segment_start_indices: np.ndarray, segment_length: int
) -> np.ndarray:
    """Assign each segment the majority label of its samples.

    Args:
        labels: Per-sample label timeline from extract_labels
        segment_start_indices: Segment start positions from ECGSegmenter
        segment_length: Segment length in samples

    Returns:
        One label per segment.
    """
    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)

convert_to_tensor

convert_to_tensor(
    data: ndarray | Tensor, dtype: dtype = torch.float32
) -> torch.Tensor

Convert an array to a tensor of the given dtype, passing tensors through.

Source code in deepecgkit/datasets/preprocessing.py

def convert_to_tensor(
    data: np.ndarray | torch.Tensor, dtype: torch.dtype = torch.float32
) -> torch.Tensor:
    """Convert an array to a tensor of the given dtype, passing tensors through."""
    if isinstance(data, torch.Tensor):
        return data.to(dtype)
    return torch.from_numpy(data).to(dtype)

Dataset Implementations

AFClassificationDataset

Bases: BaseECGDataset

PhysioNet/Computing in Cardiology Challenge 2017 AF Classification Dataset.

This dataset contains over 10,000 single-lead ECG recordings of 30-60 seconds duration for atrial fibrillation (AF) classification. Each recording is labeled as one of four categories: Normal (N), Atrial Fibrillation (A), Other rhythm (O), or Noisy (~).

The recordings are from AliveCor device and represent patient-initiated recordings.

Reference

Clifford GD, Liu C, Moody B, Li-wei HL, Silva I, Li Q, Johnson AE, Mark RG. AF classification from a short single lead ECG recording: The PhysioNet/computing in cardiology challenge 2017. In 2017 Computing in Cardiology (CinC) 2017 Sep 24 (pp. 1-4). IEEE.

URL

https://physionet.org/content/challenge-2017/1.0.0/

Source code in deepecgkit/datasets/af_classification.py

@register_dataset(
    name="af-classification",
    input_channels=1,
    description="PhysioNet 2017 AF Classification (4 classes, single-lead)",
)
class AFClassificationDataset(BaseECGDataset):
    """PhysioNet/Computing in Cardiology Challenge 2017 AF Classification Dataset.

    This dataset contains over 10,000 single-lead ECG recordings of 30-60 seconds duration
    for atrial fibrillation (AF) classification. Each recording is labeled as one of four
    categories: Normal (N), Atrial Fibrillation (A), Other rhythm (O), or Noisy (~).

    The recordings are from AliveCor device and represent patient-initiated recordings.

    Reference:
        Clifford GD, Liu C, Moody B, Li-wei HL, Silva I, Li Q, Johnson AE, Mark RG.
        AF classification from a short single lead ECG recording: The PhysioNet/computing
        in cardiology challenge 2017. In 2017 Computing in Cardiology (CinC) 2017 Sep 24 (pp. 1-4). IEEE.

    URL:
        https://physionet.org/content/challenge-2017/1.0.0/
    """

    CLASS_LABELS: ClassVar[list[str]] = [
        "Normal",
        "AF",
        "Other",
        "Noisy",
    ]

    REFERENCE_FILE: ClassVar[str] = "REFERENCE-v3.csv"

    LABEL_MAPPING: ClassVar[dict[str, int]] = {
        "N": 0,
        "A": 1,
        "O": 2,
        "~": 3,
    }

    LEADS: ClassVar[list[str]] = ["ECG"]
    DOWNLOAD_URL: ClassVar[str] = "https://physionet.org/content/challenge-2017/get-zip/1.0.0/"
    #: Mirrors of the project's file directory, tried in order. The AWS
    #: open-data mirror comes first because physionet.org is often throttled.
    DOWNLOAD_MIRRORS: ClassVar[tuple[str, ...]] = (
        "https://physionet-open.s3.amazonaws.com/challenge-2017/1.0.0/",
        "https://physionet.org/files/challenge-2017/1.0.0/",
    )
    #: The files the loader reads. The whole-project archive behind
    #: DOWNLOAD_URL is 1.4 GB and repeats the same recordings several times.
    DOWNLOAD_FILES: ClassVar[tuple[str, ...]] = (
        "REFERENCE-v3.csv",
        "training2017.zip",
        "sample2017.zip",
    )
    DOWNLOAD_SIZE: ClassVar[str | None] = "about 100 MB"
    INNER_ARCHIVES: ClassVar[tuple[str, ...]] = ("training2017.zip", "sample2017.zip")
    #: Archive name used by releases that downloaded the whole project.
    LEGACY_ARCHIVE: ClassVar[str] = "challenge-2017-1.0.0.zip"
    REFERENCE_FALLBACKS: ClassVar[tuple[str, ...]] = (
        "REFERENCE.csv",
        "REFERENCE-v2.csv",
        "REFERENCE-v1.csv",
    )
    NATIVE_SAMPLING_RATE: ClassVar[int] = 300

    def __init__(
        self,
        data_dir: str | Path | None = None,
        sampling_rate: int = 300,
        segment_duration_seconds: float = 10.0,
        segment_overlap: float = 0.0,
        normalization: str = "zscore",
        strict: bool = False,
        subset: str = "training",
        transform: Callable[[torch.Tensor], torch.Tensor] | None = None,
        target_transform: Callable[[torch.Tensor], torch.Tensor] | None = None,
        download: bool = False,
        force_download: bool = False,
        version: str = "1.0.0",
        verbose: bool = True,
    ):
        if subset not in ["training", "validation"]:
            raise ValueError("subset must be either 'training' or 'validation'")

        self.version = version
        self.subset = subset
        self.strict = strict
        self.verbose = verbose
        self.segment_duration_seconds = segment_duration_seconds
        self.segment_overlap = segment_overlap

        self.standardizer = ECGStandardizer(
            target_sampling_rate=sampling_rate,
            normalization=normalization,
        )

        self.segmenter = ECGSegmenter(
            segment_duration_seconds=segment_duration_seconds,
            sampling_rate=sampling_rate,
            overlap=segment_overlap,
        )

        self.signals: list[np.ndarray] = []
        self.labels: list[int] = []
        self.record_names: list[str] = []
        self.reference_data: pd.DataFrame | None = None

        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,
        )

    @classmethod
    def _find_reference(cls, data_dir: Path) -> Path | None:
        """Locate the label file, preferring the corrected v3 annotations."""
        for name in (cls.REFERENCE_FILE, *cls.REFERENCE_FALLBACKS):
            candidate = data_dir / name
            if candidate.exists():
                return candidate
        return None

    @classmethod
    def is_downloaded(cls, data_dir: str | Path | None = None) -> bool:
        """Report whether the labels and the training recordings are present.

        A directory holding only an interrupted download (a partial file, or a
        truncated archive) does not count, so download=True repeats the
        download instead of failing on the leftovers.
        """
        target = Path(data_dir) if data_dir is not None else cls.get_default_data_dir()
        if not target.is_dir() or cls._find_reference(target) is None:
            return False
        if (target / IN_PROGRESS_MARKER).exists():
            return False
        archive = target / "training2017.zip"
        if archive.exists():
            # Archives are deleted once extracted, so one that is still here
            # marks an extraction that never finished; it is usable only if
            # it is a whole archive rather than a truncated one.
            return zipfile.is_zipfile(archive)
        return (target / "training2017").is_dir()

    @classmethod
    def _download(cls, data_dir: Path, verbose: bool = True) -> None:
        """Download the label file and the recording archives from PhysioNet."""
        if verbose:
            logger.info(f"Downloading PhysioNet Challenge 2017 dataset to {data_dir}")
            logger.info(f"Note: this fetches {cls.DOWNLOAD_SIZE}; an interrupted download resumes.")
        data_dir.mkdir(parents=True, exist_ok=True)

        # Earlier releases downloaded the whole project into one archive; an
        # interrupted run left a truncated copy that nothing ever cleaned up.
        legacy_archive = data_dir / cls.LEGACY_ARCHIVE
        if legacy_archive.exists() and not zipfile.is_zipfile(legacy_archive):
            if verbose:
                logger.info(f"Removing incomplete download {legacy_archive.name}...")
            legacy_archive.unlink()

        with marking_download(data_dir):
            for name in cls.DOWNLOAD_FILES:
                already_extracted = (data_dir / Path(name).stem).is_dir()
                if (
                    name in cls.INNER_ARCHIVES
                    and already_extracted
                    and not (data_dir / name).exists()
                ):
                    continue
                download_file(
                    [mirror + name for mirror in cls.DOWNLOAD_MIRRORS],
                    data_dir / name,
                    desc=f"Downloading {name}",
                    max_retries=5,
                )

            cls._flatten_nested_dir(data_dir, verbose)
            cls._extract_inner_zips(data_dir, verbose)
        if verbose:
            logger.info("Download complete!")

    @classmethod
    def _flatten_nested_dir(cls, data_dir: Path, verbose: bool = True) -> None:
        """Flatten a nested extraction directory if present from a previous download."""
        if not data_dir.is_dir():
            return
        for child in data_dir.iterdir():
            if not child.is_dir() or child.name in ("__MACOSX", "training2017", "sample2017"):
                continue
            holds_dataset = (
                any(child.glob("REFERENCE*.csv"))
                or any(child.glob("*2017.zip"))
                or (child / "training2017").is_dir()
            )
            if not holds_dataset:
                continue
            if verbose:
                logger.info(f"Flattening nested directory {child.name}...")
            flatten_directory(child, data_dir)
            return

    @classmethod
    def _extract_inner_zips(cls, data_dir: Path, verbose: bool = True) -> None:
        """Extract training2017.zip and sample2017.zip, then delete them.

        An archive that is still present next to its extracted folder means
        the previous extraction was cut short, so the folder is rebuilt.

        Raises:
            FileNotFoundError: If an archive is truncated or corrupted.
        """
        if not data_dir.is_dir():
            return
        for name in cls.INNER_ARCHIVES:
            zf_path = data_dir / name
            if not zf_path.exists():
                continue
            target_dir = data_dir / zf_path.stem
            if target_dir.exists():
                shutil.rmtree(target_dir)
            if verbose:
                logger.info(f"Extracting {zf_path.name}...")
            try:
                with zipfile.ZipFile(zf_path, "r") as zf:
                    zf.extractall(data_dir)
            except (zipfile.BadZipFile, zlib.error, EOFError) as error:
                raise FileNotFoundError(
                    cls._not_found_message(
                        data_dir, f"a complete {name}, but the one there is truncated or corrupted"
                    )
                ) from error
            zf_path.unlink()

    def _load_data(self) -> None:
        """Load the AF Classification dataset into memory."""
        if self.verbose:
            logger.info(f"Loading AF Classification dataset ({self.subset} subset)...")

        self._flatten_nested_dir(self.data_dir, self.verbose)
        self._extract_inner_zips(self.data_dir, self.verbose)

        if self.subset == "training":
            expected = f"{self.REFERENCE_FILE} and a training2017/ folder"
            data_folder = self.data_dir / "training2017"
            reference_path = self._find_reference(self.data_dir)
        else:
            # sample2017.zip keeps its recordings and their own labels in a
            # validation/ subfolder.
            expected = "a sample2017/ folder with its REFERENCE.csv"
            data_folder = self.data_dir / "sample2017"
            if (data_folder / "validation").is_dir():
                data_folder = data_folder / "validation"
            reference_path = data_folder / "REFERENCE.csv"
            if not reference_path.exists():
                reference_path = self._find_reference(self.data_dir)

        if reference_path is None or not reference_path.exists() or not data_folder.is_dir():
            raise FileNotFoundError(self._not_found_message(self.data_dir, expected))

        self.reference_data = pd.read_csv(
            reference_path, header=None, names=["record_name", "label"]
        )

        iterator = (
            tqdm(
                self.reference_data.iterrows(),
                total=len(self.reference_data),
                desc="Loading records",
            )
            if self.verbose
            else self.reference_data.iterrows()
        )

        attempted_records = 0
        loaded_records = 0
        for _, row in iterator:
            record_name = row["record_name"]
            label = row["label"]

            mat_file = data_folder / f"{record_name}.mat"

            if mat_file.exists():
                attempted_records += 1
                try:
                    mat_data = scipy.io.loadmat(str(mat_file))

                    if "val" in mat_data:
                        signal = mat_data["val"].flatten()
                    elif "ecg" in mat_data:
                        signal = mat_data["ecg"].flatten()
                    elif "data" in mat_data:
                        signal = mat_data["data"].flatten()
                    else:
                        signal_key = [k for k in mat_data.keys() if not k.startswith("__")][0]
                        signal = mat_data[signal_key].flatten()

                    if label not in self.LABEL_MAPPING:
                        if self.verbose:
                            logger.warning(f"Unknown label '{label}' for record {record_name}")
                        continue

                    label_idx = self.LABEL_MAPPING[label]

                    signal = signal[np.newaxis, :]
                    resampled = self.standardizer.resample(signal, self.NATIVE_SAMPLING_RATE)

                    segments, _ = self.segmenter.segment(resampled)
                    if len(segments) == 0:
                        continue

                    for segment in segments:
                        normalized = self.standardizer.normalize(segment)
                        self.signals.append(normalized.astype(np.float32))
                        self.labels.append(label_idx)
                        self.record_names.append(record_name)
                    loaded_records += 1

                except self.RECORD_ERRORS as e:
                    self._record_load_error(record_name, e)
                    continue

        self._finish_load(loaded_records, attempted_records)

        if self.verbose:
            logger.info(
                f"Successfully loaded {len(self.signals)} segments from {self.subset} subset"
            )

is_downloaded classmethod

is_downloaded(data_dir: str | Path | None = None) -> bool

Report whether the labels and the training recordings are present.

A directory holding only an interrupted download (a partial file, or a truncated archive) does not count, so download=True repeats the download instead of failing on the leftovers.

Source code in deepecgkit/datasets/af_classification.py

@classmethod
def is_downloaded(cls, data_dir: str | Path | None = None) -> bool:
    """Report whether the labels and the training recordings are present.

    A directory holding only an interrupted download (a partial file, or a
    truncated archive) does not count, so download=True repeats the
    download instead of failing on the leftovers.
    """
    target = Path(data_dir) if data_dir is not None else cls.get_default_data_dir()
    if not target.is_dir() or cls._find_reference(target) is None:
        return False
    if (target / IN_PROGRESS_MARKER).exists():
        return False
    archive = target / "training2017.zip"
    if archive.exists():
        # Archives are deleted once extracted, so one that is still here
        # marks an extraction that never finished; it is usable only if
        # it is a whole archive rather than a truncated one.
        return zipfile.is_zipfile(archive)
    return (target / "training2017").is_dir()

LTAFDBDataset

Bases: BaseECGDataset

Long-Term AF Database (LTAFDB) Dataset.

Contains 84 long-term (typically 24-hour) two-lead ECG recordings of subjects with paroxysmal or sustained atrial fibrillation. Rhythm annotations indicate Normal, AF, Atrial Flutter, and Junctional rhythm segments.

Reference

Petrutiu S, Sahakian AV, Swiryn S. Abrupt changes in fibrillatory wave characteristics at the termination of paroxysmal atrial fibrillation in humans. Europace. 2007;9(7):466-470.

URL

https://physionet.org/content/ltafdb/1.0.0/

Source code in deepecgkit/datasets/ltafdb.py

@register_dataset(
    name="ltafdb",
    input_channels=1,
    num_classes=4,
    description="Long-Term AF Database (first lead by default, binary or 4-class)",
)
class LTAFDBDataset(BaseECGDataset):
    """Long-Term AF Database (LTAFDB) Dataset.

    Contains 84 long-term (typically 24-hour) two-lead ECG recordings of subjects
    with paroxysmal or sustained atrial fibrillation. Rhythm annotations indicate
    Normal, AF, Atrial Flutter, and Junctional rhythm segments.

    Reference:
        Petrutiu S, Sahakian AV, Swiryn S. Abrupt changes in fibrillatory wave
        characteristics at the termination of paroxysmal atrial fibrillation in humans.
        Europace. 2007;9(7):466-470.

    URL:
        https://physionet.org/content/ltafdb/1.0.0/
    """

    CLASS_LABELS: ClassVar[list[str]] = ["Normal", "AF", "AFL", "J"]
    LEADS: ClassVar[list[str]] = ["ECG1", "ECG2"]
    SAMPLING_RATE: ClassVar[int] = 128

    DOWNLOAD_URL: ClassVar[str] = "https://physionet.org/content/ltafdb/get-zip/1.0.0/"
    DOWNLOAD_MIRROR: ClassVar[str] = "https://physionet-open.s3.amazonaws.com/ltafdb/1.0.0/"
    DOWNLOAD_SIZE: ClassVar[str | None] = "about 3.7 GB"
    RECORD_SUFFIXES: ClassVar[tuple[str, ...]] = (".hea", ".atr", ".dat")
    #: Part of the cache key; raised whenever cached segments or labels would
    #: differ for the same parameters, so stale caches are never read back.
    #: 2: rhythm labels are carried across the per-beat annotations (version 1
    #: caches label almost every segment Normal).
    CACHE_VERSION: ClassVar[int] = 2
    #: Written into every cache; a cache folder without it predates version 2.
    CACHE_INFO_FILE: ClassVar[str] = "cache_info.json"
    #: The records are tens of megabytes each, so a few are fetched at a time.
    DOWNLOAD_WORKERS: ClassVar[int] = 4

    def __init__(
        self,
        data_dir: str | Path | None = None,
        sampling_rate: int = 300,
        segment_duration_seconds: float = 10.0,
        segment_overlap: float = 0.0,
        binary_classification: bool = False,
        use_both_leads: bool = False,
        normalization: str = "zscore",
        max_segments_per_record: int | None = None,
        subsample_seed: int = 0,
        strict: bool = False,
        transform: Callable[[torch.Tensor], torch.Tensor] | None = None,
        target_transform: Callable[[torch.Tensor], torch.Tensor] | None = None,
        download: bool = False,
        force_download: bool = False,
        verbose: bool = True,
    ):
        self.segment_duration_seconds = segment_duration_seconds
        self.segment_overlap = segment_overlap
        self.binary_classification = binary_classification
        self.use_both_leads = use_both_leads
        self.max_segments_per_record = max_segments_per_record
        self.subsample_seed = subsample_seed
        self.strict = strict
        self.verbose = verbose

        self.standardizer = ECGStandardizer(
            target_sampling_rate=sampling_rate,
            normalization=normalization,
        )

        self.segmenter = ECGSegmenter(
            segment_duration_seconds=segment_duration_seconds,
            sampling_rate=sampling_rate,
            overlap=segment_overlap,
        )

        self.rhythm_extractor = RhythmAnnotationExtractor(
            sampling_rate=sampling_rate, binary_classification=binary_classification
        )

        self.signals: np.ndarray = np.array([])
        self.labels: np.ndarray = np.array([], dtype=np.int64)
        self.record_names: list[str] = []

        super().__init__(
            data_dir=data_dir,
            sampling_rate=sampling_rate,
            leads=self.LEADS if use_both_leads else [self.LEADS[0]],
            transform=transform,
            target_transform=target_transform,
            download=download,
            force_download=force_download,
        )

    def _cache_key(self) -> str:
        params = (
            f"sr{self.sampling_rate}_dur{self.segment_duration_seconds}_"
            f"ovlp{self.segment_overlap}_bin{self.binary_classification}_"
            f"leads{self.use_both_leads}_norm{self.standardizer.normalization}_"
            f"max{self.max_segments_per_record}"
        )
        if self.max_segments_per_record is not None:
            # The seed shapes which segments survive subsampling, so it is part
            # of the cache identity. Unsubsampled caches are seed-independent
            # and share one key.
            params += f"_seed{self.subsample_seed}"
        params += f"_v{self.CACHE_VERSION}"
        return hashlib.md5(params.encode()).hexdigest()[:12]

    def _cache_dir(self) -> Path:
        return self.data_dir / f"cache_{self._cache_key()}"

    @staticmethod
    def _grow_memmap(path: Path, writer: np.memmap, new_rows: int) -> np.memmap:
        """Grow an on-disk .npy memmap to new_rows, preserving its content.

        open_memmap in r+ mode reads the shape from the file header and
        ignores a passed shape, so "reopening larger" silently keeps the old
        capacity — writes past it then raise, and the per-record error
        handling drops those records. A real grow needs a new file.

        Args:
            path: Path of the memmap file being grown
            writer: The currently open memmap over path
            new_rows: New first-dimension size, larger than the current one

        Returns:
            A writable memmap over the grown file.
        """
        old_rows = writer.shape[0]
        tail_shape = writer.shape[1:]
        dtype = writer.dtype
        grow_path = path.with_suffix(".grow.npy")
        grown = np.lib.format.open_memmap(
            str(grow_path), mode="w+", dtype=dtype, shape=(new_rows, *tail_shape)
        )
        grown[:old_rows] = writer[:old_rows]
        grown.flush()
        del writer
        del grown
        grow_path.replace(path)
        return np.lib.format.open_memmap(str(path), mode="r+")

    @classmethod
    def is_downloaded(cls, data_dir: str | Path | None = None) -> bool:
        """Report whether record files, or a finished segment cache, are present.

        A directory still holding the archive or a partial download does not
        count: the archive is deleted once extraction finishes.
        """
        target = Path(data_dir) if data_dir is not None else cls.get_default_data_dir()
        if has_unfinished_download(target):
            return False
        if any(target.glob("cache_*/signals.npy")):
            return True
        return any(
            any(folder.glob("*.hea")) and any(folder.glob("*.atr"))
            for folder in (target, target / "files")
        )

    @classmethod
    def _download(cls, data_dir: Path, verbose: bool = True) -> None:
        if verbose:
            logger.info(f"Downloading Long-Term AF Database to {data_dir}")
            logger.info(f"Note: this fetches {cls.DOWNLOAD_SIZE}; an interrupted download resumes.")
        cls._download_records(
            data_dir,
            lambda: cls._mirrored_files(data_dir),
            mirror=cls.DOWNLOAD_MIRROR,
            desc="Downloading LTAFDB",
            archive_url=cls.DOWNLOAD_URL,
            archive_name="ltafdb-1.0.0.zip",
            workers=cls.DOWNLOAD_WORKERS,
            keep_alive=False,
            verbose=verbose,
        )
        if verbose:
            logger.info("Download complete!")

    @classmethod
    def _mirrored_files(cls, data_dir: Path) -> list[str]:
        """List the record files to fetch, from the project's own index.

        Raises:
            RuntimeError: If the index cannot be fetched.
            ValueError: If the index names no records.
        """
        index = data_dir / "RECORDS"
        download_file(cls.DOWNLOAD_MIRROR + "RECORDS", index, progress=False)
        records = [line.strip() for line in index.read_text().splitlines() if line.strip()]
        if not records or not all(record.isalnum() for record in records):
            raise ValueError(f"unexpected content in the record index {index}")
        return [f"{record}{suffix}" for record in records for suffix in cls.RECORD_SUFFIXES]

    def _resolve_record_dir(self) -> Path:
        """Find the directory containing .hea/.atr record files."""
        if any(self.data_dir.glob("*.hea")):
            return self.data_dir
        files_dir = self.data_dir / "files"
        if files_dir.is_dir() and any(files_dir.glob("*.hea")):
            return files_dir
        return self.data_dir

    def _discover_records(self, record_dir: Path) -> list[str]:
        """Discover available records by finding files with both .hea and .atr extensions."""
        hea_stems = {p.stem for p in record_dir.glob("*.hea")}
        atr_stems = {p.stem for p in record_dir.glob("*.atr")}
        return sorted(hea_stems & atr_stems, key=lambda x: (len(x), x))

    def _warn_about_stale_caches(self, current: Path) -> None:
        """Point out cache folders that are no longer read and only take up space.

        Caches written before CACHE_VERSION 2 carry labels in which almost
        every segment is Normal. Their keys differ from today's, so they are
        never loaded, but each holds several gigabytes.

        Args:
            current: The cache folder this instance uses
        """
        if not self.data_dir.is_dir():
            return
        stale = [
            folder
            for folder in sorted(self.data_dir.glob("cache_*"))
            if folder.is_dir()
            and folder != current
            and (folder / "signals.npy").exists()
            and not (folder / self.CACHE_INFO_FILE).exists()
        ]
        if not stale:
            return
        size = sum(f.stat().st_size for folder in stale for f in folder.iterdir() if f.is_file())
        logger.warning(
            f"{len(stale)} LTAFDB cache folder(s) in {self.data_dir} were written by an earlier "
            f"release that labelled almost every segment Normal. They are no longer read and "
            f"can be deleted to free {size / 1e9:.1f} GB: {', '.join(f.name for f in stale)}"
        )

    def _load_data(self) -> None:
        if self.verbose:
            logger.info("Loading Long-Term AF Database...")

        cache_dir = self._cache_dir()
        signals_path = cache_dir / "signals.npy"
        labels_path = cache_dir / "labels.npy"
        record_names_path = cache_dir / "record_names.npy"
        self._warn_about_stale_caches(cache_dir)

        if signals_path.exists() and labels_path.exists():
            if self.verbose:
                logger.info(f"Loading from cache: {cache_dir.name}")
            self.signals = np.load(signals_path, mmap_mode="r")
            self.labels = np.load(labels_path, mmap_mode="r")
            if record_names_path.exists():
                self.record_names = np.load(record_names_path, allow_pickle=True).tolist()
            else:
                self.record_names = ["unknown"] * len(self.labels)
            if self.verbose:
                logger.info(f"Loaded {len(self.labels)} segments from cache (memory-mapped)")
                self._print_class_distribution()
            return

        record_dir = self._resolve_record_dir()
        record_names = self._discover_records(record_dir)
        if not record_names:
            raise FileNotFoundError(
                self._not_found_message(self.data_dir, ".hea and .atr record files from LTAFDB")
            )

        if self.verbose:
            logger.info(f"Found {len(record_names)} records — processing (first run only)...")

        cache_dir.mkdir(parents=True, exist_ok=True)
        tmp_signals = cache_dir / "signals_tmp.npy"

        num_leads = 2 if self.use_both_leads else 1
        seg_len = self.segmenter.segment_length

        writer = None
        write_offset = 0
        label_chunks: list[np.ndarray] = []
        all_record_names: list[str] = []
        total_segments = 0
        loaded_records = 0

        for record_name in record_names:
            record_path = record_dir / record_name

            try:
                segments, start_indices, label_timeline = read_rhythm_record(
                    record_path,
                    standardizer=self.standardizer,
                    segmenter=self.segmenter,
                    rhythm_extractor=self.rhythm_extractor,
                    native_sampling_rate=self.SAMPLING_RATE,
                    use_both_leads=self.use_both_leads,
                )

                if len(segments) == 0:
                    continue

                if (
                    self.max_segments_per_record is not None
                    and len(segments) > self.max_segments_per_record
                ):
                    rng = np.random.default_rng(self.subsample_seed)
                    indices = rng.choice(
                        len(segments),
                        self.max_segments_per_record,
                        replace=False,
                    )
                    segments = segments[indices]
                    start_indices = start_indices[indices]

                segment_labels = self.rhythm_extractor.segment_with_labels(
                    label_timeline, start_indices, self.segmenter.segment_length
                )

                normalized = np.stack([self.standardizer.normalize(seg) for seg in segments])

                if writer is None:
                    estimated_total = len(record_names) * len(segments)
                    writer = np.lib.format.open_memmap(
                        str(tmp_signals),
                        mode="w+",
                        dtype=np.float32,
                        shape=(estimated_total, num_leads, seg_len),
                    )

                needed = write_offset + len(normalized)
                if needed > writer.shape[0]:
                    new_size = max(needed, writer.shape[0] * 2)
                    writer = self._grow_memmap(tmp_signals, writer, new_size)

                writer[write_offset : write_offset + len(normalized)] = normalized
                label_chunks.append(np.array(segment_labels, dtype=np.int64))
                all_record_names.extend([record_name] * len(segments))
                write_offset += len(normalized)
                total_segments += len(segments)
                loaded_records += 1

                if self.verbose:
                    logger.info(
                        f"Loaded {record_name}: {len(segments)} segments (total: {total_segments})"
                    )

            except self.RECORD_ERRORS as e:
                self._record_load_error(record_name, e)
                continue

        self._finish_load(loaded_records, len(record_names))

        if write_offset == 0:
            shutil.rmtree(cache_dir)
            self.signals = np.array([])
            self.labels = np.array([], dtype=np.int64)
            self.record_names = []
            return

        del writer

        final_signals = np.lib.format.open_memmap(str(tmp_signals), mode="r")
        np.save(signals_path, final_signals[:write_offset])
        del final_signals
        tmp_signals.unlink(missing_ok=True)

        all_labels = np.concatenate(label_chunks)
        np.save(labels_path, all_labels)
        np.save(record_names_path, np.array(all_record_names))
        (cache_dir / self.CACHE_INFO_FILE).write_text(
            json.dumps({"cache_version": self.CACHE_VERSION, "segments": int(write_offset)}) + "\n"
        )

        self.signals = np.load(signals_path, mmap_mode="r")
        self.labels = np.load(labels_path, mmap_mode="r")
        self.record_names = all_record_names

        if self.verbose:
            logger.info(f"\nCached {len(self.labels)} segments to {cache_dir.name} (memory-mapped)")
            self._print_class_distribution()

is_downloaded classmethod

is_downloaded(data_dir: str | Path | None = None) -> bool

Report whether record files, or a finished segment cache, are present.

A directory still holding the archive or a partial download does not count: the archive is deleted once extraction finishes.

Source code in deepecgkit/datasets/ltafdb.py

@classmethod
def is_downloaded(cls, data_dir: str | Path | None = None) -> bool:
    """Report whether record files, or a finished segment cache, are present.

    A directory still holding the archive or a partial download does not
    count: the archive is deleted once extraction finishes.
    """
    target = Path(data_dir) if data_dir is not None else cls.get_default_data_dir()
    if has_unfinished_download(target):
        return False
    if any(target.glob("cache_*/signals.npy")):
        return True
    return any(
        any(folder.glob("*.hea")) and any(folder.glob("*.atr"))
        for folder in (target, target / "files")
    )

MITBIHAFDBDataset

Bases: BaseECGDataset

MIT-BIH Atrial Fibrillation Database (AFDB) Dataset.

Contains 25 long-term (10-hour) two-lead ECG recordings from subjects with atrial fibrillation (mostly paroxysmal). Rhythm annotations indicate Normal, AF, Atrial Flutter, and Junctional rhythm segments.

Reference

Moody GB, Mark RG. A new method for detecting atrial fibrillation using R-R intervals. Computers in Cardiology. 1983;10:227-230.

URL

https://physionet.org/content/afdb/1.0.0/

Source code in deepecgkit/datasets/mitbih_afdb.py

@register_dataset(
    name="mitbih-afdb",
    input_channels=1,
    num_classes=4,
    description="MIT-BIH AF Database (first lead by default, binary or 4-class)",
)
class MITBIHAFDBDataset(BaseECGDataset):
    """MIT-BIH Atrial Fibrillation Database (AFDB) Dataset.

    Contains 25 long-term (10-hour) two-lead ECG recordings from subjects with
    atrial fibrillation (mostly paroxysmal). Rhythm annotations indicate Normal,
    AF, Atrial Flutter, and Junctional rhythm segments.

    Reference:
        Moody GB, Mark RG. A new method for detecting atrial fibrillation using R-R
        intervals. Computers in Cardiology. 1983;10:227-230.

    URL:
        https://physionet.org/content/afdb/1.0.0/
    """

    CLASS_LABELS: ClassVar[list[str]] = ["Normal", "AF", "AFL", "J"]
    LEADS: ClassVar[list[str]] = ["ECG1", "ECG2"]
    SAMPLING_RATE: ClassVar[int] = 250
    DOWNLOAD_URL: ClassVar[str] = "https://physionet.org/content/afdb/get-zip/1.0.0/"
    DOWNLOAD_MIRROR: ClassVar[str] = "https://physionet-open.s3.amazonaws.com/afdb/1.0.0/"
    DOWNLOAD_SIZE: ClassVar[str | None] = "about 640 MB"
    RECORD_SUFFIXES: ClassVar[tuple[str, ...]] = (".hea", ".atr", ".dat")
    #: The records are tens of megabytes each, so a few are fetched at a time.
    DOWNLOAD_WORKERS: ClassVar[int] = 4

    RECORD_NAMES: ClassVar[list[str]] = [
        "04015",
        "04043",
        "04048",
        "04126",
        "04746",
        "04908",
        "04936",
        "05091",
        "05121",
        "05261",
        "06426",
        "06453",
        "06995",
        "07162",
        "07859",
        "07879",
        "07910",
        "08215",
        "08219",
        "08378",
        "08405",
        "08434",
        "08455",
    ]

    def __init__(
        self,
        data_dir: str | Path | None = None,
        sampling_rate: int = 300,
        segment_duration_seconds: float = 10.0,
        segment_overlap: float = 0.0,
        binary_classification: bool = False,
        use_both_leads: bool = False,
        normalization: str = "zscore",
        strict: bool = False,
        transform: Callable[[torch.Tensor], torch.Tensor] | None = None,
        target_transform: Callable[[torch.Tensor], torch.Tensor] | None = None,
        download: bool = False,
        force_download: bool = False,
        verbose: bool = True,
    ):
        self.segment_duration_seconds = segment_duration_seconds
        self.segment_overlap = segment_overlap
        self.binary_classification = binary_classification
        self.use_both_leads = use_both_leads
        self.strict = strict
        self.verbose = verbose

        self.standardizer = ECGStandardizer(
            target_sampling_rate=sampling_rate,
            normalization=normalization,
        )

        self.segmenter = ECGSegmenter(
            segment_duration_seconds=segment_duration_seconds,
            sampling_rate=sampling_rate,
            overlap=segment_overlap,
        )

        self.rhythm_extractor = RhythmAnnotationExtractor(
            sampling_rate=sampling_rate, binary_classification=binary_classification
        )

        self.signals: list[np.ndarray] = []
        self.labels: list[int] = []
        self.record_names: list[str] = []

        super().__init__(
            data_dir=data_dir,
            sampling_rate=sampling_rate,
            leads=self.LEADS if use_both_leads else [self.LEADS[0]],
            transform=transform,
            target_transform=target_transform,
            download=download,
            force_download=force_download,
        )

    @classmethod
    def is_downloaded(cls, data_dir: str | Path | None = None) -> bool:
        """Report whether extracted record files are present.

        A directory still holding the archive or a partial download does not
        count: the archive is deleted once extraction finishes.
        """
        target = Path(data_dir) if data_dir is not None else cls.get_default_data_dir()
        if has_unfinished_download(target):
            return False
        return any(target.glob("*.hea")) or any((target / "files").glob("*.hea"))

    @classmethod
    def _download(cls, data_dir: Path, verbose: bool = True) -> None:
        if verbose:
            logger.info(f"Downloading MIT-BIH AFDB to {data_dir}")
            logger.info(f"Note: this fetches {cls.DOWNLOAD_SIZE}; an interrupted download resumes.")
        cls._download_records(
            data_dir,
            [f"{record}{suffix}" for record in cls.RECORD_NAMES for suffix in cls.RECORD_SUFFIXES],
            mirror=cls.DOWNLOAD_MIRROR,
            desc="Downloading MIT-BIH AFDB",
            archive_url=cls.DOWNLOAD_URL,
            archive_name="afdb-1.0.0.zip",
            workers=cls.DOWNLOAD_WORKERS,
            keep_alive=False,
            verbose=verbose,
        )
        if verbose:
            logger.info("Download complete!")

    def _load_data(self) -> None:
        if self.verbose:
            logger.info("Loading MIT-BIH AFDB data...")

        # PhysioNet zips may nest records inside a files/ subdirectory
        files_dir = self.data_dir / "files"
        if (
            not any(self.data_dir.glob("*.hea"))
            and files_dir.is_dir()
            and any(files_dir.glob("*.hea"))
        ):
            self.data_dir = files_dir

        if not any(self.data_dir.glob("*.hea")):
            raise FileNotFoundError(
                self._not_found_message(self.data_dir, ".hea record files from MIT-BIH AFDB")
            )

        attempted_records = 0
        loaded_records = 0
        for record_name in self.RECORD_NAMES:
            record_path = self.data_dir / record_name

            if not (self.data_dir / f"{record_name}.hea").exists():
                if self.verbose:
                    logger.info(f"Skipping {record_name} (not found)")
                continue

            attempted_records += 1
            try:
                segments, start_indices, label_timeline = read_rhythm_record(
                    record_path,
                    standardizer=self.standardizer,
                    segmenter=self.segmenter,
                    rhythm_extractor=self.rhythm_extractor,
                    native_sampling_rate=self.SAMPLING_RATE,
                    use_both_leads=self.use_both_leads,
                )

                if len(segments) == 0:
                    continue

                segment_labels = self.rhythm_extractor.segment_with_labels(
                    label_timeline, start_indices, self.segmenter.segment_length
                )

                for segment, label in zip(segments, segment_labels, strict=False):
                    normalized_segment = self.standardizer.normalize(segment)
                    self.signals.append(normalized_segment.astype(np.float32))
                    self.labels.append(label)
                    self.record_names.append(record_name)

                loaded_records += 1
                if self.verbose:
                    logger.info(f"Loaded {record_name}: {len(segments)} segments")

            except self.RECORD_ERRORS as e:
                self._record_load_error(record_name, e)

        self._finish_load(loaded_records, attempted_records)

        if self.verbose:
            logger.info(f"Total segments loaded: {len(self.signals)}")
            self._print_class_distribution()

is_downloaded classmethod

is_downloaded(data_dir: str | Path | None = None) -> bool

Report whether extracted record files are present.

A directory still holding the archive or a partial download does not count: the archive is deleted once extraction finishes.

Source code in deepecgkit/datasets/mitbih_afdb.py

@classmethod
def is_downloaded(cls, data_dir: str | Path | None = None) -> bool:
    """Report whether extracted record files are present.

    A directory still holding the archive or a partial download does not
    count: the archive is deleted once extraction finishes.
    """
    target = Path(data_dir) if data_dir is not None else cls.get_default_data_dir()
    if has_unfinished_download(target):
        return False
    return any(target.glob("*.hea")) or any((target / "files").glob("*.hea"))

PTBXLDataset

Bases: BaseECGDataset

PTB-XL ECG Dataset.

PTB-XL is a large publicly available electrocardiography dataset containing 21,837 clinical 12-lead ECGs from 18,885 patients of 10 second length. Each ECG is annotated with up to 71 different diagnostic statements conforming to the SCP-ECG standard.

The dataset supports multiple diagnostic classification tasks: - Superclass: 5 diagnostic superclasses (NORM, MI, STTC, CD, HYP) - Subclass: 23 diagnostic subclasses - Diagnostic: All individual diagnostic SCP codes (~44 statements) - Form: 19 form statements - Rhythm: 12 rhythm statements - All: All SCP statement codes (~71 statements)

Reference

Wagner P, Strodthoff N, Bousseljot RD, Kreiseler D, Lunze FI, Samek W, Schaeffter T. PTB-XL, a large publicly available electrocardiography dataset. Scientific Data. 2020 May 25;7(1):154.

URL

https://physionet.org/content/ptb-xl/1.0.3/

Source code in deepecgkit/datasets/ptbxl.py

@register_dataset(
    name="ptbxl",
    input_channels=12,
    num_classes=5,
    description="PTB-XL 12-lead ECG dataset (multi-task)",
)
class PTBXLDataset(BaseECGDataset):
    """PTB-XL ECG Dataset.

    PTB-XL is a large publicly available electrocardiography dataset containing
    21,837 clinical 12-lead ECGs from 18,885 patients of 10 second length.
    Each ECG is annotated with up to 71 different diagnostic statements conforming
    to the SCP-ECG standard.

    The dataset supports multiple diagnostic classification tasks:
    - Superclass: 5 diagnostic superclasses (NORM, MI, STTC, CD, HYP)
    - Subclass: 23 diagnostic subclasses
    - Diagnostic: All individual diagnostic SCP codes (~44 statements)
    - Form: 19 form statements
    - Rhythm: 12 rhythm statements
    - All: All SCP statement codes (~71 statements)

    Reference:
        Wagner P, Strodthoff N, Bousseljot RD, Kreiseler D, Lunze FI, Samek W, Schaeffter T.
        PTB-XL, a large publicly available electrocardiography dataset.
        Scientific Data. 2020 May 25;7(1):154.

    URL:
        https://physionet.org/content/ptb-xl/1.0.3/
    """

    CLASS_LABELS_SUPERCLASS: ClassVar[list[str]] = ["NORM", "MI", "STTC", "CD", "HYP"]
    CLASS_LABELS_SUBCLASS: ClassVar[list[str]] = [
        "NORM",
        "IMI",
        "ASMI",
        "ILMI",
        "AMI",
        "ALMI",
        "INJAS",
        "LMI",
        "INJAL",
        "ISCAL",
        "ISCAN",
        "INJIN",
        "INJLA",
        "PMI",
        "INJIL",
        "ISCIN",
        "ISCIL",
        "ISCAS",
        "LAFB",
        "IRBBB",
        "LPFB",
        "CRBBB",
        "CLBBB",
    ]
    CLASS_LABELS_FORM: ClassVar[list[str]] = [
        "NDT",
        "NST_",
        "DIG",
        "LNGQT",
        "ABQRS",
        "PVC",
        "STD_",
        "VCLVH",
        "QWAVE",
        "LOWT",
        "NT_",
        "PAC",
        "LPR",
        "INVT",
        "LVOLT",
        "HVOLT",
        "TAB_",
        "STE_",
        "PRC(S)",
    ]
    CLASS_LABELS_RHYTHM: ClassVar[list[str]] = [
        "SR",
        "AFIB",
        "STACH",
        "SARRH",
        "SBRAD",
        "PACE",
        "SVARR",
        "BIGU",
        "AFLT",
        "SVTAC",
        "PSVT",
        "TRIGU",
    ]

    LEADS: ClassVar[list[str]] = [
        "I",
        "II",
        "III",
        "AVR",
        "AVL",
        "AVF",
        "V1",
        "V2",
        "V3",
        "V4",
        "V5",
        "V6",
    ]
    SAMPLING_RATE_HR: ClassVar[int] = 500
    SAMPLING_RATE_LR: ClassVar[int] = 100

    def __init__(
        self,
        data_dir: str | Path | None = None,
        sampling_rate: int = 500,
        task: str = "superclass",
        use_high_resolution: bool = True,
        folds: list[int] | None = None,
        leads: list[str] | None = None,
        normalization: str = "zscore",
        multi_label: bool = True,
        transform: Callable[[torch.Tensor], torch.Tensor] | None = None,
        target_transform: Callable[[torch.Tensor], torch.Tensor] | None = None,
        download: bool = False,
        force_download: bool = False,
        strict: bool = False,
        verbose: bool = True,
    ):
        """Initialize the PTB-XL dataset.

        Args:
            data_dir: Directory where the dataset is stored or will be downloaded.
                     If None, uses ~/.deepecgkit/datasets/ptbxldataset
            sampling_rate: Target sampling rate for the ECG signals (Hz)
            task: Classification task - one of "superclass", "subclass",
                  "diagnostic", "form", "rhythm", or "all"
            use_high_resolution: Whether to use 500Hz (True) or 100Hz (False) recordings
            folds: List of folds to include (1-10). None means all folds.
                   Recommended: folds 1-8 for training, 9 for validation, 10 for testing
            leads: List of lead names to use. None means all 12 leads.
            normalization: Normalization method - "zscore", "minmax", or "none"
            multi_label: If True, returns multi-hot encoded labels. If False, returns
                        single label (first/primary diagnosis)
            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
            verbose: Whether to print progress information
        """
        valid_tasks = ["superclass", "subclass", "diagnostic", "form", "rhythm", "all"]
        if task not in valid_tasks:
            raise ValueError(f"task must be one of {valid_tasks}, got '{task}'")

        self.task = task
        self.use_high_resolution = use_high_resolution
        self.folds = folds
        self.multi_label = multi_label
        self.strict = strict
        self.verbose = verbose
        self._leads = leads

        source_sampling_rate = (
            self.SAMPLING_RATE_HR if use_high_resolution else self.SAMPLING_RATE_LR
        )

        self.standardizer = ECGStandardizer(
            target_sampling_rate=sampling_rate,
            target_duration_seconds=10.0,
            normalization=normalization,
            clip_method="center",
        )

        self.signals: list[np.ndarray] = []
        self.labels: list[np.ndarray] = []
        self.record_names: list[str] = []
        self.patient_ids: list[str] = []
        self.metadata_df: pd.DataFrame | None = None
        self.scp_statements: pd.DataFrame | None = None
        self.source_sampling_rate = source_sampling_rate

        super().__init__(
            data_dir=data_dir,
            sampling_rate=sampling_rate,
            leads=leads if leads else self.LEADS,
            transform=transform,
            target_transform=target_transform,
            download=download,
            force_download=force_download,
        )

    DOWNLOAD_URL: ClassVar[str] = (
        "https://physionet.org/static/published-projects/ptb-xl/"
        "ptb-xl-a-large-publicly-available-electrocardiography-dataset-1.0.3.zip"
    )
    DOWNLOAD_MIRROR: ClassVar[str] = "https://physionet-open.s3.amazonaws.com/ptb-xl/1.0.3/"
    DOWNLOAD_SIZE: ClassVar[str | None] = "about 3.2 GB"
    #: The mirror holds one small file per record and resolution, so many
    #: are fetched at a time.
    DOWNLOAD_WORKERS: ClassVar[int] = 32
    METADATA_FILES: ClassVar[tuple[str, ...]] = ("ptbxl_database.csv", "scp_statements.csv")
    #: Recordings are 10 seconds long; inference cuts longer ones to match.
    RECORDING_DURATION_SECONDS: ClassVar[float] = 10.0

    @classmethod
    def is_downloaded(cls, data_dir: str | Path | None = None) -> bool:
        """Report whether the extracted metadata table is present.

        A directory still holding the archive or a partial download does not
        count: the archive is deleted once extraction finishes.
        """
        target = Path(data_dir) if data_dir is not None else cls.get_default_data_dir()
        if has_unfinished_download(target):
            return False
        return (target / "ptbxl_database.csv").exists()

    def get_preprocessing(self) -> dict:
        """Describe the preprocessing, adding the fixed recording duration."""
        preprocessing = super().get_preprocessing()
        preprocessing["segment_duration_seconds"] = self.RECORDING_DURATION_SECONDS
        return preprocessing

    @classmethod
    def _download(cls, data_dir: Path, verbose: bool = True) -> None:
        """Download PTB-XL: record by record from the mirror, or as one archive."""
        if verbose:
            logger.info(f"Downloading PTB-XL dataset to {data_dir}")
            logger.info(f"Note: this fetches {cls.DOWNLOAD_SIZE}; an interrupted download resumes.")
        cls._download_records(
            data_dir,
            lambda: cls._mirrored_files(data_dir),
            mirror=cls.DOWNLOAD_MIRROR,
            desc="Downloading PTB-XL",
            archive_url=cls.DOWNLOAD_URL,
            archive_name="ptb-xl-1.0.3.zip",
            workers=cls.DOWNLOAD_WORKERS,
            verbose=verbose,
            nested_dirs=("ptb-xl-a-large-publicly-available-electrocardiography-dataset-1.0.3",),
        )
        if verbose:
            logger.info("Download complete!")

    @classmethod
    def _mirrored_files(cls, data_dir: Path) -> list[str]:
        """List the metadata tables and record files to fetch.

        The record paths are read from ptbxl_database.csv, the table the
        loader itself follows, rather than from the project's RECORDS index,
        which runs two paths together on one line.

        Raises:
            RuntimeError: If the metadata tables cannot be fetched.
        """
        for name in cls.METADATA_FILES:
            download_file(cls.DOWNLOAD_MIRROR + name, data_dir / name, progress=False)
        table = pd.read_csv(data_dir / "ptbxl_database.csv", usecols=["filename_lr", "filename_hr"])
        records = [*table["filename_lr"], *table["filename_hr"]]
        return [
            *cls.METADATA_FILES,
            *(f"{record}{suffix}" for record in records for suffix in (".hea", ".dat")),
        ]

    def _load_data(self) -> None:
        """Load the PTB-XL dataset into memory."""
        if self.verbose:
            logger.info("Loading PTB-XL dataset...")

        metadata_path = self.data_dir / "ptbxl_database.csv"
        if not metadata_path.exists():
            raise FileNotFoundError(
                self._not_found_message(self.data_dir, "ptbxl_database.csv and the record folders")
            )

        self.metadata_df = pd.read_csv(metadata_path, index_col="ecg_id")
        self.metadata_df.scp_codes = self.metadata_df.scp_codes.apply(ast.literal_eval)

        scp_path = self.data_dir / "scp_statements.csv"
        if scp_path.exists():
            self.scp_statements = pd.read_csv(scp_path, index_col=0)
        else:
            self.scp_statements = None

        if self.folds is not None:
            self.metadata_df = self.metadata_df[self.metadata_df.strat_fold.isin(self.folds)]
            if self.verbose:
                logger.info(f"Using folds: {self.folds} ({len(self.metadata_df)} records)")

        label_columns = self._get_label_columns()

        if self._leads:
            lead_indices = [self.LEADS.index(lead) for lead in self._leads if lead in self.LEADS]
        else:
            lead_indices = list(range(len(self.LEADS)))

        iterator = (
            tqdm(self.metadata_df.iterrows(), total=len(self.metadata_df), desc="Loading records")
            if self.verbose
            else self.metadata_df.iterrows()
        )

        attempted_records = 0
        loaded_records = 0
        for ecg_id, row in iterator:
            attempted_records += 1
            try:
                record_path = (
                    self.data_dir / row.filename_hr
                    if self.use_high_resolution
                    else self.data_dir / row.filename_lr
                )
                record_path = str(record_path).replace(".hea", "").replace(".dat", "")

                record = wfdb.rdrecord(record_path)
                signal = record.p_signal.T

                signal = signal[lead_indices, :]

                if self.source_sampling_rate != self.sampling_rate:
                    signal = self.standardizer.resample(signal, self.source_sampling_rate)

                signal = self.standardizer.normalize(signal)

                labels = self._extract_labels(row.scp_codes, label_columns)

                self.signals.append(signal.astype(np.float32))
                self.labels.append(labels)
                self.record_names.append(str(ecg_id))
                self.patient_ids.append(str(row.patient_id))
                loaded_records += 1

            except self.RECORD_ERRORS as e:
                self._record_load_error(str(ecg_id), e)
                continue

        self._finish_load(loaded_records, attempted_records)

        if self.verbose:
            logger.info(f"Successfully loaded {len(self.signals)} records")
            self._print_class_distribution()

    def _get_label_columns(self) -> list[str]:
        """Get the label columns based on the task."""
        if self.task == "superclass":
            return self.CLASS_LABELS_SUPERCLASS
        elif self.task == "subclass":
            return self.CLASS_LABELS_SUBCLASS
        elif self.task == "form":
            return self.CLASS_LABELS_FORM
        elif self.task == "rhythm":
            return self.CLASS_LABELS_RHYTHM
        elif self.task == "diagnostic":
            if self.scp_statements is not None:
                diag = self.scp_statements[self.scp_statements.diagnostic == 1.0]
                return sorted(diag.index.tolist())
            raise ValueError(
                "scp_statements.csv is required for the 'diagnostic' task. "
                "Ensure the dataset is fully downloaded."
            )
        else:  # all
            if self.scp_statements is not None:
                return sorted(self.scp_statements.index.tolist())
            # Fallback: union of hardcoded lists
            all_labels = (
                self.CLASS_LABELS_SUPERCLASS
                + self.CLASS_LABELS_SUBCLASS
                + self.CLASS_LABELS_FORM
                + self.CLASS_LABELS_RHYTHM
            )
            return sorted(set(all_labels))

    def _extract_labels(self, scp_codes: dict[str, float], label_columns: list[str]) -> np.ndarray:
        """Extract labels from SCP codes."""
        if self.multi_label:
            labels = np.zeros(len(label_columns), dtype=np.float32)
            if self.task == "superclass" and self.scp_statements is not None:
                for code, likelihood in scp_codes.items():
                    if likelihood >= 50.0 and code in self.scp_statements.index:
                        stmt = self.scp_statements.loc[code]
                        if pd.notna(stmt.diagnostic_class):
                            superclass = stmt.diagnostic_class
                            if superclass in label_columns:
                                idx = label_columns.index(superclass)
                                labels[idx] = 1.0
            else:
                for code, likelihood in scp_codes.items():
                    if code in label_columns:
                        idx = label_columns.index(code)
                        labels[idx] = 1.0 if likelihood >= 50.0 else 0.0
            return labels
        else:
            if self.task == "superclass" and self.scp_statements is not None:
                superclass_counts = {cls: 0.0 for cls in self.CLASS_LABELS_SUPERCLASS}
                for code, likelihood in scp_codes.items():
                    if likelihood >= 50.0 and code in self.scp_statements.index:
                        stmt = self.scp_statements.loc[code]
                        if pd.notna(stmt.diagnostic_class):
                            superclass = stmt.diagnostic_class
                            if superclass in superclass_counts:
                                superclass_counts[superclass] += likelihood
                if any(v > 0 for v in superclass_counts.values()):
                    primary_class = max(superclass_counts, key=superclass_counts.get)
                    return np.array(
                        self.CLASS_LABELS_SUPERCLASS.index(primary_class), dtype=np.int64
                    )
            for code, likelihood in scp_codes.items():
                if code in label_columns and likelihood >= 50.0:
                    return np.array(label_columns.index(code), dtype=np.int64)
            return np.array(0, dtype=np.int64)

    def _print_class_distribution(self) -> None:
        """Print class distribution statistics."""
        if len(self.labels) == 0:
            return

        label_columns = self._get_label_columns()
        logger.info(f"\nClass distribution ({self.task}):")

        if self.multi_label:
            labels_array = np.stack(self.labels)
            for i, class_name in enumerate(label_columns):
                count = int(labels_array[:, i].sum())
                if count > 0:
                    logger.info(f"  {class_name}: {count}")
        else:
            unique, counts = np.unique(self.labels, return_counts=True)
            for label_idx, count in zip(unique, counts, strict=False):
                if label_idx < len(label_columns):
                    logger.info(f"  {label_columns[label_idx]}: {count}")

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
        """Get a sample from the dataset.

        Args:
            idx: Index of the sample to get

        Returns:
            Tuple of (ecg_signal, label) where:
                - ecg_signal: Tensor of shape (num_leads, signal_length)
                - label: Tensor of shape (num_classes,) for multi-label or scalar for single-label
        """
        signal = self.signals[idx]
        label = self.labels[idx]

        signal = convert_to_tensor(signal, dtype=torch.float32)

        if self.multi_label:
            label = torch.tensor(label, dtype=torch.float32)
        else:
            label = torch.tensor(label, dtype=torch.long)

        if self.transform is not None:
            signal = self.transform(signal)

        if self.target_transform is not None:
            label = self.target_transform(label)

        return signal, label

    def __len__(self) -> int:
        return len(self.labels)

    @property
    def num_classes(self) -> int:
        """Get the number of classes in the dataset."""
        return len(self._get_label_columns())

    @property
    def class_names(self) -> list[str]:
        """Get the names of the classes in the dataset."""
        return self._get_label_columns()

    @property
    def group_ids(self) -> list[str]:
        """Patient identifiers, used so a patient never spans multiple splits.

        PTB-XL contains multiple ECGs for some patients, so grouping by ecg_id
        would place different recordings of the same patient in train and test.
        The official strat_fold assignment is likewise made per patient.

        Returns:
            One patient identifier per loaded record.
        """
        return list(self.patient_ids)

    def get_record_info(self, idx: int) -> dict:
        """Get record information for a specific sample.

        Args:
            idx: Index of the sample

        Returns:
            Dictionary containing record information
        """
        ecg_id = int(self.record_names[idx])
        row = self.metadata_df.loc[ecg_id]

        return {
            "ecg_id": ecg_id,
            "patient_id": row.patient_id,
            "age": row.age,
            "sex": row.sex,
            "signal_shape": self.signals[idx].shape,
            "scp_codes": row.scp_codes,
            "strat_fold": row.strat_fold,
            "labels": self.labels[idx].tolist() if self.multi_label else int(self.labels[idx]),
        }

    def get_class_distribution(self) -> dict[str, int]:
        """Get the distribution of classes in the dataset.

        Returns:
            Dictionary mapping class names to their counts
        """
        label_columns = self._get_label_columns()
        distribution = {name: 0 for name in label_columns}

        if self.multi_label:
            labels_array = np.stack(self.labels)
            for i, class_name in enumerate(label_columns):
                distribution[class_name] = int(labels_array[:, i].sum())
        else:
            for lbl in self.labels:
                label_idx = int(lbl) if isinstance(lbl, np.ndarray) else lbl
                if label_idx < len(label_columns):
                    distribution[label_columns[label_idx]] += 1

        return distribution

    def get_folds_split(
        self,
        train_folds: list[int] | None = None,
        val_folds: list[int] | None = None,
        test_folds: list[int] | None = None,
    ) -> dict[str, "PTBXLDataset"]:
        """Create train/val/test splits based on stratified folds.

        The PTB-XL dataset comes with 10 pre-defined stratified folds.
        Recommended split: folds 1-8 for training, 9 for validation, 10 for testing.

        Args:
            train_folds: Folds for training (default: 1-8)
            val_folds: Folds for validation (default: 9)
            test_folds: Folds for testing (default: 10)

        Returns:
            Dictionary with 'train', 'val', 'test' PTBXLDataset instances
        """
        if train_folds is None:
            train_folds = list(range(1, 9))
        if val_folds is None:
            val_folds = [9]
        if test_folds is None:
            test_folds = [10]

        common_kwargs = {
            "data_dir": self.data_dir,
            "sampling_rate": self.sampling_rate,
            "task": self.task,
            "use_high_resolution": self.use_high_resolution,
            "leads": self._leads,
            "normalization": self.standardizer.normalization,
            "multi_label": self.multi_label,
            "transform": self.transform,
            "target_transform": self.target_transform,
            "download": False,
            "verbose": self.verbose,
        }

        return {
            "train": PTBXLDataset(folds=train_folds, **common_kwargs),
            "val": PTBXLDataset(folds=val_folds, **common_kwargs),
            "test": PTBXLDataset(folds=test_folds, **common_kwargs),
        }

num_classes property

num_classes: int

Get the number of classes in the dataset.

class_names property

class_names: list[str]

Get the names of the classes in the dataset.

group_ids property

group_ids: list[str]

Patient identifiers, used so a patient never spans multiple splits.

PTB-XL contains multiple ECGs for some patients, so grouping by ecg_id would place different recordings of the same patient in train and test. The official strat_fold assignment is likewise made per patient.

Returns:

Type	Description
list[str]	One patient identifier per loaded record.

is_downloaded classmethod

is_downloaded(data_dir: str | Path | None = None) -> bool

Report whether the extracted metadata table is present.

A directory still holding the archive or a partial download does not count: the archive is deleted once extraction finishes.

Source code in deepecgkit/datasets/ptbxl.py

@classmethod
def is_downloaded(cls, data_dir: str | Path | None = None) -> bool:
    """Report whether the extracted metadata table is present.

    A directory still holding the archive or a partial download does not
    count: the archive is deleted once extraction finishes.
    """
    target = Path(data_dir) if data_dir is not None else cls.get_default_data_dir()
    if has_unfinished_download(target):
        return False
    return (target / "ptbxl_database.csv").exists()

get_preprocessing

get_preprocessing() -> dict

Describe the preprocessing, adding the fixed recording duration.

Source code in deepecgkit/datasets/ptbxl.py

def get_preprocessing(self) -> dict:
    """Describe the preprocessing, adding the fixed recording duration."""
    preprocessing = super().get_preprocessing()
    preprocessing["segment_duration_seconds"] = self.RECORDING_DURATION_SECONDS
    return preprocessing

get_record_info

get_record_info(idx: int) -> dict

Get record information for a specific sample.

Parameters:

Name	Type	Description	Default
idx	int	Index of the sample	required

Returns:

Type	Description
dict	Dictionary containing record information

Source code in deepecgkit/datasets/ptbxl.py

def get_record_info(self, idx: int) -> dict:
    """Get record information for a specific sample.

    Args:
        idx: Index of the sample

    Returns:
        Dictionary containing record information
    """
    ecg_id = int(self.record_names[idx])
    row = self.metadata_df.loc[ecg_id]

    return {
        "ecg_id": ecg_id,
        "patient_id": row.patient_id,
        "age": row.age,
        "sex": row.sex,
        "signal_shape": self.signals[idx].shape,
        "scp_codes": row.scp_codes,
        "strat_fold": row.strat_fold,
        "labels": self.labels[idx].tolist() if self.multi_label else int(self.labels[idx]),
    }

get_class_distribution

get_class_distribution() -> dict[str, int]

Get the distribution of classes in the dataset.

Returns:

Type	Description
dict[str, int]	Dictionary mapping class names to their counts

Source code in deepecgkit/datasets/ptbxl.py

def get_class_distribution(self) -> dict[str, int]:
    """Get the distribution of classes in the dataset.

    Returns:
        Dictionary mapping class names to their counts
    """
    label_columns = self._get_label_columns()
    distribution = {name: 0 for name in label_columns}

    if self.multi_label:
        labels_array = np.stack(self.labels)
        for i, class_name in enumerate(label_columns):
            distribution[class_name] = int(labels_array[:, i].sum())
    else:
        for lbl in self.labels:
            label_idx = int(lbl) if isinstance(lbl, np.ndarray) else lbl
            if label_idx < len(label_columns):
                distribution[label_columns[label_idx]] += 1

    return distribution

get_folds_split

get_folds_split(
    train_folds: list[int] | None = None,
    val_folds: list[int] | None = None,
    test_folds: list[int] | None = None,
) -> dict[str, PTBXLDataset]

Create train/val/test splits based on stratified folds.

The PTB-XL dataset comes with 10 pre-defined stratified folds. Recommended split: folds 1-8 for training, 9 for validation, 10 for testing.

Parameters:

Name	Type	Description	Default
train_folds	list[int] | None	Folds for training (default: 1-8)	None
val_folds	list[int] | None	Folds for validation (default: 9)	None
test_folds	list[int] | None	Folds for testing (default: 10)	None

Returns:

Type	Description
dict[str, PTBXLDataset]	Dictionary with 'train', 'val', 'test' PTBXLDataset instances

Source code in deepecgkit/datasets/ptbxl.py

def get_folds_split(
    self,
    train_folds: list[int] | None = None,
    val_folds: list[int] | None = None,
    test_folds: list[int] | None = None,
) -> dict[str, "PTBXLDataset"]:
    """Create train/val/test splits based on stratified folds.

    The PTB-XL dataset comes with 10 pre-defined stratified folds.
    Recommended split: folds 1-8 for training, 9 for validation, 10 for testing.

    Args:
        train_folds: Folds for training (default: 1-8)
        val_folds: Folds for validation (default: 9)
        test_folds: Folds for testing (default: 10)

    Returns:
        Dictionary with 'train', 'val', 'test' PTBXLDataset instances
    """
    if train_folds is None:
        train_folds = list(range(1, 9))
    if val_folds is None:
        val_folds = [9]
    if test_folds is None:
        test_folds = [10]

    common_kwargs = {
        "data_dir": self.data_dir,
        "sampling_rate": self.sampling_rate,
        "task": self.task,
        "use_high_resolution": self.use_high_resolution,
        "leads": self._leads,
        "normalization": self.standardizer.normalization,
        "multi_label": self.multi_label,
        "transform": self.transform,
        "target_transform": self.target_transform,
        "download": False,
        "verbose": self.verbose,
    }

    return {
        "train": PTBXLDataset(folds=train_folds, **common_kwargs),
        "val": PTBXLDataset(folds=val_folds, **common_kwargs),
        "test": PTBXLDataset(folds=test_folds, **common_kwargs),
    }

UnifiedAFDataset

Bases: 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

A source that cannot be loaded is an error: a dataset built from two of its three sources is a different dataset under the same name. Pass skip_unavailable=True to build it from whichever sources do load, or name the wanted sources in datasets.

Source code in deepecgkit/datasets/unified_af.py

@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

    A source that cannot be loaded is an error: a dataset built from two of
    its three sources is a different dataset under the same name. Pass
    skip_unavailable=True to build it from whichever sources do load, or name
    the wanted sources in datasets.
    """

    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,
    }
    DOWNLOAD_SIZE: ClassVar[str | None] = "about 4.4 GB across its three sources"

    @classmethod
    def _source_dir(
        cls, dataset_class: type[BaseECGDataset], data_dir: str | Path | None = None
    ) -> Path:
        """Locate one source dataset.

        Each source lives in its own default directory, so a copy fetched for
        the standalone dataset is shared. A data_dir given explicitly is the
        folder that holds the sources instead, each under its usual name.
        """
        default = dataset_class.get_default_data_dir()
        if data_dir is None or Path(data_dir) == cls.get_default_data_dir():
            return default
        return Path(data_dir) / default.name

    @classmethod
    def is_downloaded(cls, data_dir: str | Path | None = None) -> bool:
        """Report whether every source dataset is present in its own directory."""
        return all(
            dataset_class.is_downloaded(cls._source_dir(dataset_class, data_dir))
            for dataset_class in cls.AVAILABLE_DATASETS.values()
        )

    def _resolve_dataset_dir(self, dataset_name: str) -> Path:
        return self._source_dir(self.AVAILABLE_DATASETS[dataset_name], self.data_dir)

    def _clear_dataset(self) -> None:
        """Clear every source's directory; nothing is kept in a directory of its own."""
        for dataset_name in self.dataset_names:
            directory = self._resolve_dataset_dir(dataset_name)
            if directory.exists():
                self.AVAILABLE_DATASETS[dataset_name]._require_clearable(directory)
                logger.info(f"Clearing existing dataset at {directory}...")
                shutil.rmtree(directory)

    def __init__(
        self,
        data_dir: str | Path | None = None,
        sampling_rate: int = 300,
        segment_duration_seconds: float = 10.0,
        datasets: list[str] | None = None,
        binary_classification: bool = False,
        normalization: str = "zscore",
        transform: Callable[[torch.Tensor], torch.Tensor] | None = None,
        target_transform: Callable[[torch.Tensor], torch.Tensor] | None = None,
        download: bool = False,
        force_download: bool = False,
        verbose: bool = True,
        dataset_kwargs: dict[str, dict] | None = None,
        skip_unavailable: bool = False,
    ):
        self.skip_unavailable = skip_unavailable
        self.segment_duration_seconds = segment_duration_seconds
        self.binary_classification = binary_classification
        self.normalization = normalization
        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: list[Any] = []
        self.dataset_sizes: list[int] = []
        self.loaded_dataset_names: list[str] = []
        self.record_names: list[str] = []

        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) -> None:
        for dataset_name in self.dataset_names:
            dataset_class = self.AVAILABLE_DATASETS[dataset_name]
            dataset_dir = self._resolve_dataset_dir(dataset_name)

            if dataset_class.is_downloaded(dataset_dir):
                if self.verbose:
                    logger.info(f"\n{'=' * 60}")
                    logger.info(f"Skipping {dataset_name}: already exists at {dataset_dir}")
                    logger.info(f"{'=' * 60}")
                continue

            if self.verbose:
                logger.info(f"\n{'=' * 60}")
                logger.info(f"Downloading dataset: {dataset_name}")
                logger.info(f"{'=' * 60}")

            dataset_class.download_to(dataset_dir, verbose=self.verbose)

    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,
            "normalization": self.normalization,
            "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) -> None:
        if self.verbose:
            logger.info("\nLoading unified AF dataset...")

        for dataset_name in self.dataset_names:
            if self.verbose:
                logger.info(f"\n{'=' * 60}")
                logger.info(f"Loading: {dataset_name}")
                logger.info(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:
                        logger.info(
                            f"  Dropped {dropped} samples with unmappable labels (Other/Noisy)"
                        )

                self.datasets.append(dataset)
                self.dataset_sizes.append(len(dataset))
                self.loaded_dataset_names.append(dataset_name)

                if self.verbose:
                    logger.info(f"Loaded {len(dataset)} samples from {dataset_name}")
                    if hasattr(dataset, "get_class_distribution"):
                        dist = dataset.get_class_distribution()
                        logger.info(f"Class distribution: {dist}")

            except Exception as e:
                if not self.skip_unavailable:
                    raise RuntimeError(
                        f"Could not load '{dataset_name}', one of the sources of the unified "
                        f"AF dataset: {e}"
                    ) from e
                logger.warning("Failed to load %s: %s", dataset_name, e)
                continue

        if len(self.datasets) == 0:
            raise RuntimeError("No datasets were successfully loaded")

        self.concat_dataset: ConcatDataset = ConcatDataset(self.datasets)
        self._collect_record_names()

        if self.verbose:
            logger.info(f"\n{'=' * 60}")
            logger.info("Unified dataset created")
            logger.info(f"Total datasets: {len(self.datasets)}")
            logger.info(f"Total samples: {len(self.concat_dataset)}")
            logger.info(f"{'=' * 60}")
            self._print_overall_distribution()

    def _collect_record_names(self) -> None:
        """Build globally unique record identifiers across the merged source datasets.

        Record names are namespaced by source dataset so that identically named
        records in different sources are never treated as the same group. These
        identifiers are what ECGDataModule uses to keep every segment of a
        recording inside a single split; without them the merged dataset would be
        split at segment level and leak between train and test.

        Raises:
            RuntimeError: If any source dataset does not expose usable record names.
        """
        record_names: list[str] = []
        for dataset_name, dataset in zip(self.loaded_dataset_names, self.datasets, strict=False):
            source_records = getattr(dataset, "record_names", None)
            if not source_records or len(source_records) != len(dataset):
                raise RuntimeError(
                    f"Source dataset '{dataset_name}' exposes "
                    f"{0 if not source_records else len(source_records)} record names for "
                    f"{len(dataset)} samples. Record-level splitting cannot be performed "
                    "safely, and segment-level splitting would leak recordings between "
                    "train and test."
                )
            record_names.extend(f"{dataset_name}/{name}" for name in source_records)

        self.record_names = record_names

    def _print_overall_distribution(self) -> None:
        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)
        logger.info("\nOverall class distribution:")
        for label_idx, count in zip(unique, counts, strict=False):
            class_name = self.class_names[label_idx]
            percentage = (count / len(all_labels)) * 100
            logger.info(f"  {class_name}: {count} ({percentage:.1f}%)")

    def stratification_labels(self) -> np.ndarray | None:
        """Return every source's labels in order, without reading a signal."""
        parts = [dataset.stratification_labels() for dataset in self.datasets]
        if not parts or any(part is None for part in parts):
            return None
        joined: np.ndarray = np.concatenate(parts)
        return joined

    def __getitem__(self, idx: int) -> tuple[torch.Tensor, torch.Tensor]:
        return self.concat_dataset[idx]

    def __len__(self) -> int:
        return len(self.concat_dataset)

    def get_preprocessing(self) -> dict:
        """Describe the preprocessing shared by every source dataset."""
        preprocessing = super().get_preprocessing()
        preprocessing["normalization"] = self.normalization
        return preprocessing

    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, strict=False)
        }

is_downloaded classmethod

is_downloaded(data_dir: str | Path | None = None) -> bool

Report whether every source dataset is present in its own directory.

Source code in deepecgkit/datasets/unified_af.py

@classmethod
def is_downloaded(cls, data_dir: str | Path | None = None) -> bool:
    """Report whether every source dataset is present in its own directory."""
    return all(
        dataset_class.is_downloaded(cls._source_dir(dataset_class, data_dir))
        for dataset_class in cls.AVAILABLE_DATASETS.values()
    )

stratification_labels

stratification_labels() -> np.ndarray | None

Return every source's labels in order, without reading a signal.

Source code in deepecgkit/datasets/unified_af.py

def stratification_labels(self) -> np.ndarray | None:
    """Return every source's labels in order, without reading a signal."""
    parts = [dataset.stratification_labels() for dataset in self.datasets]
    if not parts or any(part is None for part in parts):
        return None
    joined: np.ndarray = np.concatenate(parts)
    return joined

get_preprocessing

get_preprocessing() -> dict

Describe the preprocessing shared by every source dataset.

Source code in deepecgkit/datasets/unified_af.py

def get_preprocessing(self) -> dict:
    """Describe the preprocessing shared by every source dataset."""
    preprocessing = super().get_preprocessing()
    preprocessing["normalization"] = self.normalization
    return preprocessing

File Reading and Inference Preprocessing

read_ecg_recording

read_ecg_recording(path: str | Path) -> np.ndarray

Read an ECG recording from a supported file format.

Parameters:

Name	Type	Description	Default
path	str | Path	Path to a .npy, .csv, or WFDB (.dat/.hea) recording	required

Returns:

Type	Description
ndarray	The decoded signal array. WFDB records are transposed to put leads
ndarray	first; .npy and .csv arrays are returned as stored.

Raises:

Type	Description
ValueError	If the file suffix is not one of SUPPORTED_ECG_SUFFIXES.

Source code in deepecgkit/datasets/io.py

def read_ecg_recording(path: str | Path) -> np.ndarray:
    """Read an ECG recording from a supported file format.

    Args:
        path: Path to a .npy, .csv, or WFDB (.dat/.hea) recording

    Returns:
        The decoded signal array. WFDB records are transposed to put leads
        first; .npy and .csv arrays are returned as stored.

    Raises:
        ValueError: If the file suffix is not one of SUPPORTED_ECG_SUFFIXES.
    """
    return read_ecg_recording_with_rate(path)[0]

read_ecg_recording_with_rate

read_ecg_recording_with_rate(
    path: str | Path,
) -> tuple[np.ndarray, float | None]

Read an ECG recording together with its sampling rate, where the file records one.

Parameters:

Name	Type	Description	Default
path	str | Path	Path to a .npy, .csv, or WFDB (.dat/.hea) recording	required

Returns:

Type	Description
ndarray	Tuple of the decoded signal array and its sampling rate in Hz. WFDB
float | None	records carry their rate in the header; .npy and .csv files carry none,
tuple[ndarray, float | None]	so the rate is None for them. WFDB records are transposed to put leads
tuple[ndarray, float | None]	first; .npy and .csv arrays are returned as stored.

Raises:

Type	Description
ValueError	If the file suffix is not one of SUPPORTED_ECG_SUFFIXES.

Source code in deepecgkit/datasets/io.py

def read_ecg_recording_with_rate(path: str | Path) -> tuple[np.ndarray, float | None]:
    """Read an ECG recording together with its sampling rate, where the file records one.

    Args:
        path: Path to a .npy, .csv, or WFDB (.dat/.hea) recording

    Returns:
        Tuple of the decoded signal array and its sampling rate in Hz. WFDB
        records carry their rate in the header; .npy and .csv files carry none,
        so the rate is None for them. WFDB records are transposed to put leads
        first; .npy and .csv arrays are returned as stored.

    Raises:
        ValueError: If the file suffix is not one of SUPPORTED_ECG_SUFFIXES.
    """
    file_path = Path(path)
    if file_path.suffix not in SUPPORTED_ECG_SUFFIXES:
        raise ValueError(
            f"Unsupported input format: {file_path.suffix}. "
            f"Supported formats: .npy, .csv, .dat/.hea (WFDB)"
        )

    if file_path.suffix == ".npy":
        return np.load(str(file_path)), None
    if file_path.suffix == ".csv":
        return _read_csv_signal(file_path), None

    record = wfdb.rdrecord(str(file_path.with_suffix("")))
    return record.p_signal.T, float(record.fs)

as_model_input

as_model_input(
    data: ndarray, input_channels: int | None = None
) -> np.ndarray

Reshape an ECG array to the (batch, leads, time) layout models expect.

A one-dimensional array is treated as a single single-lead recording. A two-dimensional array is ambiguous, so the lead count decides the orientation: whichever axis matches input_channels becomes the lead axis. When the lead count is unknown or one, rows are treated as independent single-lead recordings.

Parameters:

Name	Type	Description	Default
data	ndarray	Signal array of one, two, or three dimensions	required
input_channels	int | None	Number of leads the target model consumes, used to disambiguate two-dimensional input	None

Returns:

Type	Description
ndarray	Array of shape (batch, leads, time).

Raises:

Type	Description
ValueError	If a two-dimensional array matches neither orientation, or the array has more than three dimensions.

Source code in deepecgkit/datasets/io.py

def as_model_input(data: np.ndarray, input_channels: int | None = None) -> np.ndarray:
    """Reshape an ECG array to the (batch, leads, time) layout models expect.

    A one-dimensional array is treated as a single single-lead recording. A
    two-dimensional array is ambiguous, so the lead count decides the
    orientation: whichever axis matches input_channels becomes the lead axis.
    When the lead count is unknown or one, rows are treated as independent
    single-lead recordings.

    Args:
        data: Signal array of one, two, or three dimensions
        input_channels: Number of leads the target model consumes, used to
            disambiguate two-dimensional input

    Returns:
        Array of shape (batch, leads, time).

    Raises:
        ValueError: If a two-dimensional array matches neither orientation, or
            the array has more than three dimensions.
    """
    if data.ndim == 1:
        return data.reshape(1, 1, -1)

    if data.ndim == 2:
        if input_channels is not None and data.shape[0] == input_channels:
            return data.reshape(1, data.shape[0], -1)
        if input_channels is not None and data.shape[1] == input_channels:
            return data.T.reshape(1, input_channels, -1)
        if input_channels in (None, 1):
            return data.reshape(data.shape[0], 1, -1)
        raise ValueError(
            f"Input of shape {data.shape} is not compatible with a model expecting "
            f"{input_channels} leads."
        )

    if data.ndim == 3:
        return data

    raise ValueError(
        f"Input of shape {data.shape} has {data.ndim} dimensions; expected at most three "
        f"(batch, leads, time)."
    )

prepare_recording

prepare_recording(
    recording: ndarray,
    preprocessing: Mapping[str, Any] | None,
    source_sampling_rate: float | None = None,
) -> PreparedRecording

Put one raw recording through a model's training-time preprocessing.

The steps mirror what the dataset loaders do: resample to the training rate, cut into training-length segments, and normalise each segment.

Parameters:

Name	Type	Description	Default
recording	ndarray	Signal of shape (leads, samples)	required
preprocessing	Mapping[str, Any] | None	Mapping as returned by BaseECGDataset.get_preprocessing and recorded in checkpoints, with sampling_rate, segment_duration_seconds and normalization. None leaves the recording untouched and returns it as a single segment.	required
source_sampling_rate	float | None	Rate the recording was sampled at, in Hz. None assumes it is already at the training rate.	None

Returns:

Type	Description
PreparedRecording	The prepared segments, with notes describing what was done.

Source code in deepecgkit/datasets/io.py

def prepare_recording(
    recording: np.ndarray,
    preprocessing: Mapping[str, Any] | None,
    source_sampling_rate: float | None = None,
) -> PreparedRecording:
    """Put one raw recording through a model's training-time preprocessing.

    The steps mirror what the dataset loaders do: resample to the training
    rate, cut into training-length segments, and normalise each segment.

    Args:
        recording: Signal of shape (leads, samples)
        preprocessing: Mapping as returned by BaseECGDataset.get_preprocessing
            and recorded in checkpoints, with sampling_rate,
            segment_duration_seconds and normalization. None leaves the
            recording untouched and returns it as a single segment.
        source_sampling_rate: Rate the recording was sampled at, in Hz. None
            assumes it is already at the training rate.

    Returns:
        The prepared segments, with notes describing what was done.
    """
    signal = np.asarray(recording, dtype=np.float64)
    if signal.ndim == 1:
        signal = signal[np.newaxis, :]
    notes: list[str] = []

    if not preprocessing:
        return PreparedRecording(
            segments=signal[np.newaxis].astype(np.float32),
            start_seconds=np.zeros(1),
            sampling_rate=source_sampling_rate,
            notes=notes,
        )

    target_rate = preprocessing.get("sampling_rate")
    normalization = preprocessing.get("normalization") or "none"
    duration = preprocessing.get("segment_duration_seconds")
    standardizer = ECGStandardizer(
        target_sampling_rate=target_rate or 0, normalization=normalization
    )

    if target_rate and source_sampling_rate and source_sampling_rate != target_rate:
        signal = standardizer.resample(signal, source_sampling_rate)
        notes.append(f"resampled from {source_sampling_rate:g} Hz to {target_rate:g} Hz")
    elif target_rate and not source_sampling_rate:
        notes.append(f"assumed to be sampled at {target_rate:g} Hz, the rate the model saw")

    starts = np.zeros(1)
    segments = signal[np.newaxis]
    if duration and target_rate:
        segmenter = ECGSegmenter(segment_duration_seconds=duration, sampling_rate=target_rate)
        cut, start_indices = segmenter.segment(signal)
        if len(cut):
            segments = cut
            starts = np.asarray(start_indices, dtype=np.float64) / float(target_rate)
            if len(cut) > 1:
                notes.append(f"cut into {len(cut)} segments of {duration:g} s")
        else:
            seconds = signal.shape[-1] / float(target_rate)
            notes.append(
                f"only {seconds:.1f} s long, shorter than the {duration:g} s segments the "
                f"model was trained on, so it is used whole"
            )

    if normalization != "none":
        segments = np.stack([standardizer.normalize(segment) for segment in segments])
        notes.append(f"{normalization} normalised")

    return PreparedRecording(
        segments=segments.astype(np.float32),
        start_seconds=starts,
        sampling_rate=float(target_rate) if target_rate else source_sampling_rate,
        notes=notes,
    )

PreparedRecording dataclass

One recording after the preprocessing a model was trained with.

Attributes:

Name	Type	Description
segments	ndarray	Array of shape (num_segments, leads, samples), float32
start_seconds	ndarray	Start of each segment within the recording, in seconds
sampling_rate	float | None	Rate of the segments in Hz, or None when unknown
notes	list[str]	What was applied or assumed, for reporting to the user

Source code in deepecgkit/datasets/io.py

@dataclass
class PreparedRecording:
    """One recording after the preprocessing a model was trained with.

    Attributes:
        segments: Array of shape (num_segments, leads, samples), float32
        start_seconds: Start of each segment within the recording, in seconds
        sampling_rate: Rate of the segments in Hz, or None when unknown
        notes: What was applied or assumed, for reporting to the user
    """

    segments: np.ndarray
    start_seconds: np.ndarray
    sampling_rate: float | None = None
    notes: list[str] = field(default_factory=list)