gen_eeg_event_files.py

#!/usr/bin/env python

"""
Parse a SF logfile and BDF file and generate EEGLAB compat event file
"""

from __future__ import print_function, division

import sys, os
import csv
import glob
import mne
import argparse
import operator
import json
import numpy as np
import zlib

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('subject_dir')
    parser.add_argument('output_dir')
    args = parser.parse_args()
    if os.path.isdir(args.subject_dir):
        base, sid = os.path.split(args.subject_dir)
        base, _ = os.path.split(base)
        base, _ = os.path.split(base)
        bdf_file = os.path.join(base,'raw-eeg','%s.bdf' % sid)
        print(bdf_file)
        has_bdf = os.path.isfile(bdf_file)
        if has_bdf:

            raw = mne.io.read_raw_edf(bdf_file, stim_channel=-1)
            eeg_events = mne.find_events(raw, stim_channel='STI 014', shortest_event=1)

            start_times = []
            for e in eeg_events:
                if e[1]==0 and e[2]==1:
                    start_times.append([e[0], e[0]/raw.info["sfreq"]])

        game_events2 = []
        death_times = []

        for f in glob.glob("%s/*.evt" % args.subject_dir):
            g = int(os.path.split(f)[-1].split(".")[0].split("-")[-1])
            ueid = 0
            with open(f, "r") as fin:
                lines = [l.strip().split(" ") for l in fin.readlines()]
                gs = float(lines[0][1])
                N = len(lines)
                # epoch = 0
                for i in range(N):
                    e = lines[i]
                    game_marker = None
                    et = float(e[1]) - gs
                    if i == 0:
                        game_marker = "game-start"
                    elif i == (N-1):
                        game_marker = "game-end"
                    if game_marker != None:
                        if len(e) < 3:
                            e.append(game_marker)
                        else:
                            e[2] = "%s,%s" % (e[2], game_marker)
                    if len(e) == 3:
                        for ee in e[2].split(","):
                            ueid += 1
                            if has_bdf:
                                game_events2.append([et+start_times[g-1][1], g, ee, ueid, e[0]])
                            else:
                                game_events2.append([et, g, ee, ueid, e[0]])

        game_events2 = sorted(game_events2, key=operator.itemgetter(0, 1))

        if not os.path.isdir(args.output_dir):
            os.makedirs(args.output_dir)
        with open(os.path.join(args.output_dir, "%s.ext" % (sid)), "w") as fout:
            writer = csv.writer(fout, delimiter="\t")
            writer.writerow(["latency","game","type","euid","gametime"])
            for ge in game_events2:
                writer.writerow(ge)