1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229
| """ EEG 认知负荷计算算法
核心方法:θ/α 比值 + 前额叶 β 波功率 + 多模态融合 """
import numpy as np from scipy.signal import butter, filtfilt, welch from typing import Dict, Tuple
def extract_eeg_bands(eeg_signal: np.ndarray, fps: int = 256) -> Dict[str, np.ndarray]: """ 提取 EEG 频带信号 Args: eeg_signal: 原始 EEG 信号, shape=(N,) fps: 采样率 Returns: bands: { 'delta': Delta 波段信号, 'theta': Theta 波段信号, 'alpha': Alpha 波段信号, 'beta': Beta 波段信号, 'gamma': Gamma 波段信号 } Bitbrain 方法: 使用 Butterworth 带通滤波器提取各频带 """ nyquist = fps / 2 band_defs = { 'delta': (0.5, 4.0), 'theta': (4.0, 8.0), 'alpha': (8.0, 13.0), 'beta': (13.0, 30.0), 'gamma': (30.0, 45.0) } bands = {} for band_name, (low, high) in band_defs.items(): low_norm = low / nyquist high_norm = high / nyquist b, a = butter(4, [low_norm, high_norm], btype='band') bands[band_name] = filtfilt(b, a, eeg_signal) return bands
def calculate_band_power(signal: np.ndarray, fps: int = 256) -> float: """ 计算信号功率 Args: signal: 频带信号 fps: 采样率 Returns: power: 信号功率(μV²) """ freqs, psd = welch(signal, fs=fPS, nperseg=min(256, len(signal))) power = np.sum(psd) return power
def calculate_cognitive_load_index(eeg_channels: Dict[str, np.ndarray], fps: int = 256) -> Dict[str, float]: """ 计算认知负荷指数 Args: eeg_channels: 多通道 EEG 信号 { 'Fz': 前额叶信号, 'Cz': 中央信号, 'Pz': 顶叶信号, 'Oz': 枕叶信号 } fps: 采样率 Returns: indices: { 'theta_alpha_ratio': θ/α 比值, 'frontal_beta_power': 前额叶 β 功率, 'cognitive_load_score': 综合认知负荷分数, 'fatigue_index': 疲劳指数 } Bitbrain 方法: 认知负荷 = θ/α 比值 * 权重 + β 抑制 * 权重 疲劳指数 = θ 功率增强 + α 功率增强 """ frontal_bands = extract_eeg_bands(eeg_channels['Fz'], fps) theta_power = calculate_band_power(frontal_bands['theta'], fps) alpha_power = calculate_band_power(frontal_bands['alpha'], fps) beta_power = calculate_band_power(frontal_bands['beta'], fps) theta_alpha_ratio = theta_power / (alpha_power + 1e-6) frontal_beta_power = beta_power baseline_ratio = 1.0 cognitive_load_score = min(100, (theta_alpha_ratio / baseline_ratio - 1) * 50) baseline_theta = calculate_band_power( extract_eeg_bands(eeg_channels['Fz'][:256], fps)['theta'], fps ) baseline_alpha = calculate_band_power( extract_eeg_bands(eeg_channels['Fz'][:256], fps)['alpha'], fps ) fatigue_index = (theta_power / baseline_theta - 1) * 50 + (alpha_power / baseline_alpha - 1) * 50 return { 'theta_alpha_ratio': theta_alpha_ratio, 'frontal_beta_power': frontal_beta_power, 'cognitive_load_score': cognitive_load_score, 'fatigue_index': fatigue_index }
def detect_cognitive_distraction_eeg( cognitive_load_score: float, fatigue_index: float, threshold_load: float = 60.0, threshold_fatigue: float = 50.0 ) -> Tuple[bool, str]: """ 检测认知分心 Args: cognitive_load_score: 认知负荷分数 fatigue_index: 疲劳指数 threshold_load: 认知负荷阈值 threshold_fatigue: 疲劳阈值 Returns: is_distraction: 是否认知分心 level: 分心等级(normal/high_load/fatigue/extreme) IMS落地应用: 认知分心判定逻辑: 1. 认知负荷分数 >60 → 高负荷警告 2. 疲劳指数 >50 → 疲劳警告 3. 两者同时触发 → 极端状态 """ is_distraction = False level = "normal" if cognitive_load_score > threshold_load and fatigue_index > threshold_fatigue: is_distraction = True level = "extreme" elif cognitive_load_score > threshold_load: is_distraction = True level = "high_load" elif fatigue_index > threshold_fatigue: is_distraction = True level = "fatigue" return is_distraction, level
if __name__ == "__main__": fps = 256 duration = 10 n_samples = fps * duration t = np.linspace(0, duration, n_samples) normal_eeg = ( 10 * np.sin(2 * np.pi * 10 * t) + 5 * np.sin(2 * np.pi * 6 * t) + 3 * np.sin(2 * np.pi * 20 * t) + np.random.randn(n_samples) * 2 ) high_load_eeg = ( 5 * np.sin(2 * np.pi * 10 * t) + 15 * np.sin(2 * np.pi * 6 * t) + 2 * np.sin(2 * np.pi * 20 * t) + np.random.randn(n_samples) * 2 ) print("="*60) print("EEG 认知负荷检测测试") print("="*60) indices_normal = calculate_cognitive_load_index({'Fz': normal_eeg}, fps) print(f"\n正常状态:") print(f" θ/α 比值: {indices_normal['theta_alpha_ratio']:.2f}") print(f" 认知负荷分数: {indices_normal['cognitive_load_score']:.1f}") print(f" 疲劳指数: {indices_normal['fatigue_index']:.1f}") is_distraction, level = detect_cognitive_distraction_eeg( indices_normal['cognitive_load_score'], indices_normal['fatigue_index'] ) print(f" 状态: {level}") indices_high_load = calculate_cognitive_load_index({'Fz': high_load_eeg}, fps) print(f"\n高负荷状态:") print(f" θ/α 比值: {indices_high_load['theta_alpha_ratio']:.2f}") print(f" 认知负荷分数: {indices_high_load['cognitive_load_score']:.1f}") print(f" 疲劳指数: {indices_high_load['fatigue_index']:.1f}") is_distraction_high, level_high = detect_cognitive_distraction_eeg( indices_high_load['cognitive_load_score'], indices_high_load['fatigue_index'] ) print(f" 状态: {level_high}")
|