beta calculate func
This commit is contained in:
73
Tools/beta_calculate.py
Normal file
73
Tools/beta_calculate.py
Normal file
@@ -0,0 +1,73 @@
|
||||
import numpy as np
|
||||
from scipy.signal import welch
|
||||
from collections import deque
|
||||
|
||||
|
||||
|
||||
class Beta_Calculate():
|
||||
def __init__(self, Threshold_value_low, Threshold_value_high, fs=250, win_len=5, config=None):
|
||||
self.Threshold_value_low = Threshold_value_low
|
||||
self.Threshold_value_high = Threshold_value_high
|
||||
self.fs = fs
|
||||
self.beta_result = []
|
||||
self.eegQueue = deque(maxlen=win_len)
|
||||
|
||||
def calculate_all(self, data, fs, nperseg=1000):
|
||||
mean_x = np.mean(data, axis=-1, keepdims=True)
|
||||
data = data - mean_x
|
||||
freqs, psd = self.compute_psd_multichannel(data, fs, nperseg)
|
||||
beta_psd = np.sum(self.band_psd(freqs, psd, (13, 30)))
|
||||
alpha_psd = np.sum(self.band_psd(freqs, psd, (8, 13)))
|
||||
theta_psd = np.sum(self.band_psd(freqs, psd, (4, 8)))
|
||||
|
||||
print(f"[功率] β={beta_psd:.2f} | α={alpha_psd:.2f} | θ={theta_psd:.2f}")
|
||||
|
||||
return beta_psd, alpha_psd, theta_psd
|
||||
|
||||
def compute_psd_multichannel(self, data, fs=250, nperseg=1000):
|
||||
n_samples = data.shape[-1]
|
||||
if n_samples < nperseg:
|
||||
nperseg = n_samples
|
||||
|
||||
noverlap = 500
|
||||
if noverlap >= nperseg:
|
||||
noverlap = int(nperseg / 2)
|
||||
|
||||
if nperseg == 0:
|
||||
return np.array([]), np.zeros((data.shape[0], 0))
|
||||
|
||||
freqs, psd = welch(data, fs=fs, nperseg=nperseg, noverlap=noverlap, axis=-1)
|
||||
return freqs, psd
|
||||
|
||||
def band_psd(self, freqs, psd, band):
|
||||
idx = np.logical_and(freqs >= band[0], freqs <= band[1])
|
||||
return np.sum(psd[:, idx], axis=-1)
|
||||
|
||||
|
||||
def reset_queue(self):
|
||||
self.eegQueue.clear()
|
||||
|
||||
|
||||
def queueOpt(self, data):
|
||||
if data is None or data.size == 0:
|
||||
return None
|
||||
if len(self.eegQueue) < self.eegQueue.maxlen:
|
||||
self.eegQueue.append(data)
|
||||
else:
|
||||
self.eegQueue.append(data)
|
||||
|
||||
if len(self.eegQueue) == self.eegQueue.maxlen:
|
||||
eegData = np.hstack([self.eegQueue[i] for i in range(len(self.eegQueue))])
|
||||
if eegData.size == 0:
|
||||
return None
|
||||
eegData -= np.mean(eegData, axis=-1, keepdims=True)
|
||||
|
||||
beta_psd, alpha_psd, theta_psd = self.calculate_all(eegData, fs=self.fs, nperseg=1000)
|
||||
|
||||
return (beta_psd)
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user