Benchmark results

import argparse
import os
import pathlib
from tkinter import BOTTOM
import urllib.request
import numpy as np
import pandas as pd
from joblib import Parallel, delayed
import mne
from mne.io.brainvision.brainvision import _aux_vhdr_info

from mne_bids import write_raw_bids, print_dir_tree, make_report, BIDSPath

lemon_info = pd.read_csv(
    "./META_File_IDs_Age_Gender_Education_Drug_Smoke_SKID_LEMON.csv")
lemon_info = lemon_info.set_index("ID")
eeg_subjects = pd.read_csv('./lemon_eeg_subjects.csv')
lemon_info = lemon_info.loc[eeg_subjects.subject]
lemon_info['gender'] = lemon_info['Gender_ 1=female_2=male'].map({1: 2, 2: 1})
lemon_info['age_guess'] = np.array(
  lemon_info['Age'].str.split('-').tolist(), dtype=np.int).mean(1)
subjects = list(lemon_info.index)

def convert_lemon_to_bids(lemon_data_dir, bids_save_dir, n_jobs=1, DEBUG=False):
    """Convert TUAB dataset to BIDS format.

    Parameters
    ----------
    lemon_data_dir : str
        Directory where the original LEMON dataset is saved, e.g.
        `/storage/store3/data/LEMON_RAW`.
    bids_save_dir : str
        Directory where to save the BIDS version of the dataset.
    n_jobs : None | int
        Number of jobs for parallelization.
    """
    subjects_ = subjects
    if DEBUG:
        subjects_ = subjects[:1]

    good_subjects = Parallel(n_jobs=n_jobs)(
        delayed(_convert_subject)(subject, lemon_data_dir, bids_save_dir)
        for subject in subjects_) 
    subjects_ = [sub for sub in good_subjects if not isinstance(sub, tuple)]
    _, bad_subjects, errs = zip(*[
        sub for sub in good_subjects if isinstance(sub, tuple)])
    bad_subjects = pd.DataFrame(
        dict(subjects= bad_subjects, error=errs))
    bad_subjects.to_csv(
        '/storage/store3/data/LEMON_EEG_BIDS/bids_conv_erros.csv')
    # update the participants file as LEMON has no official age data
    participants = pd.read_csv(
        "/storage/store3/data/LEMON_EEG_BIDS/participants.tsv", sep='\t')
    participants = participants.set_index("participant_id")
    participants.loc[subjects_, 'age'] = lemon_info.loc[subjects_, 'age_guess']
    participants.to_csv(
        "/storage/store3/data/LEMON_EEG_BIDS/participants.tsv", sep='\t')


def _convert_subject(subject, data_path, bids_save_dir):
    """Get the work done for one subject"""
    try:
        fname = pathlib.Path(data_path) / subject / "RSEEG" / f"{subject}.vhdr"    
        raw = mne.io.read_raw_brainvision(fname)

        raw.set_channel_types({"VEOG": "eog"})
        montage = mne.channels.make_standard_montage('standard_1005')
        raw.set_montage(montage)
        sub_id = subject.strip("sub-")
        raw.info['subject_info'] = {
            'participant_id': sub_id,
            'sex': lemon_info.loc[subject, 'gender'],
            'age': lemon_info.loc[subject, 'age_guess'],
            # XXX LEMON shares no public age 
            'hand': lemon_info.loc[subject, 'Handedness']
        }
        events, _ = mne.events_from_annotations(raw)

        events = events[(events[:, 2] == 200) | (events[:, 2] == 210)]
        event_id = {"eyes/open": 200, "eyes/closed": 210}
        bids_path = BIDSPath(
            subject=sub_id, session=None, task='RSEEG',
            run=None,
            root=bids_save_dir, datatype='eeg', check=True)

        write_raw_bids(
            raw,
            bids_path,
            events_data=events,
            event_id=event_id,
            overwrite=True
        )
    except Exception as err:
        print(err)
        return ("BAD", subject, err)
    return subject


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Convert LEMON to BIDS.')
    parser.add_argument(
        '--lemon_data_dir', type=str,
        default='/storage/store3/data/LEMON_RAW',
        help='Path to the original data.')
    parser.add_argument(
        '--bids_data_dir', type=str,
        default=pathlib.Path("/storage/store3/data/LEMON_EEG_BIDS"),
        help='Path to where the converted data should be saved.')
    parser.add_argument(
        '--n_jobs', type=int, default=1,
        help='number of parallel processes to use (default: 1)')
    parser.add_argument(
        '--DEBUG', type=bool, default=False,
        help='activate debugging mode')
    args = parser.parse_args()

    convert_lemon_to_bids(
        args.lemon_data_dir, args.bids_data_dir, n_jobs=args.n_jobs,
        DEBUG=args.DEBUG)

    print_dir_tree(args.bids_data_dir)
    print(make_report(args.bids_data_dir))