90 lines
3.9 KiB
Python
90 lines
3.9 KiB
Python
import os
|
|
import mne
|
|
|
|
|
|
def load_bad_annotations(filepath, fileending="badSegments.csv"):
|
|
"""
|
|
Loads the annotations for bad segments
|
|
:param filepath: The path to the file we want to load
|
|
:param fileending: Depending if the subject, for which we load the annotations, was preprocessed manually,
|
|
the ending of the filename will be different.
|
|
The default are file endings of the given preprocessed annotations: "badSegments.csv"
|
|
For the manual preprocessed annotations, the file endings are: "badannotations.csv"
|
|
:return: The mne annotations
|
|
"""
|
|
if os.path.isfile(filepath + "_" + fileending):
|
|
return mne.read_annotations(filepath + "_" + fileending)
|
|
|
|
|
|
def load_preprocessed_data(subject, dataset):
|
|
"""
|
|
Load the raw object as well as the annotations of the preprocessed file
|
|
:param subject: The subject, for which we want to load the raw object
|
|
:param dataset: The currently viewed dataset
|
|
:param selected_subjects: The manually preprocessed subjects
|
|
:return: The raw object
|
|
"""
|
|
folder = "Dataset/" + dataset + "/sub-" + subject + "/ses-" + dataset + "/eeg/"
|
|
filepath = folder + "sub-" + subject + "_task-" + dataset
|
|
raw = mne.io.read_raw_fif(filepath + "_cleaned.fif")
|
|
return raw
|
|
|
|
|
|
def get_keys_for_events(stimulus=None, condition=None):
|
|
"""
|
|
For a given stimulus and condition get all the event keys.
|
|
:param stimulus: Either 'face' or 'car' or 'None' for no stimulus
|
|
:param condition: Either 'intact' or 'scrambled' or 'None' for no condition
|
|
:return: A list of keys or 'stimulus' if neither stimulus or condition was given
|
|
"""
|
|
if stimulus == 'face':
|
|
if condition == 'intact':
|
|
return ["stimulus:{}".format(k) for k in range(1, 41)]
|
|
elif condition == 'scrambled':
|
|
return ["stimulus:{}".format(k) for k in range(101, 141)]
|
|
else: # All faces
|
|
return ["stimulus:{}".format(k) for k in list(range(1, 41)) + list(range(101, 141))]
|
|
elif stimulus == 'car':
|
|
if condition == 'intact':
|
|
return ["stimulus:{}".format(k) for k in range(41, 81)]
|
|
elif condition == 'scrambled':
|
|
return ["stimulus:{}".format(k) for k in range(141, 181)]
|
|
else: # All cars
|
|
return ["stimulus:{}".format(k) for k in list(range(41, 81)) + list(range(141, 181))]
|
|
else: # If no stimulus is given
|
|
if condition == 'intact':
|
|
return ["stimulus:{}".format(k) for k in range(1, 41) and range(41, 81)]
|
|
elif condition == 'scrambled':
|
|
return ["stimulus:{}".format(k) for k in list(range(101, 141)) + list(range(141, 181))]
|
|
else: # Every stimulus
|
|
return "stimulus"
|
|
|
|
|
|
def get_epochs(raw, conditions=None, picks=None, tmin=-0.1, tmax=1):
|
|
"""
|
|
Returns the epochs for a given dataset
|
|
:param raw: the dataset
|
|
:param conditions: A List of tuples, of the Form [(stimulus, condition), (stimulus,condition)]
|
|
i.e. [('face',None), ('car', 'scrambled')] returns the epochs where the stimulus is face and the stim+condition is car+scrambled
|
|
default is None, i.e. everything
|
|
:param picks: a list. Additional criteria for picking the epochs, e.g. channels, etc.
|
|
:param tmin: onset before the event
|
|
:param tmax: end after the event
|
|
:return:
|
|
"""
|
|
|
|
events, events_dict = mne.events_from_annotations(raw)
|
|
events_dict_key = {}
|
|
if conditions is None:
|
|
conditions = [(None, None)]
|
|
|
|
for condition in conditions:
|
|
wanted_keys = get_keys_for_events(condition[0], condition[1])
|
|
if wanted_keys == "stimulus":
|
|
wanted_keys = [e for e in events_dict.keys() if "stimulus" in e]
|
|
events_dict_key.update(dict((k, events_dict[k]) for k in wanted_keys if k in events_dict))
|
|
epochs = mne.Epochs(raw, events, events_dict_key, tmin=tmin, tmax=tmax, reject_by_annotation=False, picks=picks)
|
|
epochs.drop_bad()
|
|
|
|
return epochs, events_dict_key
|