braindecode 0.8.1__py3-none-any.whl → 1.1.0__py3-none-any.whl

braindecode/datasets/mne.py

@@ -3,18 +3,28 @@
 #
 # License: BSD (3-clause)
 
+from __future__ import annotations
+
+import mne
 import numpy as np
 import pandas as pd
-import mne
 
-from .base import BaseDataset, BaseConcatDataset, WindowsDataset
+from .base import BaseConcatDataset, BaseDataset, WindowsDataset
 
 
 def create_from_mne_raw(
-        raws, trial_start_offset_samples, trial_stop_offset_samples,
-        window_size_samples, window_stride_samples, drop_last_window,
-        descriptions=None, mapping=None, preload=False, drop_bad_windows=True,
-        accepted_bads_ratio=0.0):
+    raws: list[mne.io.BaseRaw],
+    trial_start_offset_samples: int,
+    trial_stop_offset_samples: int,
+    window_size_samples: int,
+    window_stride_samples: int,
+    drop_last_window: bool,
+    descriptions: list[dict | pd.Series] | None = None,
+    mapping: dict[str, int] | None = None,
+    preload: bool = False,
+    drop_bad_windows: bool = True,
+    accepted_bads_ratio: float = 0.0,
+) -> BaseConcatDataset:
     """Create WindowsDatasets from mne.RawArrays
 
     Parameters
@@ -58,13 +68,16 @@ def create_from_mne_raw(
     """
     # Prevent circular import
     from ..preprocessing.windowers import create_windows_from_events
+
     if descriptions is not None:
         if len(descriptions) != len(raws):
             raise ValueError(
                 f"length of 'raws' ({len(raws)}) and 'description' "
-                f"({len(descriptions)}) has to match")
-        base_datasets = [BaseDataset(raw, desc) for raw, desc in
-                         zip(raws, descriptions)]
+                f"({len(descriptions)}) has to match"
+            )
+        base_datasets = [
+            BaseDataset(raw, desc) for raw, desc in zip(raws, descriptions)
+        ]
     else:
         base_datasets = [BaseDataset(raw) for raw in raws]
 
@@ -84,8 +97,12 @@ def create_from_mne_raw(
     return windows_datasets
 
 
-def create_from_mne_epochs(list_of_epochs, window_size_samples,
-                           window_stride_samples, drop_last_window):
+def create_from_mne_epochs(
+    list_of_epochs: list[mne.BaseEpochs],
+    window_size_samples: int,
+    window_stride_samples: int,
+    drop_last_window: bool,
+) -> BaseConcatDataset:
     """Create WindowsDatasets from mne.Epochs
 
     Parameters
@@ -108,8 +125,8 @@ def create_from_mne_epochs(list_of_epochs, window_size_samples,
     """
     # Prevent circular import
     from ..preprocessing.windowers import _check_windowing_arguments
-    _check_windowing_arguments(0, 0, window_size_samples,
-                               window_stride_samples)
+
+    _check_windowing_arguments(0, 0, window_size_samples, window_stride_samples)
 
     list_of_windows_ds = []
     for epochs in list_of_epochs:
@@ -124,24 +141,28 @@ def create_from_mne_epochs(list_of_epochs, window_size_samples,
             # if last window does not end at trial stop, make it stop there
             starts = np.append(starts, stop)
 
-        fake_events = [[start, window_size_samples, -1] for start in
-                       starts]
+        fake_events = [[start, window_size_samples, -1] for start in starts]
 
         for trial_i, trial in enumerate(epochs):
-            metadata = pd.DataFrame({
-                'i_window_in_trial': np.arange(len(fake_events)),
-                'i_start_in_trial': starts + original_trial_starts[trial_i],
-                'i_stop_in_trial': starts + original_trial_starts[
-                    trial_i] + window_size_samples,
-                'target': len(fake_events) * [event_descriptions[trial_i]]
-            })
+            metadata = pd.DataFrame(
+                {
+                    "i_window_in_trial": np.arange(len(fake_events)),
+                    "i_start_in_trial": starts + original_trial_starts[trial_i],
+                    "i_stop_in_trial": starts
+                    + original_trial_starts[trial_i]
+                    + window_size_samples,
+                    "target": len(fake_events) * [event_descriptions[trial_i]],
+                }
+            )
             # window size - 1, since tmax is inclusive
             mne_epochs = mne.Epochs(
-                mne.io.RawArray(trial, epochs.info), fake_events,
+                mne.io.RawArray(trial, epochs.info),
+                fake_events,
                 baseline=None,
                 tmin=0,
                 tmax=(window_size_samples - 1) / epochs.info["sfreq"],
-                metadata=metadata)
+                metadata=metadata,
+            )
 
             mne_epochs.drop_bad(reject=None, flat=None)
 

braindecode/datasets/moabb.py

@@ -1,5 +1,4 @@
-"""Dataset objects for some public datasets.
-"""
+"""Dataset objects for some public datasets."""
 
 # Authors: Hubert Banville <hubert.jbanville@gmail.com>
 #          Lukas Gemein <l.gemein@gmail.com>
@@ -9,16 +8,23 @@
 #
 # License: BSD (3-clause)
 
-import pandas as pd
+from __future__ import annotations
+
+import warnings
+from typing import Any
+
 import mne
+import pandas as pd
 
-from .base import BaseDataset, BaseConcatDataset
 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
@@ -29,35 +35,41 @@ def _find_dataset_in_moabb(dataset_name, dataset_kwargs=None):
     raise ValueError(f"{dataset_name} not found in moabb datasets")
 
 
-def _fetch_and_unpack_moabb_data(dataset, subject_ids):
-    data = dataset.get_data(subject_ids)
+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():
-                # set annotation if empty
-                if len(raw.annotations) == 0:
-                    annots = _annotations_from_moabb_stim_channel(raw, dataset)
-                    raw.set_annotations(annots)
+                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
-    })
+    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 stim channel
-    events = mne.find_events(raw)
+    # 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 dataset.event_id.items()}
-    annots = mne.annotations_from_events(events, raw.info['sfreq'], event_desc)
+    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
@@ -66,28 +78,47 @@ def _annotations_from_moabb_stim_channel(raw, dataset):
     return annots
 
 
-def fetch_data_with_moabb(dataset_name, subject_ids, dataset_kwargs=None):
+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
+    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
     """
-    dataset = _find_dataset_in_moabb(dataset_name, dataset_kwargs)
+    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)
+    return _fetch_and_unpack_moabb_data(
+        dataset, subject_id, dataset_load_kwargs=dataset_load_kwargs
+    )
 
 
 class MOABBDataset(BaseConcatDataset):
@@ -103,11 +134,38 @@ class MOABBDataset(BaseConcatDataset):
     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, subject_ids, dataset_kwargs=None):
-        raws, description = fetch_data_with_moabb(dataset_name, subject_ids, dataset_kwargs)
-        all_base_ds = [BaseDataset(raw, row)
-                       for raw, (_, row) in zip(raws, description.iterrows())]
+
+    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)
 
 
@@ -122,6 +180,7 @@ class BNCI2014001(MOABBDataset):
     """
     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
@@ -141,6 +200,7 @@ class HGD(MOABBDataset):
     """
     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

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