braindecode 0.8__py3-none-any.whl → 1.0.0__py3-none-any.whl

braindecode/datasets/moabb.py

@@ -0,0 +1,209 @@
+"""Dataset objects for some public datasets."""
+
+# Authors: Hubert Banville <hubert.jbanville@gmail.com>
+#          Lukas Gemein <l.gemein@gmail.com>
+#          Simon Brandt <simonbrandt@protonmail.com>
+#          David Sabbagh <dav.sabbagh@gmail.com>
+#          Pierre Guetschel <pierre.guetschel@gmail.com>
+#
+# License: BSD (3-clause)
+
+from __future__ import annotations
+
+import warnings
+from typing import Any
+
+import mne
+import pandas as pd
+
+from braindecode.util import _update_moabb_docstring
+
+from .base import BaseConcatDataset, BaseDataset
+
+
+def _find_dataset_in_moabb(dataset_name, dataset_kwargs=None):
+    # soft dependency on moabb
+    from moabb.datasets.utils import dataset_list
+
+    for dataset in dataset_list:
+        if dataset_name == dataset.__name__:
+            # return an instance of the found dataset class
+            if dataset_kwargs is None:
+                return dataset()
+            else:
+                return dataset(**dataset_kwargs)
+    raise ValueError(f"{dataset_name} not found in moabb datasets")
+
+
+def _fetch_and_unpack_moabb_data(dataset, subject_ids=None, dataset_load_kwargs=None):
+    if dataset_load_kwargs is None:
+        data = dataset.get_data(subject_ids)
+    else:
+        data = dataset.get_data(subjects=subject_ids, **dataset_load_kwargs)
+
+    raws, subject_ids, session_ids, run_ids = [], [], [], []
+    for subj_id, subj_data in data.items():
+        for sess_id, sess_data in subj_data.items():
+            for run_id, raw in sess_data.items():
+                annots = _annotations_from_moabb_stim_channel(raw, dataset)
+                raw.set_annotations(annots)
+                raws.append(raw)
+                subject_ids.append(subj_id)
+                session_ids.append(sess_id)
+                run_ids.append(run_id)
+    description = pd.DataFrame(
+        {"subject": subject_ids, "session": session_ids, "run": run_ids}
+    )
+    return raws, description
+
+
+def _annotations_from_moabb_stim_channel(raw, dataset):
+    # find events from the stim channel
+    stim_channels = mne.utils._get_stim_channel(None, raw.info, raise_error=False)
+    if len(stim_channels) > 0:
+        # returns an empty array if none found
+        events = mne.find_events(raw, shortest_event=0, verbose=False)
+        event_id = dataset.event_id
+    else:
+        events, event_id = mne.events_from_annotations(raw, verbose=False)
+
+    # get annotations from events
+    event_desc = {k: v for v, k in event_id.items()}
+    annots = mne.annotations_from_events(events, raw.info["sfreq"], event_desc)
+
+    # set trial on and offset given by moabb
+    onset, offset = dataset.interval
+    annots.onset += onset
+    annots.duration += offset - onset
+    return annots
+
+
+def fetch_data_with_moabb(
+    dataset_name: str,
+    subject_ids: list[int] | int | None = None,
+    dataset_kwargs: dict[str, Any] | None = None,
+    dataset_load_kwargs: dict[str, Any] | None = None,
+) -> tuple[list[mne.io.Raw], pd.DataFrame]:
+    # ToDo: update path to where moabb downloads / looks for the data
+    """Fetch data using moabb.
+
+    Parameters
+    ----------
+    dataset_name: str | moabb.datasets.base.BaseDataset
+        the name of a dataset included in moabb
+    subject_ids: list(int) | int
+        (list of) int of subject(s) to be fetched
+    dataset_kwargs: dict, optional
+        optional dictionary containing keyword arguments
+        to pass to the moabb dataset when instantiating it.
+    data_load_kwargs: dict, optional
+        optional dictionary containing keyword arguments
+        to pass to the moabb dataset's load_data method.
+        Allows using the moabb cache_config=None and
+        process_pipeline=None.
+
+    Returns
+    -------
+    raws: mne.Raw
+    info: pandas.DataFrame
+    """
+    if isinstance(dataset_name, str):
+        dataset = _find_dataset_in_moabb(dataset_name, dataset_kwargs)
+    else:
+        from moabb.datasets.base import BaseDataset
+
+        if isinstance(dataset_name, BaseDataset):
+            dataset = dataset_name
+
+    subject_id = [subject_ids] if isinstance(subject_ids, int) else subject_ids
+    return _fetch_and_unpack_moabb_data(
+        dataset, subject_id, dataset_load_kwargs=dataset_load_kwargs
+    )
+
+
+class MOABBDataset(BaseConcatDataset):
+    """A class for moabb datasets.
+
+    Parameters
+    ----------
+    dataset_name: str
+        name of dataset included in moabb to be fetched
+    subject_ids: list(int) | int | None
+        (list of) int of subject(s) to be fetched. If None, data of all
+        subjects is fetched.
+    dataset_kwargs: dict, optional
+        optional dictionary containing keyword arguments
+        to pass to the moabb dataset when instantiating it.
+    dataset_load_kwargs: dict, optional
+        optional dictionary containing keyword arguments
+        to pass to the moabb dataset's load_data method.
+        Allows using the moabb cache_config=None and
+        process_pipeline=None.
+    """
+
+    def __init__(
+        self,
+        dataset_name: str,
+        subject_ids: list[int] | int | None = None,
+        dataset_kwargs: dict[str, Any] | None = None,
+        dataset_load_kwargs: dict[str, Any] | None = None,
+    ):
+        # soft dependency on moabb
+        from moabb import __version__ as moabb_version  # type: ignore
+
+        if moabb_version == "1.0.0":
+            warnings.warn(
+                "moabb version 1.0.0 generates incorrect annotations. "
+                "Please update to another version, version 0.5 or 1.1.0 "
+            )
+
+        raws, description = fetch_data_with_moabb(
+            dataset_name,
+            subject_ids,
+            dataset_kwargs,
+            dataset_load_kwargs=dataset_load_kwargs,
+        )
+        all_base_ds = [
+            BaseDataset(raw, row) for raw, (_, row) in zip(raws, description.iterrows())
+        ]
+        super().__init__(all_base_ds)
+
+
+class BNCI2014001(MOABBDataset):
+    doc = """See moabb.datasets.bnci.BNCI2014001
+
+    Parameters
+    ----------
+    subject_ids: list(int) | int | None
+        (list of) int of subject(s) to be fetched. If None, data of all
+        subjects is fetched.
+    """
+    try:
+        from moabb.datasets import BNCI2014001
+
+        __doc__ = _update_moabb_docstring(BNCI2014001, doc)
+    except ModuleNotFoundError:
+        pass  # keep moabb soft dependency, otherwise crash on loading of datasets.__init__.py
+
+    def __init__(self, subject_ids):
+        super().__init__("BNCI2014001", subject_ids=subject_ids)
+
+
+class HGD(MOABBDataset):
+    doc = """See moabb.datasets.schirrmeister2017.Schirrmeister2017
+
+    Parameters
+    ----------
+    subject_ids: list(int) | int | None
+        (list of) int of subject(s) to be fetched. If None, data of all
+        subjects is fetched.
+    """
+    try:
+        from moabb.datasets import Schirrmeister2017
+
+        __doc__ = _update_moabb_docstring(Schirrmeister2017, doc)
+    except ModuleNotFoundError:
+        pass  # keep moabb soft dependency, otherwise crash on loading of datasets.__init__.py
+
+    def __init__(self, subject_ids):
+        super().__init__("Schirrmeister2017", subject_ids=subject_ids)

