Dataset

ecgbench.dataset — the single PyTorch Dataset every supported dataset loads through, plus the collate function and the two errors it raises.

The read-time adapters (window=, leads=, units=, transform=) shape the returned tensor only, in that order. They never touch the source files, the exported fold CSVs, or validation — which reads whole records through its own window-less copy of _load_signal in ecgbench/validation/engine.py.

dataset

Unified PyTorch Dataset for loading any ECG dataset supported by ECGBench.

Uses the dataset's YAML config to determine how to load signals and metadata. Adding a new dataset requires only a config file — no changes to this class.

ECGDataset

ECGDataset(dataset: str | Any, split: str | None = 'train', version: str = 'clean', data_path: Path | str | None = None, sampling_rate: int | None = None, fold_numbers: list[int] | None = None, transform: Callable | None = None, metadata_source: str = 'hf', labels: bool = False, leads: list[str] | None = None, units: str = 'mV', window: tuple[int, int | None] | None = None)

Bases: Dataset

PyTorch Dataset for loading any ECG dataset supported by ECGBench.

This class uses the dataset's YAML config to determine how to load signals and metadata. Adding a new dataset requires only a config file.

Parameters:

Name	Type	Description	Default
dataset	str | Any	Dataset slug (e.g., "ptbxl") or a DatasetConfig object	required
split	str | None	"train", "val", "test", or None. None selects records purely by fold_numbers, ignoring the default split boundaries — use it for custom cross-validation. It requires fold_numbers.	'train'
version	str	"clean" (default) or "original"	'clean'
data_path	Path | str | None	Path to the dataset's signal files on disk. If None, attempts auto-download from config.download_url.	None
sampling_rate	int | None	Which sampling rate to load (default: config.default_sampling_rate)	None
fold_numbers	list[int] | None	Specific fold(s) to load. None = all folds for the split. With a named split the folds must belong to it (folds 1-8 are train, 9 val, 10 test); pass split=None to cross that boundary.	None
window	tuple[int, int | None] | None	(start, length) in samples, e.g. (0, 2500) for the first 2500 and (2500, 2500) for the next. length=None reads to the end. Pushed down into the reader, so only these samples are decoded — much faster than cropping afterwards on long records, and unlike a lambda transform it survives a DataLoader with num_workers>0 under the "spawn" start method. Raises WindowOutOfRangeError if the window does not fit.	None
transform	Callable | None	Optional callable applied to the signal tensor, after window, leads and units	None
metadata_source	str	"hf" (download fold CSVs from HuggingFace) or "local".	'hf'
leads	list[str] | None	select and reorder leads by name, e.g. ["I", "II", "V5"]. Case-insensitive; needs lead_names in the dataset's config. The signal's first dimension becomes len(leads). Where a release uses more than one layout (zzu_pecg, mitdb) the names are re-resolved per record, and a record whose layout lacks a requested lead raises rather than substituting another.	None
units	str	"mV" (default) or "uV" — applied after lead selection and before transform. Never affects validation or the exported folds.	'mV'
labels	bool	attach per-record labels and metadata as sample["labels"]. Needs a local copy of the source dataset — fold CSVs on the Hub carry identifiers only, never labels.	False

Example

>>> train_ds = ECGDataset("ptbxl", split="train", data_path="/data/ptb-xl/1.0.3/")
>>> loader = DataLoader(train_ds, batch_size=32, collate_fn=ecg_collate_fn)

>>> ds = ECGDataset("ptbxl", split="train", data_path="...", labels=True)
>>> ds[0]["labels"]["superclasses"]
['NORM']

Source code in ecgbench/dataset.py

