初始化zmq 项目

concentration/algorithm/calculate_focus.py

@@ -0,0 +1,396 @@
+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)