braindecode/datasets/nmt.py

@@ -0,0 +1,311 @@
+"""
+Dataset classes for the NMT EEG Corpus dataset.
+
+The NMT Scalp EEG Dataset is an open-source annotated dataset of healthy and
+pathological EEG recordings for predictive modeling. This dataset contains
+2,417 recordings from unique participants spanning almost 625 h.
+
+Note:
+    - The signal unit may not be uV and further examination is required.
+    - The spectrum shows that the signal may have been band-pass filtered from about 2 - 33Hz,
+    which needs to be further determined.
+
+"""
+
+# Authors: Mohammad Bayazi <mj.darvishi92@gmail.com>
+#          Bruno Aristimunha <b.aristimunha@gmail.com>
+#
+# License: BSD (3-clause)
+
+from __future__ import annotations
+
+import glob
+import os
+import warnings
+from pathlib import Path
+from unittest import mock
+
+import mne
+import numpy as np
+import pandas as pd
+from joblib import Parallel, delayed
+from mne.datasets import fetch_dataset
+
+from braindecode.datasets.base import BaseConcatDataset, BaseDataset
+
+NMT_URL = "https://zenodo.org/record/10909103/files/NMT.zip"
+NMT_archive_name = "NMT.zip"
+NMT_folder_name = "MNE-NMT-eeg-dataset"
+NMT_dataset_name = "NMT-EEG-Corpus"
+
+NMT_dataset_params = {
+    "dataset_name": NMT_dataset_name,
+    "url": NMT_URL,
+    "archive_name": NMT_archive_name,
+    "folder_name": NMT_folder_name,
+    "hash": "77b3ce12bcaf6c6cce4e6690ea89cb22bed55af10c525077b430f6e1d2e3c6bf",
+    "config_key": NMT_dataset_name,
+}
+
+
+class NMT(BaseConcatDataset):
+    """The NMT Scalp EEG Dataset.
+
+    An Open-Source Annotated Dataset of Healthy and Pathological EEG
+    Recordings for Predictive Modeling.
+
+    This dataset contains 2,417 recordings from unique participants spanning
+    almost 625 h.
+
+    Here, the dataset can be used for three tasks, brain-age, gender prediction,
+    abnormality detection.
+
+    The dataset is described in [Khan2022]_.
+
+    .. versionadded:: 0.9
+
+    Parameters
+    ----------
+    path: str
+        Parent directory of the dataset.
+    recording_ids: list(int) | int
+        A (list of) int of recording id(s) to be read (order matters and will
+        overwrite default chronological order, e.g. if recording_ids=[1,0],
+        then the first recording returned by this class will be chronologically
+        later than the second recording. Provide recording_ids in ascending
+        order to preserve chronological order.).
+    target_name: str
+        Can be "pathological", "gender", or "age".
+    preload: bool
+        If True, preload the data of the Raw objects.
+
+    References
+    ----------
+    .. [Khan2022] Khan, H.A.,Ul Ain, R., Kamboh, A.M., Butt, H.T.,Shafait,S.,
+        Alamgir, W., Stricker, D. and Shafait, F., 2022. The NMT scalp EEG
+        dataset: an open-source annotated dataset of healthy and pathological
+        EEG recordings for predictive modeling. Frontiers in neuroscience,
+        15, p.755817.
+    """
+
+    def __init__(
+        self,
+        path=None,
+        target_name="pathological",
+        recording_ids=None,
+        preload=False,
+        n_jobs=1,
+    ):
+        # correct the path if needed
+        if path is not None:
+            list_csv = glob.glob(f"{path}/**/Labels.csv", recursive=True)
+            if isinstance(list_csv, list) and len(list_csv) > 0:
+                path = Path(list_csv[0]).parent
+
+        if path is None or len(list_csv) == 0:
+            path = fetch_dataset(
+                dataset_params=NMT_dataset_params,
+                path=Path(path) if path is not None else None,
+                processor="unzip",
+                force_update=False,
+            )
+            # First time we fetch the dataset, we need to move the files to the
+            # correct directory.
+            path = _correct_path(path)
+
+        # Get all file paths
+        file_paths = glob.glob(
+            os.path.join(path, "**" + os.sep + "*.edf"), recursive=True
+        )
+
+        # sort by subject id
+        file_paths = [
+            file_path
+            for file_path in file_paths
+            if os.path.splitext(file_path)[1] == ".edf"
+        ]
+
+        # sort by subject id
+        file_paths = sorted(
+            file_paths, key=lambda p: int(os.path.splitext(p)[0].split(os.sep)[-1])
+        )
+        if recording_ids is not None:
+            file_paths = [file_paths[rec_id] for rec_id in recording_ids]
+
+        # read labels and rearrange them to match TUH Abnormal EEG Corpus
+        description = pd.read_csv(
+            os.path.join(path, "Labels.csv"), index_col="recordname"
+        )
+        if recording_ids is not None:
+            description = description.iloc[recording_ids]
+        description.replace(
+            {
+                "not specified": "X",
+                "female": "F",
+                "male": "M",
+                "abnormal": True,
+                "normal": False,
+            },
+            inplace=True,
+        )
+        description.rename(columns={"label": "pathological"}, inplace=True)
+        description.reset_index(drop=True, inplace=True)
+        description["path"] = file_paths
+        description = description[["path", "pathological", "age", "gender"]]
+
+        if n_jobs == 1:
+            base_datasets = [
+                self._create_dataset(d, target_name, preload)
+                for recording_id, d in description.iterrows()
+            ]
+        else:
+            base_datasets = Parallel(n_jobs)(
+                delayed(self._create_dataset)(d, target_name, preload)
+                for recording_id, d in description.iterrows()
+            )
+
+        super().__init__(base_datasets)
+
+    @staticmethod
+    def _create_dataset(d, target_name, preload):
+        raw = mne.io.read_raw_edf(d.path, preload=preload)
+        d["n_samples"] = raw.n_times
+        d["sfreq"] = raw.info["sfreq"]
+        d["train"] = "train" in d.path.split(os.sep)
+        base_dataset = BaseDataset(raw, d, target_name)
+        return base_dataset
+
+
+def _correct_path(path: str):
+    """
+    Check if the path is correct and rename the file if needed.
+
+    Parameters
+    ----------
+    path: basestring
+        Path to the file.
+
+    Returns
+    -------
+    path: basestring
+        Corrected path.
+    """
+    if not Path(path).exists():
+        unzip_file_name = f"{NMT_archive_name}.unzip"
+        if (Path(path).parent / unzip_file_name).exists():
+            try:
+                os.rename(
+                    src=Path(path).parent / unzip_file_name,
+                    dst=Path(path),
+                )
+
+            except PermissionError:
+                raise PermissionError(
+                    f"Please rename {Path(path).parent / unzip_file_name}"
+                    + f"manually to {path} and try again."
+                )
+        path = os.path.join(path, "nmt_scalp_eeg_dataset")
+
+    return path
+
+
+def _get_header(*args):
+    all_paths = {**_NMT_PATHS}
+    return all_paths[args[0]]
+
+
+def _fake_pd_read_csv(*args, **kwargs):
+    # Create a list of lists to hold the data
+    data = [
+        ["0000001.edf", "normal", 35, "male", "train"],
+        ["0000002.edf", "abnormal", 28, "female", "test"],
+        ["0000003.edf", "normal", 62, "male", "train"],
+        ["0000004.edf", "abnormal", 41, "female", "test"],
+        ["0000005.edf", "normal", 19, "male", "train"],
+        ["0000006.edf", "abnormal", 55, "female", "test"],
+        ["0000007.edf", "normal", 71, "male", "train"],
+    ]
+
+    # Create the DataFrame, specifying column names
+    df = pd.DataFrame(data, columns=["recordname", "label", "age", "gender", "loc"])
+
+    return df
+
+
+def _fake_raw(*args, **kwargs):
+    sfreq = 10
+    ch_names = [
+        "EEG A1-REF",
+        "EEG A2-REF",
+        "EEG FP1-REF",
+        "EEG FP2-REF",
+        "EEG F3-REF",
+        "EEG F4-REF",
+        "EEG C3-REF",
+        "EEG C4-REF",
+        "EEG P3-REF",
+        "EEG P4-REF",
+        "EEG O1-REF",
+        "EEG O2-REF",
+        "EEG F7-REF",
+        "EEG F8-REF",
+        "EEG T3-REF",
+        "EEG T4-REF",
+        "EEG T5-REF",
+        "EEG T6-REF",
+        "EEG FZ-REF",
+        "EEG CZ-REF",
+        "EEG PZ-REF",
+    ]
+    duration_min = 6
+    data = np.random.randn(len(ch_names), duration_min * sfreq * 60)
+    info = mne.create_info(ch_names=ch_names, sfreq=sfreq, ch_types="eeg")
+    raw = mne.io.RawArray(data=data, info=info)
+    return raw
+
+
+_NMT_PATHS = {
+    # these are actual file paths and edf headers from NMT EEG Corpus
+    "nmt_scalp_eeg_dataset/abnormal/train/0000036.edf": b"0       0000036  M 13-May-1951 0000036  Age:32                                          ",
+    # noqa E501
+    "nmt_scalp_eeg_dataset/abnormal/eval/0000037.edf": b"0       0000037  M 13-May-1951 0000037  Age:32                                          ",
+    # noqa E501
+    "nmt_scalp_eeg_dataset/abnormal/eval/0000038.edf": b"0       0000038  M 13-May-1951 0000038  Age:32                                          ",
+    # noqa E501
+    "nmt_scalp_eeg_dataset/normal/train/0000039.edf": b"0       0000039  M 13-May-1951 0000039  Age:32                                          ",
+    # noqa E501
+    "nmt_scalp_eeg_dataset/normal/eval/0000040.edf": b"0       0000040  M 13-May-1951 0000040  Age:32                                          ",
+    # noqa E501
+    "nmt_scalp_eeg_dataset/normal/eval/0000041.edf": b"0       0000041  M 13-May-1951 0000041  Age:32                                          ",
+    # noqa E501
+    "nmt_scalp_eeg_dataset/abnormal/train/0000042.edf": b"0       0000042  M 13-May-1951 0000042  Age:32                                          ",
+    # noqa E501
+    "Labels.csv": b"0       recordname,label,age,gender,loc       1 0000001.edf,normal,22,not specified,train                                                                      ",
+    # noqa E501
+}
+
+
+class _NMTMock(NMT):
+    """Mocked class for testing and examples."""
+
+    @mock.patch("glob.glob", return_value=_NMT_PATHS.keys())
+    @mock.patch("mne.io.read_raw_edf", new=_fake_raw)
+    @mock.patch("pandas.read_csv", new=_fake_pd_read_csv)
+    def __init__(
+        self,
+        mock_glob,
+        path,
+        recording_ids=None,
+        target_name="pathological",
+        preload=False,
+        n_jobs=1,
+    ):
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", message="Cannot save date file")
+            super().__init__(
+                path=path,
+                recording_ids=recording_ids,
+                target_name=target_name,
+                preload=preload,
+                n_jobs=n_jobs,
+            )