braindecode 0.5.1__py3-none-any.whl

braindecode/datasets/moabb.py

@@ -0,0 +1,110 @@
+"""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>
+#
+# License: BSD (3-clause)
+
+import pandas as pd
+import mne
+
+from .base import BaseDataset, BaseConcatDataset, WindowsDataset
+
+
+def _find_dataset_in_moabb(dataset_name):
+    # 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
+            return dataset()
+    raise ValueError("'dataset_name' not found in moabb datasets")
+
+
+def _fetch_and_unpack_moabb_data(dataset, subject_ids):
+    data = dataset.get_data(subject_ids)
+    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)
+                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 stim channel
+    events = mne.find_events(raw)
+
+    # 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)
+
+    # 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, subject_ids):
+    # ToDo: update path to where moabb downloads / looks for the data
+    """Fetch data using moabb.
+
+    Parameters
+    ----------
+    dataset_name: str
+        the name of a dataset included in moabb
+    subject_ids: list(int) | int
+        (list of) int of subject(s) to be fetched
+
+    Returns
+    -------
+    raws: mne.Raw
+    info: pandas.DataFrame
+    """
+    dataset = _find_dataset_in_moabb(dataset_name)
+    subject_id = [subject_ids] if isinstance(subject_ids, int) else subject_ids
+    return _fetch_and_unpack_moabb_data(dataset, subject_id)
+
+
+class MOABBDataset(BaseConcatDataset):
+    """A class for moabb datasets.
+
+    Parameters
+    ----------
+    dataset_name: name of dataset included in moabb to be fetched
+    subject_ids: list(int) | int
+        (list of) int of subject(s) to be fetched
+    """
+    def __init__(self, dataset_name, subject_ids):
+        raws, description = fetch_data_with_moabb(dataset_name, subject_ids)
+        all_base_ds = [BaseDataset(raw, row)
+                       for raw, (_, row) in zip(raws, description.iterrows())]
+        super().__init__(all_base_ds)
+
+
+class BNCI2014001(MOABBDataset):
+    """See moabb.datasets.bnci.BNCI2014001"""
+    def __init__(self, *args, **kwargs):
+        super().__init__("BNCI2014001", *args, **kwargs)
+
+
+class HGD(MOABBDataset):
+    """See moabb.datasets.schirrmeister2017.Schirrmeister2017"""
+    def __init__(self, *args, **kwargs):
+        super().__init__("Schirrmeister2017", *args, **kwargs)

braindecode/datasets/sleep_physionet.py

@@ -0,0 +1,103 @@
+# Authors: Hubert Banville <hubert.jbanville@gmail.com>
+#
+# License: BSD (3-clause)
+
+
+import os
+
+import numpy as np
+import pandas as pd
+import mne
+from mne.datasets.sleep_physionet.age import fetch_data
+
+from .base import BaseDataset, BaseConcatDataset
+
+
+class SleepPhysionet(BaseConcatDataset):
+    """Sleep Physionet dataset.
+
+    Sleep dataset from https://physionet.org/content/sleep-edfx/1.0.0/.
+    Contains overnight recordings from 78 healthy subjects.
+
+    See [MNE example](https://mne.tools/stable/auto_tutorials/sample-datasets/plot_sleep.html).
+
+    Parameters
+    ----------
+    subject_ids: list(int) | int | None
+        (list of) int of subject(s) to be loaded. If None, load all available
+        subjects.
+    recording_ids: list(int) | None
+        Recordings to load per subject (each subject except 13 has two
+        recordings). Can be [1], [2] or [1, 2] (same as None).
+    preload: bool
+        If True, preload the data of the Raw objects.
+    load_eeg_only: bool
+        If True, only load the EEG channels and discard the others (EOG, EMG,
+        temperature, respiration) to avoid resampling the other signals.
+    crop_wake_mins: float
+        Number of minutes of wake time to keep before the first sleep event
+        and after the last sleep event. Used to reduce the imbalance in this
+        dataset. Default of 30 mins.
+    """
+    def __init__(self, subject_ids=None, recording_ids=None, preload=False,
+                 load_eeg_only=True, crop_wake_mins=30):
+        if subject_ids is None:
+            subject_ids = range(83)
+        if recording_ids is None:
+            recording_ids = [1, 2]
+
+        paths = fetch_data(
+            subject_ids, recording=recording_ids, on_missing='warn')
+
+        all_base_ds = list()
+        for p in paths:
+            raw, desc = self._load_raw(
+                p[0], p[1], preload=preload, load_eeg_only=load_eeg_only,
+                crop_wake_mins=crop_wake_mins)
+            base_ds = BaseDataset(raw, desc)
+            all_base_ds.append(base_ds)
+        super().__init__(all_base_ds)
+
+    @staticmethod
+    def _load_raw(raw_fname, ann_fname, preload, load_eeg_only=True,
+                  crop_wake_mins=False):
+        ch_mapping = {
+            'EOG horizontal': 'eog',
+            'Resp oro-nasal': 'misc',
+            'EMG submental': 'misc',
+            'Temp rectal': 'misc',
+            'Event marker': 'misc'
+        }
+        exclude = ch_mapping.keys() if load_eeg_only else ()
+
+        raw = mne.io.read_raw_edf(raw_fname, preload=preload, exclude=exclude)
+        annots = mne.read_annotations(ann_fname)
+        raw.set_annotations(annots, emit_warning=False)
+
+        if crop_wake_mins > 0:
+            # Find first and last sleep stages
+            mask = [
+                x[-1] in ['1', '2', '3', '4', 'R'] for x in annots.description]
+            sleep_event_inds = np.where(mask)[0]
+
+            # Crop raw
+            tmin = annots[int(sleep_event_inds[0])]['onset'] - \
+                   crop_wake_mins * 60
+            tmax = annots[int(sleep_event_inds[-1])]['onset'] + \
+                   crop_wake_mins * 60
+            raw.crop(tmin=tmin, tmax=tmax)
+
+        # Rename EEG channels
+        ch_names = {
+            i: i.replace('EEG ', '') for i in raw.ch_names if 'EEG' in i}
+        mne.rename_channels(raw.info, ch_names)
+
+        if not load_eeg_only:
+            raw.set_channel_types(ch_mapping)
+
+        basename = os.path.basename(raw_fname)
+        subj_nb = int(basename[3:5])
+        sess_nb = int(basename[5])
+        desc = pd.Series({'subject': subj_nb, 'recording': sess_nb}, name='')
+
+        return raw, desc

braindecode/datasets/tuh.py

@@ -0,0 +1,230 @@
+"""
+Dataset classes for the Temple University Hospital (TUH) EEG Corpus and the
+TUH Abnormal EEG Corpus.
+"""
+
+# Authors: Lukas Gemein <l.gemein@gmail.com>
+#
+# License: BSD (3-clause)
+
+import re
+import os
+import glob
+
+import pandas as pd
+import mne
+
+from .base import BaseDataset, BaseConcatDataset
+
+
+class TUH(BaseConcatDataset):
+    """Temple University Hospital (TUH) EEG Corpus
+    (www.isip.piconepress.com/projects/tuh_eeg/html/downloads.shtml#c_tueg).
+
+    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 then the second recording. Provide recording_ids in ascending
+        order to preserve chronological order.).
+    target_name: str
+        Can be 'gender', or 'age'.
+    preload: bool
+        If True, preload the data of the Raw objects.
+    add_physician_reports: bool
+        If True, the physician reports will be read from disk and added to the
+        description.
+    """
+    def __init__(self, path, recording_ids=None, target_name=None,
+                 preload=False, add_physician_reports=False):
+        # create an index of all files and gather easily accessible info
+        # without actually touching the files
+        file_paths = glob.glob(os.path.join(path, '**/*.edf'), recursive=True)
+        descriptions = _create_chronological_description(file_paths)
+        # limit to specified recording ids before doing slow stuff
+        if recording_ids is not None:
+            descriptions = descriptions[recording_ids]
+        # create datasets gathering more info about the files touching them
+        # reading the raws and potentially preloading the data
+        base_datasets = self._create_datasets(
+            descriptions, target_name, preload, add_physician_reports)
+        super().__init__(base_datasets)
+
+    @staticmethod
+    def _create_datasets(descriptions, target_name, preload,
+                         add_physician_reports):
+        # this is the second loop (slow)
+        base_datasets = []
+        for file_path_i, file_path in descriptions.loc['path'].iteritems():
+            # parse age and gender information from EDF header
+            age, gender = _parse_age_and_gender_from_edf_header(file_path)
+            raw = mne.io.read_raw_edf(file_path, preload=preload)
+            # read info relevant for preprocessing from raw without loading it
+            sfreq = raw.info['sfreq']
+            n_samples = raw.n_times
+            d = {
+                'sfreq': float(sfreq),
+                'n_samples': int(n_samples),
+                'age': int(age),
+                'gender': gender,
+            }
+            if add_physician_reports:
+                physician_report = _read_physician_report(file_path)
+                d['report'] = physician_report
+            additional_description = pd.Series(d)
+            description = pd.concat(
+                [descriptions.pop(file_path_i), additional_description])
+            base_dataset = BaseDataset(raw, description,
+                                       target_name=target_name)
+            base_datasets.append(base_dataset)
+        return base_datasets
+
+
+def _create_chronological_description(file_paths):
+    # this is the first loop (fast)
+    descriptions = []
+    for file_path in file_paths:
+        description = _parse_description_from_file_path(file_path)
+        descriptions.append(pd.Series(description))
+    descriptions = pd.concat(descriptions, axis=1)
+    # order descriptions chronologically
+    descriptions.sort_values(
+        ["year", "month", "day", "subject", "session", "segment"],
+        axis=1, inplace=True)
+    # https://stackoverflow.com/questions/42284617/reset-column-index-pandas
+    descriptions = descriptions.T.reset_index(drop=True).T
+    return descriptions
+
+
+def _parse_description_from_file_path(file_path):
+    # stackoverflow.com/questions/3167154/how-to-split-a-dos-path-into-its-components-in-python  # noqa
+    file_path = os.path.normpath(file_path)
+    tokens = file_path.split(os.sep)
+    # expect file paths as tuh_eeg/version/file_type/reference/data_split/
+    #                          subject/recording session/file
+    # e.g.                 tuh_eeg/v1.1.0/edf/01_tcp_ar/027/00002729/
+    #                          s001_2006_04_12/00002729_s001.edf
+    version = tokens[-7]
+    year, month, day = tokens[-2].split('_')[1:]
+    subject_id = tokens[-3]
+    session = tokens[-2].split('_')[0]
+    segment = tokens[-1].split('_')[-1].split('.')[-2]
+    reference = tokens[-5].split('_')[2]
+    return {
+        'path': file_path,
+        'version': version,
+        'year': int(year),
+        'month': int(month),
+        'day': int(day),
+        'subject': int(subject_id),
+        'session': int(session[1:]),
+        'segment': int(segment[1:]),
+        'reference': reference,
+    }
+
+
+def _read_physician_report(file_path):
+    directory = os.path.dirname(file_path)
+    txt_file = glob.glob(os.path.join(directory, '**/*.txt'), recursive=True)
+    # check that there is at most one txt file in the same directory
+    assert len(txt_file) in [0, 1]
+    report = ''
+    if txt_file:
+        txt_file = txt_file[0]
+        # somewhere in the corpus, encoding apparently changed
+        # first try to read as utf-8, if it does not work use latin-1
+        try:
+            with open(txt_file, 'r', encoding='utf-8') as f:
+                report = f.read()
+        except UnicodeDecodeError:
+            with open(txt_file, 'r', encoding='latin-1') as f:
+                report = f.read()
+    return report
+
+
+def _read_edf_header(file_path):
+    f = open(file_path, "rb")
+    header = f.read(88)
+    f.close()
+    return header
+
+
+def _parse_age_and_gender_from_edf_header(file_path):
+    header = _read_edf_header(file_path)
+    # bytes 8 to 88 contain ascii local patient identification
+    # see https://www.teuniz.net/edfbrowser/edf%20format%20description.html
+    patient_id = header[8:].decode("ascii")
+    age = -1
+    found_age = re.findall(r"Age:(\d+)", patient_id)
+    if len(found_age) == 1:
+        age = int(found_age[0])
+    gender = "X"
+    found_gender = re.findall(r"\s([F|M])\s", patient_id)
+    if len(found_gender) == 1:
+        gender = found_gender[0]
+    return age, gender
+
+
+class TUHAbnormal(TUH):
+    """Temple University Hospital (TUH) Abnormal EEG Corpus.
+    see www.isip.piconepress.com/projects/tuh_eeg/html/downloads.shtml#c_tuab
+
+    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 then 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.
+    add_physician_reports: bool
+        If True, the physician reports will be read from disk and added to the
+        description.
+    """
+    def __init__(self, path, recording_ids=None, target_name='pathological',
+                 preload=False, add_physician_reports=False):
+        super().__init__(path=path, recording_ids=recording_ids,
+                         preload=preload,
+                         add_physician_reports=add_physician_reports)
+        additional_descriptions = []
+        for file_path in self.description.path:
+            additional_description = (
+                self._parse_additional_description_from_file_path(file_path))
+            additional_descriptions.append(additional_description)
+        additional_descriptions = pd.DataFrame(additional_descriptions)
+        self.description = pd.concat(
+            [self.description, additional_descriptions], axis=1)
+        # not 100% sure if this is required:
+        # set target name and target of base datasets
+        for ds_i, ds in enumerate(self.datasets):
+            ds.target_name = target_name
+            ds.target = self.description.loc[ds_i, target_name]
+            for k in additional_descriptions.columns:
+                ds.description[k] = self.description.loc[ds_i, k]
+
+    @staticmethod
+    def _parse_additional_description_from_file_path(file_path):
+        file_path = os.path.normpath(file_path)
+        tokens = file_path.split(os.sep)
+        # expect paths as version/file type/data_split/pathology status/
+        #                     reference/subset/subject/recording session/file
+        # e.g.            v2.0.0/edf/train/normal/01_tcp_ar/000/00000021/
+        #                     s004_2013_08_15/00000021_s004_t000.edf
+        assert ('abnormal' in tokens or 'normal' in tokens), (
+            'No pathology labels found.')
+        assert ('train' in tokens or 'eval' in tokens), (
+            'No train or eval set information found.')
+        return {
+            'train': 'train' in tokens,
+            'pathological': 'abnormal' in tokens,
+        }

braindecode/datautil/__init__.py

@@ -0,0 +1,11 @@
+"""
+Utilities for data manipulation.
+"""
+
+from .windowers import create_windows_from_events, create_fixed_length_windows
+from .preprocess import (zscore, scale, exponential_moving_demean,
+                         exponential_moving_standardize, filterbank,
+                         preprocess, Preprocessor)
+from .xy import create_from_X_y
+from .mne import create_from_mne_raw, create_from_mne_epochs
+from .serialization import save_concat_dataset, load_concat_dataset

braindecode/datautil/mne.py

@@ -0,0 +1,136 @@
+import numpy as np
+import pandas as pd
+import mne
+
+from ..datasets.base import BaseDataset, BaseConcatDataset, WindowsDataset
+from ..datautil.windowers import (
+    _check_windowing_arguments, create_windows_from_events)
+
+
+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):
+    """Create WindowsDatasets from mne.RawArrays
+
+    Parameters
+    ----------
+    raws: array-like
+        list of mne.RawArrays
+    trial_start_offset_samples: int
+        start offset from original trial onsets in samples
+    trial_stop_offset_samples: int
+        stop offset from original trial stop in samples
+    window_size_samples: int
+        window size
+    window_stride_samples: int
+        stride between windows
+    drop_last_window: bool
+        whether or not have a last overlapping window, when
+        windows do not equally divide the continuous signal
+    descriptions: array-like
+        list of dicts or pandas.Series with additional information about the raws
+    mapping: dict(str: int)
+        mapping from event description to target value
+    preload: bool
+        if True, preload the data of the Epochs objects.
+    drop_bad_windows: bool
+        If True, call `.drop_bad()` on the resulting mne.Epochs object. This
+        step allows identifying e.g., windows that fall outside of the
+        continuous recording. It is suggested to run this step here as otherwise
+        the BaseConcatDataset has to be updated as well.
+
+    Returns
+    -------
+    windows_datasets: BaseConcatDataset
+        X and y transformed to a dataset format that is compativle with skorch
+        and braindecode
+    """
+    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)]
+    else:
+        base_datasets = [BaseDataset(raw) for raw in raws]
+
+    base_datasets = BaseConcatDataset(base_datasets)
+    windows_datasets = create_windows_from_events(
+        base_datasets,
+        trial_start_offset_samples=trial_start_offset_samples,
+        trial_stop_offset_samples=trial_stop_offset_samples,
+        window_size_samples=window_size_samples,
+        window_stride_samples=window_stride_samples,
+        drop_last_window=drop_last_window,
+        mapping=mapping,
+        drop_bad_windows=drop_bad_windows,
+        preload=preload
+    )
+    return windows_datasets
+
+
+def create_from_mne_epochs(list_of_epochs, window_size_samples,
+                           window_stride_samples, drop_last_window):
+    """Create WindowsDatasets from mne.Epochs
+
+    Parameters
+    ----------
+    list_of_epochs: array-like
+        list of mne.Epochs
+    window_size_samples: int
+        window size
+    window_stride_samples: int
+        stride between windows
+    drop_last_window: bool
+        whether or not have a last overlapping window, when
+        windows do not equally divide the continuous signal
+
+    Returns
+    -------
+    windows_datasets: BaseConcatDataset
+        X and y transformed to a dataset format that is compativle with skorch
+        and braindecode
+    """
+    _check_windowing_arguments(0, 0, window_size_samples,
+                               window_stride_samples)
+
+    list_of_windows_ds = []
+    for epochs in list_of_epochs:
+        event_descriptions = epochs.events[:, 2]
+        original_trial_starts = epochs.events[:, 0]
+        stop = len(epochs.times) - window_size_samples
+
+        # already includes last incomplete window start
+        starts = np.arange(0, stop + 1, window_stride_samples)
+
+        if not drop_last_window and starts[-1] < stop:
+            # 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]
+
+        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]]
+            })
+            # window size - 1, since tmax is inclusive
+            mne_epochs = mne.Epochs(
+                mne.io.RawArray(trial, epochs.info), fake_events,
+                baseline=None,
+                tmin=0,
+                tmax=(window_size_samples - 1) / epochs.info["sfreq"],
+                metadata=metadata)
+
+            mne_epochs.drop_bad(reject=None, flat=None)
+
+            windows_ds = WindowsDataset(mne_epochs)
+            list_of_windows_ds.append(windows_ds)
+
+    return BaseConcatDataset(list_of_windows_ds)