加入STFT后,可以将设置好的不同频带的波形抽取出来做单独的分析,这里以Alpha,Beta,Delta,Theta,Gamma五种波形为例!
from mne import Epochs, pick_types, events_from_annotations
from mne.io import concatenate_raws
from mne.io import read_raw_edf
from mne.datasets import eegbci
import mne
import numpy as np
import pandas as pd
import glob
import numpy as np
import os
from scipy import signal, fft
import matplotlib.pyplot as plt
path_time = "ttt.csv" # 患者发病发病起止时间表
file_dir = "chb01"
path_save = "data"
# 选择患者共有的通道
ch = ['FP1-F7', 'F7-T7', 'T7-P7', 'P7-O1', 'FP1-F3', 'F3-C3', 'C3-P3', 'P3-O1', 'FP2-F4', 'F4-C4', 'C4-P4', 'P4-O2',
'FP2-F8', 'F8-T8', 'T8-P8-0', 'P8-O2', 'FZ-CZ', 'CZ-PZ', 'P7-T7', 'T7-FT9', 'FT9-FT10', 'FT10-T8']
sfreq = 256
bandFreqs = [
{'name': 'Delta', 'fmin': 1, 'fmax': 3},
{'name': 'Theta', 'fmin': 4, 'fmax': 7},
{'name': 'Alpha', 'fmin': 8, 'fmax': 13},
{'name': 'Beta', 'fmin': 14, 'fmax': 31},
{'name': 'Gamma', 'fmin': 31, 'fmax': 40}
]
# 定义STFT函数
def STFT(epochsData, sfreq, band=bandFreqs):
f, t, Zxx = signal.stft(epochsData, fs=sfreq)
bandResult = []
for iter_freq in band:
index = np.where((iter_freq['fmin'] < f) & (f < iter_freq['fmax']))
portion = np.zeros(Zxx.shape, dtype=np.complex_)
portion[:, :, index, :] = Zxx[:, :, index, :]
_, xrec = signal.istft(portion, fs=sfreq)
# 保存滤波后的结果
bandResult.append(xrec)
return bandResult
time = pd.read_csv(path_time,index_col="chb")
files = sorted(os.listdir(file_dir))
for file in files:
if os.path.splitext(file)[1] == '.edf':
f = os.path.splitext(file)[0]
f_str = str(os.path.splitext(os.path.splitext(file)[0])[0])
if i == 0:
raws = mne.io.read_raw_edf(file_dir+"/" + file,preload=True,verbose=False)
raws.pick_channels(ch)
raws.filter(0.1,50.,method='iir')
raw_d,raw_t = raws[:,:]
i+=1
else:
i+=1
if f_str in time.index:
time.loc[f_str]['start'] = time.loc[f_str]['start'] * 256 + len(raw_t)
time.loc[f_str]['end'] = time.loc[f_str]['end']*256 + len(raw_t)
raw = mne.io.read_raw_edf(file_dir+"/" + file, preload=True,verbose=False)
raw.pick_channels(ch)
raw.filter(0.1,50.,method='iir')
raws = concatenate_raws([raws,raw])
raws_d, raw_t = raws[:,:]
d, t = raws[:,:]
data = d*1e6
stft = STFT(d, sfreq)
pointNum = d.shape[0]
stftFreq = np.abs(fft.fft(stft[:pointNum]))
data = np.transpose(stftFreq, axes=(1,3,2,0))
np.save(path_save+"/"+file_dir+".npy",data)