def __init__(
    self,
    dataset: str | Any,  # str or DatasetConfig
    split: str | None = "train",
    version: str = "clean",
    data_path: Path | str | None = None,
    sampling_rate: int | None = None,
    fold_numbers: list[int] | None = None,
    transform: Callable | None = None,
    metadata_source: str = "hf",
    labels: bool = False,
    leads: list[str] | None = None,
    units: str = "mV",
    window: tuple[int, int | None] | None = None,
):
    super().__init__()

    from ecgbench.config import DatasetConfig, load_config

    if isinstance(dataset, str):
        self.config = load_config(dataset)
    elif isinstance(dataset, DatasetConfig):
        self.config = dataset
    else:
        raise TypeError(f"dataset must be str or DatasetConfig, got {type(dataset)}")

    self.split = split.lower() if isinstance(split, str) else None
    self.version = version
    self.sampling_rate = sampling_rate or self.config.default_sampling_rate
    self.transform = transform
    self.metadata_source = metadata_source
    self.window = _resolve_window(window)

    if self.split is not None and self.split not in ("train", "val", "test"):
        raise ValueError(f"split must be 'train', 'val', 'test', or None, got '{split}'")
    # split=None selects purely by fold number, across the default split
    # boundaries — the only way to express custom cross-validation, since
    # fold_numbers alone is scoped to one split's directory.
    if self.split is None and not fold_numbers:
        raise ValueError(
            "split=None selects records by fold number across all splits, so "
            "fold_numbers is required with it. For example "
            "ECGDataset(..., split=None, fold_numbers=[1, 2, 3])."
        )
    if self.version not in ("clean", "original"):
        raise ValueError(f"version must be 'clean' or 'original', got '{version}'")

    # Resolve signal data path
    from ecgbench.download import resolve_data_path

    self.data_path = resolve_data_path(data_path, self.config)

    # Load fold metadata
    self.metadata_df = self._load_metadata(fold_numbers)

    # Determine signal path column
    self.signal_col = self.config.signal_path_columns.get(self.sampling_rate)
    if not self.signal_col:
        raise ValueError(
            f"No signal_path_column for rate {self.sampling_rate}. "
            f"Available: {list(self.config.signal_path_columns.keys())}"
        )

    # Lead selection and output units are read-time adapters: they shape the
    # tensor __getitem__ returns and never touch the files, the fold CSVs, or
    # validation — a record excluded for a bad V6 stays excluded even if you
    # never load V6.
    self._unit_factor = _resolve_units(units, self.config.signal_units)
    if self.config.signal_units.strip().lower() in _UNIT_FACTORS:
        self.units = "mV" if self._unit_factor == 1.0 else "uV"
    else:
        # Report what the samples actually are, not the mV default nobody
        # chose. sample["units"] is what a user checks before plotting.
        self.units = self.config.signal_units

    self.lead_names = tuple(self.config.lead_names) if self.config.lead_names else None
    self._lead_index: list[int] | None = None
    # What the user asked for, kept so a record storing a different layout can
    # be re-resolved against it at read time. See _lead_index_for().
    self._requested_leads: list[str] | None = None
    self._declared_n_leads = len(self.config.lead_names) if self.config.lead_names else None
    self._alt_lead_index: dict[int, list[int]] = {}
    # Per-record-path indices, for a release whose records store the same
    # number of leads under different names. See _lead_index_for().
    self._path_lead_index: dict[str, list[int]] = {}
    if leads is not None:
        # With several layouts in play the declared order answers for none of
        # them, so the request is checked against the union — a name in no
        # layout is a typo and fails here — and the indices are re-resolved
        # per record at read time. _lead_index is then only a "selection is
        # active" marker; _lead_index_for never returns it.
        available = _declarable_lead_names(self.config)
        self._lead_index, names = _resolve_leads(leads, available, self.config.slug)
        self._requested_leads = list(leads)
        # The resolved *names* are the same whatever layout a record uses —
        # that is the point of selecting by name — so this stays valid even
        # when the indices differ per record.
        self.lead_names = tuple(names)

    # Labels: loaded once and aligned to this split, never per __getitem__.
    self.labels_df = self._load_labels() if labels else None

__getitem__

__getitem__(idx: int) -> dict[str, Any]

Get a single ECG record with signal and metadata.

The dict holds:

• "signal": torch.Tensor, float32, shape (leads, samples). leads is len(self.lead_names) after any leads= selection, and samples is the window= length when one is set.
• "record_id": record identifier
• "split": the dataset's split, or — when constructed with split=None — this record's own default_split
• "fold": int (if available)
• "labels": dict of label fields (only with labels=True)
• All other metadata columns

Returns:

Type	Description
dict[str, Any]	The record, as described above.

Raises:

Type	Description
WindowOutOfRangeError	window= does not fit this record. Record length is not constant in every dataset, so a window can fit most records and not all.

Source code in ecgbench/dataset.py

