用Python玩转WESAD和DREAMER：手把手教你读取ECG情绪识别数据集（附完整代码）

用Python玩转WESAD和DREAMER：手把手教你读取ECG情绪识别数据集（附完整代码）

情绪识别技术正逐渐从实验室走向实际应用，而ECG（心电图）信号因其客观性和连续性成为研究热点。WESAD和DREAMER作为两个重要的多模态情绪识别数据集，为研究者提供了丰富的生理信号数据。本文将带你从零开始，掌握这两个数据集的Python处理技巧。

1. 环境准备与数据获取

在开始处理数据前，我们需要搭建合适的Python环境。推荐使用Anaconda创建独立环境：

conda create -n emotion_recognition python=3.8 conda activate emotion_recognition pip install numpy pandas matplotlib scipy seaborn

WESAD数据集可直接从官网下载，而DREAMER需要提交申请。下载后建议按以下结构组织文件：

情绪识别项目/ ├── data/ │ ├── WESAD/ │ │ ├── S1.pkl │ │ ├── S2.pkl │ │ └── ... │ └── DREAMER.mat ├── scripts/ │ └── data_processing.py └── outputs/ └── figures/

提示：处理生理信号数据时，保持文件路径的规范性非常重要，可以避免许多不必要的错误。

2. WESAD数据集深度解析

WESAD数据集包含15名受试者的多模态生理信号，采样频率为700Hz。其数据结构采用Python字典形式存储，主要包含三个关键部分：

• subject: 受试者标识符（如S1, S2等）
• signal: 原始生理信号数据（分为chest和wrist两组）
• label: 情绪状态标签（0-7）

2.1 数据加载与探索

让我们创建一个专门处理WESAD数据的Python类：

import pickle import numpy as np import matplotlib.pyplot as plt class WESADLoader: def __init__(self, data_path, subject_id): self.data_path = data_path self.subject_id = subject_id self.data = self._load_data() def _load_data(self): with open(f"{self.data_path}/{self.subject_id}.pkl", 'rb') as f: return pickle.load(f, encoding='latin1') def get_chest_signals(self): return self.data['signal']['chest'] def get_labels(self): return self.data['label'] def plot_ecg_segment(self, start=0, end=5000): ecg = self.get_chest_signals()['ECG'][start:end] plt.figure(figsize=(15, 4)) plt.plot(ecg) plt.title(f"ECG Segment (Subject {self.subject_id})") plt.xlabel("Samples") plt.ylabel("Amplitude") plt.grid() plt.show()

使用示例：

loader = WESADLoader("data/WESAD", "S3") chest_data = loader.get_chest_signals() print(f"Available chest signals: {list(chest_data.keys())}") # 输出各信号长度 for signal, values in chest_data.items(): print(f"{signal}: {len(values)} samples") # 可视化ECG片段 loader.plot_ecg_segment(10000, 15000)

2.2 情绪片段提取与分析

WESAD数据集包含多种情绪状态，我们可以按标签提取特定情绪的数据：

def extract_emotion_segments(loader, emotion_code): labels = loader.get_labels() indices = np.where(labels == emotion_code)[0] chest_data = loader.get_chest_signals() segments = {} for signal in chest_data: segments[signal] = chest_data[signal][indices] return segments # 提取基线状态数据（label=1） baseline_data = extract_emotion_segments(loader, 1) # 提取压力状态数据（label=2） stress_data = extract_emotion_segments(loader, 2)

3. DREAMER数据集实战指南

DREAMER数据集记录了23名受试者在观看情绪诱发视频时的ECG和EEG信号，其数据结构更为复杂，采用MATLAB的.mat格式存储。

3.1 数据加载与结构解析

使用scipy.io加载.mat文件：

import scipy.io as sio def load_dreamer(data_path): data = sio.loadmat(data_path, simplify_cells=True) return data['DREAMER'] dreamer_data = load_dreamer('data/DREAMER.mat') print(f"数据集包含 {dreamer_data['noOfSubjects']} 名受试者") print(f"每个受试者有 {dreamer_data['noOfVideoSequences']} 个视频片段")

DREAMER数据结构关键字段：

3.2 ECG信号提取与处理

提取特定受试者的ECG数据：

