396 lines
15 KiB
Python
396 lines
15 KiB
Python
import numpy as np
|
||
from scipy.signal import welch
|
||
from scipy.fft import fft
|
||
from scipy import signal
|
||
from collections import deque
|
||
import time
|
||
import os
|
||
# import logging
|
||
import base64
|
||
import io
|
||
|
||
# logger = logging.getLogger(__name__)
|
||
#
|
||
# try:
|
||
# import matplotlib
|
||
# matplotlib.use('Agg')
|
||
# import matplotlib.pyplot as plt
|
||
# MATPLOTLIB_AVAILABLE = True
|
||
# except ImportError:
|
||
# MATPLOTLIB_AVAILABLE = False
|
||
# logger.warning("matplotlib未安装，报告图表功能不可用")
|
||
|
||
|
||
class Calculate():
|
||
def __init__(self, Threshold_value_low, Threshold_value_high, fs=250, win_len=10):
|
||
self.Threshold_value_low = Threshold_value_low
|
||
self.Threshold_value_high = Threshold_value_high
|
||
self.fs = fs
|
||
self.focus_result = []
|
||
self.CLI_result = []
|
||
self.EVI_result = []
|
||
self.eegQueue = deque(maxlen=win_len)
|
||
|
||
# # 存储历史数据用于绘图
|
||
# self.beta_history = []
|
||
# self.alpha_history = []
|
||
# self.theta_history = []
|
||
# self.focus_history = []
|
||
# self.timestamp_history = []
|
||
#
|
||
# # 记录开始时间
|
||
# self.start_time = None
|
||
# self.recording = False
|
||
#
|
||
# # 图表保存路径
|
||
# self.chart_dir = "reports"
|
||
# if not os.path.exists(self.chart_dir):
|
||
# os.makedirs(self.chart_dir)
|
||
# print(f"[调试] 创建目录: {self.chart_dir}")
|
||
|
||
# 初始化滤波器
|
||
self.b_notch, self.a_notch = signal.iirnotch(50 / (self.fs/2), 30)
|
||
self.b_design = signal.firwin(65, [2 / (self.fs/2), 40 / (self.fs/2)], pass_zero=False)
|
||
|
||
print("[调试] Calculate 类初始化完成")
|
||
|
||
def calculate_focus(self, beta, alpha, theta):
|
||
"""
|
||
专注度计算 - 固定映射版本
|
||
"""
|
||
# 原始比值
|
||
raw = beta / (alpha + theta + 1e-10)
|
||
|
||
# Sigmoid 映射：让 raw 在 0.3-1.5 区间敏感
|
||
# 参数可调：
|
||
# k = 12 (斜率，越大越陡)
|
||
# x0 = 0.6 (中心点，raw=0.6时focus≈50)
|
||
k = 12.0
|
||
x0 = 0.6
|
||
focus = 100.0 / (1.0 + np.exp(-k * (raw - x0)))
|
||
|
||
# 可选：添加滑动平均平滑
|
||
return int(focus)
|
||
|
||
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}")
|
||
|
||
|
||
focus_score = self.calculate_focus(beta_psd, alpha_psd, theta_psd)
|
||
focus_score = max(0, min(100, focus_score))
|
||
|
||
self.focus_result.append(focus_score)
|
||
if len(self.focus_result) > 3:
|
||
self.focus_result.pop(0)
|
||
final_focus = int(self.simple_moving_average(self.focus_result, window_size=5))
|
||
|
||
cli_denom = alpha_psd + beta_psd
|
||
CLI_score = np.log(theta_psd / (cli_denom + 1e-10)) if cli_denom > 0 else 0
|
||
self.CLI_result.append(CLI_score)
|
||
if len(self.CLI_result) > 5:
|
||
self.CLI_result.pop(0)
|
||
final_CLI = round(self.simple_moving_average(self.CLI_result, window_size=5), 2)
|
||
|
||
return final_focus, final_CLI, 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 simple_moving_average(self, data, window_size=5):
|
||
if len(data) == 0:
|
||
return 30
|
||
window = data[-window_size:]
|
||
return sum(window) / len(window)
|
||
|
||
def reset_queue(self):
|
||
self.eegQueue.clear()
|
||
|
||
# def start_recording(self):
|
||
# """开始记录数据"""
|
||
# self.recording = True
|
||
# self.start_time = time.time()
|
||
# self.beta_history = []
|
||
# self.alpha_history = []
|
||
# self.theta_history = []
|
||
# self.focus_history = []
|
||
# self.timestamp_history = []
|
||
# print("[调试] ========== 开始记录专注度数据 ==========")
|
||
|
||
# def stop_recording(self):
|
||
# """停止记录并生成图表"""
|
||
# print(f"[调试] stop_recording被调用, recording={self.recording}, focus_history长度={len(self.focus_history)}")
|
||
# self.recording = False
|
||
# if len(self.focus_history) > 0:
|
||
# print("[调试] 数据非空，开始生成图表...")
|
||
# # 保存到本地文件
|
||
# chart_path = self.save_chart_to_file()
|
||
# if chart_path:
|
||
# print(f"[调试] 本地文件保存成功: {chart_path}")
|
||
# else:
|
||
# print("[调试] 本地文件保存失败")
|
||
# # 生成base64编码
|
||
# base64_data = self.generate_chart_base64()
|
||
# return base64_data
|
||
# else:
|
||
# print("[调试] 没有数据可保存，focus_history为空")
|
||
# return None
|
||
|
||
# def add_data_point(self, focus, beta, alpha, theta):
|
||
# if not self.recording:
|
||
# return
|
||
# current_time = time.time()
|
||
# elapsed = current_time - self.start_time
|
||
#
|
||
# self.beta_history.append(beta)
|
||
# self.alpha_history.append(alpha)
|
||
# self.theta_history.append(theta)
|
||
# self.focus_history.append(focus)
|
||
# self.timestamp_history.append(elapsed)
|
||
# print(f"[调试] 记录数据点: time={elapsed:.1f}s, focus={focus}, beta={beta:.2f}")
|
||
|
||
# def save_chart_to_file(self):
|
||
# """
|
||
# 保存图表到本地文件（唯一实现）
|
||
# """
|
||
# print(f"[调试] save_chart_to_file被调用, MATPLOTLIB_AVAILABLE={MATPLOTLIB_AVAILABLE}")
|
||
#
|
||
# if not MATPLOTLIB_AVAILABLE:
|
||
# print("[调试] matplotlib不可用，无法保存")
|
||
# return None
|
||
#
|
||
# if len(self.focus_history) < 2:
|
||
# print(f"[调试] 数据点不足，需要至少2个点，当前{len(self.focus_history)}个点")
|
||
# return None
|
||
#
|
||
# print(f"[调试] 开始保存图表到本地文件...")
|
||
#
|
||
# # 确保所有列表长度一致
|
||
# min_len = min(len(self.beta_history), len(self.alpha_history),
|
||
# len(self.theta_history), len(self.focus_history),
|
||
# len(self.timestamp_history))
|
||
#
|
||
# print(f"[调试] 数据长度: min_len={min_len}")
|
||
#
|
||
# beta_list = self.beta_history[:min_len]
|
||
# alpha_list = self.alpha_history[:min_len]
|
||
# theta_list = self.theta_history[:min_len]
|
||
# focus_list = self.focus_history[:min_len]
|
||
# times = self.timestamp_history[:min_len]
|
||
#
|
||
# # 生成文件名
|
||
# timestamp = time.strftime("%Y%m%d_%H%M%S")
|
||
# chart_path = os.path.join(self.chart_dir, f"concentration_report_{timestamp}.png")
|
||
# print(f"[调试] 保存路径: {chart_path}")
|
||
#
|
||
# try:
|
||
# # 创建图表
|
||
# fig, ax1 = plt.subplots(figsize=(14, 8))
|
||
#
|
||
# # 左Y轴：功率数据
|
||
# ax1.plot(times, beta_list, 'b-', linewidth=1.5, alpha=0.8, label='Beta Power')
|
||
# ax1.plot(times, alpha_list, 'g-', linewidth=1.5, alpha=0.8, label='Alpha Power')
|
||
# ax1.plot(times, theta_list, 'orange', linewidth=1.5, alpha=0.8, label='Theta Power')
|
||
# ax1.set_xlabel('Time (s)', fontsize=12)
|
||
# ax1.set_ylabel('Band Power', fontsize=12, color='black')
|
||
# ax1.tick_params(axis='y', labelcolor='black')
|
||
# ax1.legend(loc='upper left')
|
||
# ax1.grid(True, alpha=0.3)
|
||
#
|
||
# # 右Y轴：专注度
|
||
# ax2 = ax1.twinx()
|
||
# ax2.plot(times, focus_list, 'r-', linewidth=2, alpha=0.9, label='Focus (%)')
|
||
# ax2.set_ylabel('Focus (%)', fontsize=12, color='red')
|
||
# ax2.tick_params(axis='y', labelcolor='red')
|
||
# ax2.set_ylim(0, 105)
|
||
# ax2.legend(loc='upper right')
|
||
#
|
||
# # 标题
|
||
# duration = times[-1] if times else 0
|
||
# avg_focus = np.mean(focus_list) if focus_list else 0
|
||
# plt.title(f'Concentration and EEG Band Power Trend\nDuration: {duration:.1f}s, Avg Focus: {avg_focus:.1f}%',
|
||
# fontsize=14)
|
||
#
|
||
# plt.tight_layout()
|
||
# plt.savefig(chart_path, dpi=150, bbox_inches='tight')
|
||
# plt.close()
|
||
#
|
||
# print(f"\n{'='*60}")
|
||
# print(f"专注度报告图片已保存到本地:")
|
||
# print(f" 文件路径: {chart_path}")
|
||
# print(f" 数据点数: {min_len}")
|
||
# print(f" 时长: {duration:.1f}秒")
|
||
# print(f" 平均专注度: {avg_focus:.1f}%")
|
||
# print(f"{'='*60}\n")
|
||
#
|
||
# return chart_path
|
||
#
|
||
# except Exception as e:
|
||
# print(f"[调试] 保存文件时出错: {e}")
|
||
# import traceback
|
||
# traceback.print_exc()
|
||
# return None
|
||
#
|
||
# def generate_chart_base64(self):
|
||
# """
|
||
# 生成图表的base64编码（用于网络传输）
|
||
# """
|
||
# if not MATPLOTLIB_AVAILABLE:
|
||
# return None
|
||
#
|
||
# if len(self.focus_history) < 2:
|
||
# return None
|
||
#
|
||
# min_len = min(len(self.beta_history), len(self.alpha_history),
|
||
# len(self.theta_history), len(self.focus_history),
|
||
# len(self.timestamp_history))
|
||
#
|
||
# beta_list = self.beta_history[:min_len]
|
||
# alpha_list = self.alpha_history[:min_len]
|
||
# theta_list = self.theta_history[:min_len]
|
||
# focus_list = self.focus_history[:min_len]
|
||
# times = self.timestamp_history[:min_len]
|
||
#
|
||
# fig, ax1 = plt.subplots(figsize=(14, 8))
|
||
#
|
||
# ax1.plot(times, beta_list, 'b-', linewidth=1.5, alpha=0.8, label='Beta Power')
|
||
# ax1.plot(times, alpha_list, 'g-', linewidth=1.5, alpha=0.8, label='Alpha Power')
|
||
# ax1.plot(times, theta_list, 'orange', linewidth=1.5, alpha=0.8, label='Theta Power')
|
||
# ax1.set_xlabel('Time (s)', fontsize=12)
|
||
# ax1.set_ylabel('Band Power', fontsize=12, color='black')
|
||
# ax1.tick_params(axis='y', labelcolor='black')
|
||
# ax1.legend(loc='upper left')
|
||
# ax1.grid(True, alpha=0.3)
|
||
#
|
||
# ax2 = ax1.twinx()
|
||
# ax2.plot(times, focus_list, 'r-', linewidth=2, alpha=0.9, label='Focus (%)')
|
||
# ax2.set_ylabel('Focus (%)', fontsize=12, color='red')
|
||
# ax2.tick_params(axis='y', labelcolor='red')
|
||
# ax2.set_ylim(0, 105)
|
||
# ax2.legend(loc='upper right')
|
||
#
|
||
# duration = times[-1] if times else 0
|
||
# avg_focus = np.mean(focus_list) if focus_list else 0
|
||
# plt.title(f'Concentration and EEG Band Power Trend\nDuration: {duration:.1f}s, Avg Focus: {avg_focus:.1f}%',
|
||
# fontsize=14)
|
||
#
|
||
# plt.tight_layout()
|
||
#
|
||
# buffer = io.BytesIO()
|
||
# plt.savefig(buffer, format='png', dpi=150, bbox_inches='tight')
|
||
# buffer.seek(0)
|
||
# image_base64 = base64.b64encode(buffer.read()).decode('utf-8')
|
||
# plt.close()
|
||
#
|
||
# return image_base64
|
||
|
||
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)
|
||
eegData = signal.lfilter(self.b_notch, self.a_notch, eegData)
|
||
eegData = signal.lfilter(self.b_design, 1, eegData)
|
||
focus_score, CLI_score, beta, alpha, theta = self.calculate_all(eegData, fs=self.fs, nperseg=1000)
|
||
|
||
# self.add_data_point(focus_score, beta, alpha, theta)
|
||
|
||
return focus_score
|
||
return None
|
||
|
||
|
||
class Calculate2():
|
||
def __init__(self, Threshold_value_low, Threshold_value_high):
|
||
self.Threshold_value_low = Threshold_value_low
|
||
self.Threshold_value_high = Threshold_value_high
|
||
self.focus_result = []
|
||
self.theta_result = []
|
||
self.alpha_result = []
|
||
self.flow_result = []
|
||
|
||
def calculate_all(self, data, fs, L=2500):
|
||
mean_x = np.mean(data, axis=-1, keepdims=True)
|
||
data = data - mean_x
|
||
|
||
Y = fft(data, axis=-1)
|
||
P2 = np.abs(Y / L)
|
||
P1 = P2[:, :L // 2 + 1]
|
||
P1[:, 1:-1] = 2 * P1[:, 1:-1]
|
||
|
||
beta_power = self.PSD(P1, L, fs, 13, 30)
|
||
alpha_power = self.PSD(P1, L, fs, 8, 13)
|
||
theta_power = self.PSD(P1, L, fs, 4, 8)
|
||
gamma_power = self.PSD(P1, L, fs, 30, 100)
|
||
|
||
focus_score = beta_power / (alpha_power + theta_power)
|
||
print('focus score:', focus_score)
|
||
focus_score = ((focus_score - self.Threshold_value_low) * 100) / (self.Threshold_value_high - self.Threshold_value_low)
|
||
self.focus_result.append(focus_score)
|
||
if len(self.focus_result) > 3:
|
||
self.focus_result.pop(0)
|
||
final_focus = int(self.simple_moving_average(self.focus_result, window_size=3))
|
||
|
||
self.theta_result.append(theta_power)
|
||
if len(self.theta_result) > 30:
|
||
self.theta_result.pop(0)
|
||
self.alpha_result.append(alpha_power)
|
||
if len(self.alpha_result) > 30:
|
||
self.alpha_result.pop(0)
|
||
rest_theta = self.simple_moving_average(self.theta_result, window_size=30)
|
||
rest_alpha = self.simple_moving_average(self.alpha_result, window_size=30)
|
||
distraction_score = (theta_power / rest_theta) * (1 - (alpha_power / rest_alpha))
|
||
|
||
flow_score = gamma_power / beta_power
|
||
flow_score = (flow_score / self.Threshold_value_high) * 100
|
||
self.flow_result.append(flow_score)
|
||
if len(self.flow_result) > 3:
|
||
self.flow_result.pop(0)
|
||
final_flow = int(self.simple_moving_average(self.flow_result, window_size=3))
|
||
|
||
return final_focus, distraction_score, final_flow
|
||
|
||
def PSD(self, P1, L, Fs, s_freq, e_freq):
|
||
s_point = round(s_freq * L / Fs)
|
||
e_point = round(e_freq * L / Fs)
|
||
x, y = P1.shape
|
||
band_PSD = 0
|
||
for i in range(x):
|
||
for j in range(s_point, e_point):
|
||
band_PSD += P1[i, j] ** 2
|
||
return band_PSD
|
||
|
||
def simple_moving_average(self, data, window_size=3):
|
||
if len(data) == 0:
|
||
return []
|
||
window = data[-window_size:]
|
||
return sum(window) / len(window) |