kraken/backend/tests/core/mi/test_erds.py

""" RED/ERS 单元测试 """
import os

from core.sig_chain.sig_reader import Reader
from core.mi.eeg_erds import ErdErs

TEST_DATA_PATH = "tests/data/"
BDF_FILE_PATH = os.path.join(TEST_DATA_PATH, "eeg_raw_data.bdf")
ERDS_FILE_PATH = os.path.join(TEST_DATA_PATH, "erds.png")
TFR_ERDS_FILE_PATH = os.path.join(TEST_DATA_PATH, "tfr_erds.png")


def setup_module():
    if not os.path.exists(TEST_DATA_PATH):
        os.makedirs(TEST_DATA_PATH)


def teardown_module():
    if os.path.exists(ERDS_FILE_PATH):
        os.remove(ERDS_FILE_PATH)
    if os.path.exists(TFR_ERDS_FILE_PATH):
        os.remove(TFR_ERDS_FILE_PATH)


def test_main():
    # ERD/ERS
    # 左右手
    # [ "C3", "C4" ]
    ch_names = ["C3", "Cz", "C4"]

    reader = Reader()
    raw = reader.read(BDF_FILE_PATH, ch_names)
    raw.annotations.rename({
        "trainSuccess": "mi",
        "trainFailed": "mi",
        "restState": "rest"
    })
    raw.resample(200)

    channels = ("C3", "C4")
    erds = ErdErs(-1, 1)
    epochs, event_id_pick = erds.get_epochs(raw, channels)
    tfr = erds.process(epochs, (-1, 0), mode="percent")
    erds.draw_image(tfr, channels, ERDS_FILE_PATH)
    tfr = erds.process(epochs, (-1, 0))
    erds.draw_tfr_image(tfr, event_id_pick, channels, TFR_ERDS_FILE_PATH)