def get_ecg_data(dreamer_data, subject_idx): subject_data = dreamer_data['Data'][subject_idx] return { 'baseline': subject_data['ECG']['baseline'], 'stimuli': subject_data['ECG']['stimuli'], 'valence': subject_data['ScoreValence'], 'arousal': subject_data['ScoreArousal'] } # 获取第一个受试者的ECG数据 subject_ecg = get_ecg_data(dreamer_data, 0) # 可视化第一个视频片段的ECG plt.figure(figsize=(15, 4)) plt.plot(subject_ecg['stimuli'][0]) plt.title("ECG during Emotion Induction (First Video)") plt.xlabel("Samples") plt.ylabel("Amplitude") plt.grid() plt.show()

4. 高级分析与可视化技巧

掌握了基础数据读取后，我们可以进行更深入的分析。

4.1 心率变异性（HRV）分析

HRV是情绪识别中的重要特征，我们可以从ECG信号中提取：

from scipy.signal import find_peaks def compute_hrv(ecg_signal, fs=700): # 检测R波峰值 peaks, _ = find_peaks(ecg_signal, height=np.mean(ecg_signal)*1.5, distance=fs*0.6) # 计算RR间期（毫秒） rr_intervals = np.diff(peaks) * (1000/fs) # 计算SDNN（HRV常用指标） sdnn = np.std(rr_intervals) return { 'peaks': peaks, 'rr_intervals': rr_intervals, 'sdnn': sdnn } # 计算基线状态的HRV baseline_ecg = loader.get_chest_signals()['ECG'] hrv_results = compute_hrv(baseline_ecg[:20000]) # 使用前20秒数据 plt.figure(figsize=(15, 4)) plt.plot(baseline_ecg[:20000]) plt.plot(hrv_results['peaks'], baseline_ecg[hrv_results['peaks']], 'rx') plt.title("R Peak Detection") plt.show() print(f"HRV (SDNN): {hrv_results['sdnn']:.2f} ms")

4.2 多受试者数据整合

对于需要同时分析多个受试者的情况，我们可以构建统一的数据结构：

def process_multiple_subjects(data_path, subject_ids): all_data = {} for subject in subject_ids: loader = WESADLoader(data_path, subject) signals = loader.get_chest_signals() labels = loader.get_labels() # 提取各情绪状态的平均ECG幅值 emotions = { 1: 'baseline', 2: 'stress', 3: 'amusement' } subject_results = {} for code, name in emotions.items(): segments = extract_emotion_segments(loader, code) subject_results[name] = np.mean(np.abs(segments['ECG'])) all_data[subject] = subject_results return pd.DataFrame(all_data).T # 处理前5名受试者 results_df = process_multiple_subjects("data/WESAD", [f"S{i}" for i in range(1, 6)]) print(results_df)

这将输出一个DataFrame，显示各受试者在不同情绪状态下的ECG平均幅值：

5. 常见问题与调试技巧

在实际操作中，你可能会遇到以下问题：

1. 数据加载错误：

   • 确保文件路径正确
   • 检查Python版本与pickle协议兼容性
   • 对于.mat文件，尝试不同的simplify_cells参数

2. 信号显示异常：

   # 检查数据范围 print(f"ECG range: {np.min(ecg)} to {np.max(ecg)}") # 标准化显示 ecg_normalized = (ecg - np.mean(ecg)) / np.std(ecg)

3. 内存不足处理：

   • 对于大型.mat文件，考虑分块加载
   • 使用memory_map选项：

     data = sio.loadmat('large_file.mat', mat_dtype=True, matlab_compatible=False)

4. 跨平台兼容性：

   • 路径处理使用os.path.join：

     import os file_path = os.path.join('data', 'WESAD', 'S1.pkl')

注意：生理信号数据通常包含高频噪声，在实际分析前应考虑适当的滤波处理。一个简单的带通滤波示例：

from scipy.signal import butter, filtfilt def bandpass_filter(data, lowcut=1.0, highcut=35.0, fs=700, order=4): nyq = 0.5 * fs low = lowcut / nyq high = highcut / nyq b, a = butter(order, [low, high], btype='band') return filtfilt(b, a, data) # 应用滤波器 filtered_ecg = bandpass_filter(chest_data['ECG'])