AutoMark.py

import tkinter as tk
from tkinter import filedialog
from tkinter import messagebox
import mne
import json
import os
import mne.export

# ====== 配置 ======
start_marker_name = "PythonStart"
cut_marker_name = "CutAfterAutoMark"


# ====== GUI 部分 ======
class EEGProcessingApp:
    def __init__(self, root):
        self.root = root
        self.root.title("EEG 数据处理工具")
        self.root.geometry("500x400")

        self.vhdr_file = None
        self.json_file = None
        self.output_prefix = "processed_data"

        self.create_widgets()

    def create_widgets(self):
        # 标题标签
        title_label = tk.Label(self.root, text="选择 EEG 数据文件和事件标记文件", font=("Arial", 14))
        title_label.pack(pady=10)

        # 输入框和按钮
        self.vhdr_button = tk.Button(self.root, text="选择 .vhdr 文件", command=self.select_vhdr)
        self.vhdr_button.pack(pady=5)

        self.json_button = tk.Button(self.root, text="选择 event_log.json 文件", command=self.select_json)
        self.json_button.pack(pady=5)

        # 标记名称输入
        self.start_marker_label = tk.Label(self.root, text="请输入开始标记名称（默认：PythonStart）")
        self.start_marker_label.pack(pady=5)
        self.start_marker_entry = tk.Entry(self.root)
        self.start_marker_entry.insert(0, "PythonStart")
        self.start_marker_entry.pack(pady=5)

        self.cut_marker_label = tk.Label(self.root, text="请输入结束标记名称（默认：CutAfterAutoMark）")
        self.cut_marker_label.pack(pady=5)
        self.cut_marker_entry = tk.Entry(self.root)
        self.cut_marker_entry.insert(0, "CutAfterAutoMark")
        self.cut_marker_entry.pack(pady=5)

        # 处理按钮
        self.process_button = tk.Button(self.root, text="开始处理", command=self.process_data)
        self.process_button.pack(pady=20)

    def select_vhdr(self):
        file_path = filedialog.askopenfilename(filetypes=[("BrainVision Header", "*.vhdr")])
        if file_path:
            self.vhdr_file = file_path
            self.vhdr_button.config(text=f"已选择: {os.path.basename(file_path)}")
            print(f"选择的 .vhdr 文件路径：{file_path}")  # 打印路径，确保路径正确

    def select_json(self):
        file_path = filedialog.askopenfilename(filetypes=[("JSON 文件", "*.json")])
        if file_path:
            self.json_file = file_path
            self.json_button.config(text=f"已选择: {os.path.basename(file_path)}")
            print(f"选择的 event_log.json 文件路径：{file_path}")  # 打印路径，确保路径正确

    def process_data(self):
        if not self.vhdr_file or not self.json_file:
            messagebox.showerror("错误", "请确保已选择 .vhdr 文件和 event_log.json 文件！")
            return

        # 获取用户输入的标记名称
        start_marker_name = self.start_marker_entry.get()
        cut_marker_name = self.cut_marker_entry.get()

        try:
            self.process_eeg_data(self.vhdr_file, self.json_file, start_marker_name, cut_marker_name)
            messagebox.showinfo("完成", "数据处理成功！")
        except Exception as e:
            messagebox.showerror("错误", f"处理过程中出现错误: {e}")

    def process_eeg_data(self, vhdr_file, json_file, start_marker_name, cut_marker_name):
        # 读取脑电数据和标记
        raw = mne.io.read_raw_brainvision(vhdr_file, preload=True)
        events, event_id = mne.events_from_annotations(raw, event_id=None, verbose=True)

        print("📌 所有标记（Annotations）：")
        print(raw.annotations)
        print("📌 所有 event_id 映射：")
        print(event_id)

        with open(json_file, 'r', encoding='utf-8') as f:
            json_data = json.load(f)

        start_time_json = json_data['start_time']
        events_json = json_data['events']

        sfreq = raw.info['sfreq']

        # 获取标记的样本点
        def get_sample_of_marker(name):
            for annot in raw.annotations:
                if annot['description'].endswith(name):
                    sample = int(annot['onset'] * sfreq)
                    print(f"✅ 找到标记 {annot['description']}，采样点位置: {sample}")
                    return sample
            raise ValueError(f"❌ 没有找到标记：{name}")

        start_sample = get_sample_of_marker(start_marker_name)
        cut_sample = get_sample_of_marker(cut_marker_name)

        # 取出现有的所有标记列表（转成可变list）
        old_ann = raw.annotations
        new_onsets = list(old_ann.onset)
        new_durations = list(old_ann.duration)
        new_descriptions = list(old_ann.description)

        # 插入新的标记
        for e in events_json:
            delta_sec = e['time'] - start_time_json
            onset = start_sample / sfreq + delta_sec
            desc = e['title'].replace(",", " ")
            if onset > raw.times[-1]:
                print(f"⚠️ 跳过标记 {desc}：onset={onset:.2f}s 超出数据长度（{raw.times[-1]:.2f}s）")
                continue
            if onset < 0:
                print(f"⚠️ 跳过标记 {desc}：onset={onset:.2f}s 为负数")
                continue
            # 追加到列表
            new_onsets.append(onset)
            new_durations.append(0)
            new_descriptions.append("Stimulus/" + desc)
            print(f"✅ 插入标记 {desc}，onset={onset:.2f}s")

        # 用新的完整列表创建Annotations并设置回raw
        new_ann = mne.Annotations(onset=new_onsets, duration=new_durations, description=new_descriptions)
        raw.set_annotations(new_ann)

        # 裁剪数据
        cut_time_sec = cut_sample / sfreq
        raw_cropped = raw.copy().crop(tmin=cut_time_sec)

        # 保存处理后的数据
        vhdr_path = os.path.abspath(self.output_prefix + ".vhdr")
        mne.export.export_raw(vhdr_path, raw_cropped, fmt="brainvision")

        # 打印输出文件路径
        print(f"\n🎉 全部完成！输出文件：")
        for ext in ['.vhdr', '.vmrk', '.eeg']:
            print(f"  {self.output_prefix}{ext}")

if __name__ == "__main__":
    # 创建主窗口并运行
    root = tk.Tk()
    app = EEGProcessingApp(root)
    root.mainloop()