def __getitem__(self, idx: int) -> dict[str, Any]:
    """Get a single ECG record with signal and metadata.

    The dict holds:

    - ``"signal"``: torch.Tensor, float32, shape (leads, samples). ``leads``
      is ``len(self.lead_names)`` after any ``leads=`` selection, and
      ``samples`` is the ``window=`` length when one is set.
    - ``"record_id"``: record identifier
    - ``"split"``: the dataset's split, or — when constructed with
      ``split=None`` — this record's own ``default_split``
    - ``"fold"``: int (if available)
    - ``"labels"``: dict of label fields (only with ``labels=True``)
    - All other metadata columns

    Returns:
        The record, as described above.

    Raises:
        WindowOutOfRangeError: ``window=`` does not fit this record. Record
            length is not constant in every dataset, so a window can fit most
            records and not all.
    """
    if idx < 0 or idx >= len(self.metadata_df):
        raise IndexError(
            f"Index {idx} out of range for dataset of size {len(self.metadata_df)}"
        )

    row = self.metadata_df.iloc[idx]

    # Load signal
    signal_path = str(row[self.signal_col])
    if self.config.signal_format == "wfdb":
        signal_path = str(Path(signal_path).with_suffix(""))
    full_path = str(self.data_path / signal_path)

    signal = _load_signal(
        full_path,
        self.config.signal_format,
        self.config.signal_unit_scale,
        self.window,
    )

    if self._lead_index is not None:
        lead_index = self._lead_index_for(
            signal.shape[0], row.get(self.config.record_id_column), full_path
        )
        if signal.shape[0] <= max(lead_index):
            raise ValueError(
                f"Record {row.get(self.config.record_id_column)!r} has "
                f"{signal.shape[0]} leads, too few for the requested "
                f"{list(self.lead_names)}"
            )
        signal = signal[lead_index]

    signal_tensor = torch.from_numpy(signal).float()
    if self._unit_factor != 1.0:
        signal_tensor = signal_tensor * self._unit_factor

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

    # Build result dict
    result: dict[str, Any] = {
        "signal": signal_tensor,
        "record_id": row.get(self.config.record_id_column),
        # With split=None the rows come from several splits, so reporting a
        # single split name would be a lie — give the record's own instead.
        "split": self.split if self.split is not None else row.get("default_split"),
    }

    # Add fold if available
    if "fold" in row.index:
        result["fold"] = int(row["fold"])

    # Labels stay in their own dict: a nested key cannot collide with
    # "signal"/"fold"/"split", and ecg_collate_fn keeps dicts as a list.
    if self.labels_df is not None:
        result["labels"] = self.labels_df.iloc[idx].to_dict()

    # Add all other metadata
    for col in self.metadata_df.columns:
        if col in (self.signal_col, self.config.record_id_column, "default_split"):
            continue
        if col in ("fold",):
            continue  # Already added

        value = row[col]
        if isinstance(value, str):
            value = _parse_dict_string(value)

        if isinstance(value, (int, float, np.integer, np.floating)):
            if not np.isnan(value) if isinstance(value, (float, np.floating)) else True:
                result[col] = torch.tensor(float(value), dtype=torch.float32)
            else:
                result[col] = value
        elif isinstance(value, dict):
            result[col] = value
        else:
            result[col] = value

    return result

WindowOutOfRangeError

Bases: ValueError

The requested sample window does not fit inside the record.

SplitsNotPublishedError

Bases: RuntimeError

The dataset's splits are deliberately not on the Hub.

Raised instead of a bare 404 for credentialed or restricted sources, whose identifiers ECGBench will not republish. The message carries the command that regenerates the identical split locally.

UnitConversionError

Bases: ValueError

The dataset's samples are not in a physical unit, so units= cannot apply.

Raised for sources whose publisher standardised the waveforms — see DatasetConfig.signal_units. Scaling them by 1000 would produce a number that looks like microvolts and means nothing.

ecg_collate_fn

ecg_collate_fn(batch: list[dict[str, Any]]) -> dict[str, Any]

Custom collate function for ECG dataset batches.

Stacks tensors, keeps dicts and strings as lists.

Parameters:

Name	Type	Description	Default
batch	list[dict[str, Any]]	List of samples from the dataset	required

Returns:

Type	Description
dict[str, Any]	Batched dictionary

Source code in ecgbench/dataset.py

def ecg_collate_fn(batch: list[dict[str, Any]]) -> dict[str, Any]:
    """Custom collate function for ECG dataset batches.

    Stacks tensors, keeps dicts and strings as lists.

    Args:
        batch: List of samples from the dataset

    Returns:
        Batched dictionary
    """
    from torch.utils.data._utils.collate import default_collate

    if not batch:
        return {}

    all_keys = set(batch[0].keys())
    collatable = {}
    non_collatable = {}

    for key in all_keys:
        values = [sample[key] for sample in batch]

        if all(isinstance(v, dict) for v in values):
            non_collatable[key] = values
        elif all(isinstance(v, (str, type(None))) for v in values):
            non_collatable[key] = values
        else:
            collatable[key] = values

    # default_collate expects a list of dicts, not a dict of lists
    if collatable:
        collatable_batch = [
            {k: collatable[k][i] for k in collatable}
            for i in range(len(batch))
        ]
        result = default_collate(collatable_batch)
    else:
        result = {}
    result.update(non_collatable)

    return result