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
| import torch import torch.nn as nn
class StressDetectionPipeline(nn.Module): """ 论文核心管道: 1. 眼动数据预处理(瞳孔直径PD + 凝视坐标) 2. 特征提取(15个PD统计特征 + 3个注视特征) 3. 时序建模(CNN/LSTM/ConvLSTM) 4. 压力分类(二分类) """ def __init__(self, config: dict): super().__init__() self.window_sec = 5 if config['model'] == 'cnn': self.model = CNNClassifier(input_dim=1) elif config['model'] == 'lstm': self.model = LSTMClassifier( input_dim=1, hidden_dim=128, num_layers=config.get('num_layers', 1) ) elif config['model'] == 'convlstm': self.model = ConvLSTMClassifier( input_dim=1, hidden_dim=64, kernel_size=3 ) def preprocess(self, pupil_diameter: torch.Tensor): """ 瞳孔直径预处理: 1. 异常值去除 2. 滤波 3. 基线校正(除法基线) """ mean = pupil_diameter.mean() std = pupil_diameter.std() mask = torch.abs(pupil_diameter - mean) < 3 * std cleaned = pupil_diameter[mask] baseline = cleaned[:30].mean() normalized = cleaned / baseline return normalized def forward(self, x: torch.Tensor): """ 前向传播 Args: x: 眼动数据,shape=(B, T) 或 (B, T, C) Returns: stress_prob: 压力概率,shape=(B, 1) """ x = self.preprocess(x) return self.model(x)
class CNNClassifier(nn.Module): """1D-CNN分类器""" def __init__(self, input_dim=1): super().__init__() self.conv = nn.Sequential( nn.Conv1d(input_dim, 32, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool1d(2), nn.Conv1d(32, 64, kernel_size=3, padding=1), nn.ReLU(), nn.AdaptiveAvgPool1d(1) ) self.fc = nn.Linear(64, 1) def forward(self, x): x = self.conv(x) x = self.fc(x.squeeze(-1)) return torch.sigmoid(x)
class LSTMClassifier(nn.Module): """LSTM分类器""" def __init__(self, input_dim=1, hidden_dim=128, num_layers=1): super().__init__() self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers=num_layers, batch_first=True) self.fc = nn.Linear(hidden_dim, 1) def forward(self, x): out, (h_n, c_n) = self.lstm(x) last_hidden = out[:, -1, :] return torch.sigmoid(self.fc(last_hidden))
def loso_cross_validation(dataset): """ 论文使用嵌套LOSO验证: - 外循环:留一受试者验证泛化性 - 内循环:超参数优化 """ subjects = dataset.get_subjects() results = [] for test_subject in subjects: train_data = dataset.get_others(test_subject) test_data = dataset.get_subject(test_subject) best_params = hyperparameter_search(train_data) model = train_model(train_data, best_params) acc = evaluate(model, test_data) results.append(acc) return results
|