987 lines
34 KiB
Python
987 lines
34 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
runDecoder.py - BDF EEG Depression Assessment
|
||
|
||
功能：
|
||
1. 读取 raw_data 文件夹的第一个 .bdf 格式文件
|
||
2. 预处理：坏通道剔除、50Hz陷波、0.8-40Hz带通、幅值过滤、ICA去伪迹
|
||
3. 调用 infer_pth.py 中的 predict_hc_mdd 进行 HC/MDD 分类预测
|
||
4. 保存图表（EEG、PSD、Topomap）
|
||
5. 生成 ResultData.txt
|
||
|
||
"""
|
||
import matplotlib
|
||
matplotlib.use('Agg')
|
||
import numpy as np
|
||
import os
|
||
import shutil
|
||
import scipy.signal as signal
|
||
import matplotlib.pyplot as plt
|
||
import mne
|
||
from mne.preprocessing import ICA
|
||
|
||
# ==========================
|
||
# Config - 预处理参数
|
||
# ==========================
|
||
# 滤波参数
|
||
BANDPASS_LOW = 0.8
|
||
BANDPASS_HIGH = 40.0
|
||
NOTCH_FREQS = [50, 100] # 工频陷波
|
||
|
||
# 幅值过滤阈值 (μV)
|
||
AMPLITUDE_MIN_UV = -200.0
|
||
AMPLITUDE_MAX_UV = 200.0
|
||
|
||
# ICA 参数
|
||
ICA_N_COMPONENTS = 20 # 使用绝对数量而非比例
|
||
ICA_RANDOM_STATE = 97
|
||
ICA_MAX_ITER = 800
|
||
|
||
# 坏段检测阈值 (μV)
|
||
BAD_SEGMENT_THRESHOLD_UV = 350.0
|
||
|
||
# 默认采样率
|
||
DEFAULT_FS = 250.0
|
||
|
||
# 画图参数
|
||
EEG_PLOT_SECONDS = 10
|
||
PSD_FMIN, PSD_FMAX = 0.8, 45.0
|
||
FIXED_EEG_IDXS = [23, 47, 39, 6, 2, 21, 35, 57] # 0-based index
|
||
FIXED_EEG_LABELS = ["C5", "O1", "TP7", "FPZ", "PO6", "P4", "AF7", "AF3"]
|
||
|
||
# 频段定义
|
||
BANDS_METRICS = {
|
||
"Delta": (1.0, 4.0),
|
||
"Theta": (4.0, 8.0),
|
||
"Alpha": (8.0, 13.0),
|
||
"Beta": (13.0, 30.0),
|
||
}
|
||
TOTAL_POWER_BAND = (1.0, 50.0)
|
||
|
||
BANDS_TOPOMAP = {
|
||
"delta": (1.0, 4.0),
|
||
"theta": (4.0, 8.0),
|
||
"alpha": (8.0, 13.0),
|
||
"beta": (13.0, 30.0),
|
||
"broad": (1.0, 30.0),
|
||
}
|
||
|
||
# PSD 参数
|
||
PSD_NPERSEG = 1024 # FFT 窗口大小，越大频率分辨率越高
|
||
|
||
EPS = 1e-12
|
||
|
||
# 脑地形图颜色范围参数
|
||
# 设置为 None 表示自动范围，设置为 (min, max) 固定范围
|
||
TOPOMAP_VMIN = None # 例如: -1.0
|
||
TOPOMAP_VMAX = None # 例如: 1.0
|
||
# 或者使用对称范围（相对于均值的倍数）
|
||
TOPOMAP_SYM_SCALE = 1.5 # 颜色范围 = 均值 ± std * SYM_SCALE
|
||
|
||
# 脑地形图圆形大小参数 (0.08 - 0.15 范围)
|
||
# 数值越小圆形越小，越大圆形越大
|
||
TOPOMAP_SPHERE_RADIUS = 0.12
|
||
|
||
# 边界处理参数
|
||
# 滤波前 padding 秒数，用于消除边界振铃效应
|
||
FILTER_PAD_SEC = 1.0
|
||
|
||
|
||
# ==========================
|
||
# 数据文件读取
|
||
# ==========================
|
||
def load_data_file(file_path: str) -> tuple:
|
||
"""根据文件扩展名读取数据，返回 MNE Raw 对象"""
|
||
ext = os.path.splitext(file_path)[1].lower()
|
||
|
||
if ext == ".bdf":
|
||
return load_bdf_file(file_path)
|
||
elif ext == ".mat":
|
||
return load_mat_file(file_path)
|
||
else:
|
||
raise ValueError(f"不支持的文件格式: {ext}")
|
||
|
||
|
||
def load_bdf_file(bdf_path: str) -> tuple:
|
||
"""读取 .bdf 格式文件，返回 MNE Raw 对象"""
|
||
print(f"[INFO] Reading BDF file: {bdf_path}")
|
||
raw = mne.io.read_raw_bdf(bdf_path, preload=True, verbose=False)
|
||
|
||
try:
|
||
raw.set_montage("standard_1020", on_missing="ignore")
|
||
except Exception as e:
|
||
print(f"[WARN] Failed to set standard_1020 montage: {e}")
|
||
|
||
sfreq = raw.info['sfreq']
|
||
ch_names = raw.ch_names
|
||
n_channels = len(ch_names)
|
||
duration = raw.times[-1] - raw.times[0]
|
||
|
||
print(f"[INFO] Channels: {n_channels}, Duration: {duration:.2f}s, SFreq: {sfreq:.2f}Hz")
|
||
|
||
return raw, sfreq, ch_names
|
||
|
||
|
||
def load_mat_file(mat_path: str) -> tuple:
|
||
"""读取 .mat 格式文件，返回 MNE Raw 对象"""
|
||
print(f"[INFO] Reading MAT file: {mat_path}")
|
||
import scipy.io
|
||
|
||
mat = scipy.io.loadmat(mat_path)
|
||
eeg = mat['eeg'][0, 0]
|
||
|
||
# 提取数据
|
||
data = eeg['data'] # (T, C)
|
||
if data.shape[0] < data.shape[1]:
|
||
data = data.T # 确保是 (T, C)
|
||
data = data.astype(np.float64) # 确保是 float
|
||
|
||
# 提取采样率
|
||
sfreq = float(eeg['sample_rate'][0, 0])
|
||
|
||
# 提取通道名称
|
||
ch_names_raw = eeg['electrode_name']
|
||
if ch_names_raw.ndim == 2:
|
||
ch_names = [str(ch[0]) if isinstance(ch[0], np.bytes_) else str(ch[0]) for ch in ch_names_raw[0]]
|
||
else:
|
||
ch_names = [f"EEG{i+1}" for i in range(data.shape[1])]
|
||
|
||
n_channels = data.shape[1]
|
||
n_samples = data.shape[0]
|
||
duration = n_samples / sfreq
|
||
|
||
print(f"[INFO] Channels: {n_channels}, Duration: {duration:.2f}s, SFreq: {sfreq:.2f}Hz")
|
||
|
||
# 创建 MNE Raw 对象
|
||
info = mne.create_info(ch_names=ch_names, sfreq=sfreq, ch_types=["eeg"] * n_channels)
|
||
raw = mne.io.RawArray(data.T, info, verbose=False) # (T, C) -> (C, T)
|
||
|
||
# 尝试设置通道位置
|
||
try:
|
||
electrode_xyz = eeg['electrode_xyz'] # (64, 3)
|
||
if electrode_xyz.shape[0] == n_channels:
|
||
ch_pos = {}
|
||
for i, name in enumerate(ch_names):
|
||
ch_pos[name] = electrode_xyz[i] / 1000.0 # 转换为米
|
||
montage = mne.channels.make_dig_montage(ch_pos=ch_pos, coord_frame='head')
|
||
info.set_montage(montage)
|
||
print("[INFO] Applied electrode positions from mat file")
|
||
else:
|
||
raw.set_montage("standard_1020", on_missing="ignore")
|
||
except Exception as e:
|
||
print(f"[WARN] Failed to set montage from mat file: {e}")
|
||
try:
|
||
raw.set_montage("standard_1020", on_missing="ignore")
|
||
except:
|
||
pass
|
||
|
||
return raw, sfreq, ch_names
|
||
|
||
|
||
# ==========================
|
||
# 坏通道检测
|
||
# ==========================
|
||
def detect_bad_channels(raw: mne.io.RawArray, z_thresh: float = 3.0) -> list:
|
||
"""检测坏通道：全零/常数通道 + MAD z-score 离群通道"""
|
||
data = raw.get_data()
|
||
ch_names = raw.ch_names
|
||
bad_chs = []
|
||
|
||
ptp = np.ptp(data, axis=1)
|
||
std = np.std(data, axis=1)
|
||
|
||
for i, (p, s) in enumerate(zip(ptp, std)):
|
||
if p < 1e-12 or s < 1e-12:
|
||
bad_chs.append(ch_names[i])
|
||
|
||
valid_mask = np.array([ch not in bad_chs for ch in ch_names])
|
||
if valid_mask.sum() > 2:
|
||
valid_ptp = ptp[valid_mask]
|
||
med = np.median(valid_ptp)
|
||
mad = np.median(np.abs(valid_ptp - med)) + 1e-30
|
||
z = np.abs(ptp - med) / (mad * 1.4826)
|
||
|
||
for i, zv in enumerate(z):
|
||
if zv > z_thresh and ch_names[i] not in bad_chs:
|
||
bad_chs.append(ch_names[i])
|
||
|
||
if bad_chs:
|
||
print(f"[INFO] Bad channels detected: {bad_chs}")
|
||
else:
|
||
print("[INFO] No bad channels detected")
|
||
|
||
return bad_chs
|
||
|
||
|
||
# ==========================
|
||
# 坏段标注
|
||
# ==========================
|
||
def annotate_bad_segments(raw: mne.io.RawArray, peak_to_peak_uv: float = 350.0):
|
||
"""简单坏段检测：按1秒窗口计算峰峰值，超过阈值标为 bad"""
|
||
peak_to_peak_v = peak_to_peak_uv * 1e-6
|
||
win = int(raw.info["sfreq"] * 1.0)
|
||
step = int(raw.info["sfreq"] * 0.5)
|
||
data = raw.get_data()
|
||
n_times = data.shape[1]
|
||
|
||
onsets = []
|
||
durations = []
|
||
|
||
for start in range(0, n_times - win, step):
|
||
seg = data[:, start:start + win]
|
||
ptp = np.ptp(seg, axis=1)
|
||
if np.any(ptp > peak_to_peak_v):
|
||
onsets.append(start / raw.info["sfreq"])
|
||
durations.append(win / raw.info["sfreq"])
|
||
|
||
if len(onsets) > 0:
|
||
ann = mne.Annotations(onset=onsets, duration=durations, description=["BAD_SEG"] * len(onsets))
|
||
raw.set_annotations(ann)
|
||
print(f"[INFO] Annotated {len(onsets)} bad segments")
|
||
|
||
|
||
# ==========================
|
||
# 核心预处理函数
|
||
# ==========================
|
||
def preprocess_bdf(raw: mne.io.RawArray) -> mne.io.RawArray:
|
||
"""BDF 数据预处理流程"""
|
||
print("[INFO] Starting preprocessing pipeline...")
|
||
|
||
# 1) 裁剪首尾 2s
|
||
crop_sec = 2.0
|
||
t_start = crop_sec
|
||
t_end = raw.times[-1] - crop_sec
|
||
if t_end > t_start:
|
||
raw = raw.crop(tmin=t_start, tmax=t_end)
|
||
print(f"[INFO] Cropped: removed first/last {crop_sec}s")
|
||
|
||
# 2) 去直流偏置
|
||
data = raw.get_data()
|
||
data -= data.mean(axis=1, keepdims=True)
|
||
raw._data = data
|
||
print("[INFO] Removed DC offset")
|
||
|
||
# 3) 坏通道检测与插值
|
||
bad_chs = detect_bad_channels(raw)
|
||
if bad_chs:
|
||
raw.info["bads"] = bad_chs
|
||
try:
|
||
raw_tmp = raw.copy()
|
||
raw_tmp.set_montage(raw.get_montage(), on_missing="ignore")
|
||
raw_tmp.interpolate_bads(reset_bads=True, verbose=False)
|
||
raw = raw_tmp
|
||
print(f"[INFO] Bad channels interpolated: {bad_chs}")
|
||
except Exception as e:
|
||
print(f"[WARN] Bad channel interpolation failed: {e}")
|
||
raw.info["bads"] = []
|
||
|
||
# 4) 50Hz 陷波滤波
|
||
print(f"[INFO] Applying notch filter: {NOTCH_FREQS}Hz")
|
||
raw.notch_filter(NOTCH_FREQS, fir_design="firwin", verbose=False)
|
||
|
||
# 5) 0.8-40Hz 带通滤波 (使用 padding 消除边界振铃)
|
||
print(f"[INFO] Applying bandpass filter: {BANDPASS_LOW}-{BANDPASS_HIGH}Hz (with {FILTER_PAD_SEC}s padding)")
|
||
pad_sec = FILTER_PAD_SEC
|
||
raw_length = raw.times[-1]
|
||
pad_start = max(0, pad_sec)
|
||
pad_end = max(0, pad_sec)
|
||
|
||
if raw_length > pad_start + pad_end + 1.0:
|
||
raw.filter(BANDPASS_LOW, BANDPASS_HIGH, fir_design="firwin",
|
||
pad='reflect', verbose=False)
|
||
raw = raw.crop(tmin=pad_sec, tmax=raw_length - pad_sec)
|
||
print(f"[INFO] Removed {pad_sec}s padding from each side after filtering")
|
||
else:
|
||
raw.filter(BANDPASS_LOW, BANDPASS_HIGH, fir_design="firwin", verbose=False)
|
||
print(f"[WARN] Data too short ({raw_length:.1f}s), skipping padding")
|
||
|
||
# 6) 幅值过滤
|
||
print(f"[INFO] Applying amplitude filter: [{AMPLITUDE_MIN_UV}, {AMPLITUDE_MAX_UV}] μV")
|
||
amplitude_thresh_v = AMPLITUDE_MAX_UV * 1e-6
|
||
d = raw.get_data()
|
||
mask = np.abs(d) > amplitude_thresh_v
|
||
n_clipped = int(mask.sum())
|
||
if n_clipped > 0:
|
||
d[mask] = 0.0
|
||
raw._data = d
|
||
print(f"[INFO] Amplitude clipping: {n_clipped} samples exceeded ±200μV, set to 0")
|
||
|
||
# 7) 坏段标注
|
||
annotate_bad_segments(raw, peak_to_peak_uv=BAD_SEGMENT_THRESHOLD_UV)
|
||
|
||
# 8) ICA 去伪迹
|
||
print("[INFO] Running ICA for artifact removal...")
|
||
ica = ICA(n_components=ICA_N_COMPONENTS, random_state=ICA_RANDOM_STATE,
|
||
max_iter=ICA_MAX_ITER, method="fastica")
|
||
ica.fit(raw, reject_by_annotation=True, verbose=False)
|
||
|
||
try:
|
||
eog_inds, _ = ica.find_bads_eog(raw, verbose=False)
|
||
if eog_inds:
|
||
ica.exclude.extend(eog_inds)
|
||
print(f"[INFO] ICA exclude EOG components: {eog_inds}")
|
||
except Exception as e:
|
||
print(f"[WARN] ICA EOG detection skipped: {e}")
|
||
|
||
raw_clean = ica.apply(raw.copy(), verbose=False)
|
||
|
||
# 9) ICA 后再次去直流
|
||
d = raw_clean.get_data()
|
||
d -= d.mean(axis=1, keepdims=True)
|
||
raw_clean._data = d
|
||
|
||
print("[INFO] Preprocessing completed")
|
||
return raw_clean
|
||
|
||
|
||
# ==========================
|
||
# 输出目录管理
|
||
# ==========================
|
||
def ensure_outdir(out_root: str) -> str:
|
||
"""确保输出目录存在，并清空旧文件（保留 ResultData.txt）"""
|
||
if os.path.exists(out_root):
|
||
for filename in os.listdir(out_root):
|
||
if filename == "ResultData.txt":
|
||
continue
|
||
file_path = os.path.join(out_root, filename)
|
||
try:
|
||
if os.path.isfile(file_path) or os.path.islink(file_path):
|
||
os.unlink(file_path)
|
||
elif os.path.isdir(file_path):
|
||
shutil.rmtree(file_path)
|
||
except Exception as e:
|
||
print(f"[WARN] Failed to delete {file_path}: {e}")
|
||
else:
|
||
os.makedirs(out_root, exist_ok=True)
|
||
return out_root
|
||
|
||
|
||
# ==========================
|
||
# 通道分区
|
||
# ==========================
|
||
def _norm_name(s: str) -> str:
|
||
return str(s).strip().upper().replace(" ", "")
|
||
|
||
|
||
def build_channel_index_map(ch_names, n_channels: int):
|
||
if not ch_names or len(ch_names) != n_channels:
|
||
return {}
|
||
return {_norm_name(nm): i for i, nm in enumerate(ch_names)}
|
||
|
||
|
||
def pick_indices_by_names(name_to_idx, names):
|
||
idx = []
|
||
for n in names:
|
||
nn = _norm_name(n)
|
||
if nn in name_to_idx:
|
||
idx.append(name_to_idx[nn])
|
||
return sorted(list(set(idx)))
|
||
|
||
|
||
def _fallback_region_indices(n_channels: int):
|
||
a = int(n_channels * 0.33)
|
||
b = int(n_channels * 0.66)
|
||
return (
|
||
list(range(0, a)), # frontal
|
||
list(range(a, b)), # central
|
||
list(range(b, n_channels)), # parietal
|
||
list(range(0, max(2, a // 2))), # prefrontal
|
||
list(range(b, n_channels)), # posterior
|
||
[i for i in range(n_channels) if i % 2 == 0], # left
|
||
[i for i in range(n_channels) if i % 2 == 1], # right
|
||
)
|
||
|
||
|
||
def get_region_indices(name_to_idx, n_channels: int):
|
||
if not name_to_idx:
|
||
return _fallback_region_indices(n_channels)
|
||
|
||
central_names = ["CZ","C1","C2","C3","C4","C5","C6","CP1","CP2","CP3","CP4","CP5","CP6","FC1","FC2","FC3","FC4","FC5","FC6"]
|
||
frontal_names = ["FZ","F1","F2","F3","F4","F5","F6","F7","F8","AF3","AF4","AF7","AF8","FPZ","FP1","FP2","FCZ"]
|
||
parietal_names = ["PZ","P1","P2","P3","P4","P5","P6","POZ","PO3","PO4","PO5","PO6","PO7","PO8","CPZ"]
|
||
prefrontal_names = ["FP1","FP2","FPZ","AF3","AF4","AF7","AF8"]
|
||
posterior_names = ["O1","O2","OZ","PO7","PO8","PO3","PO4","PZ","P3","P4","P1","P2"]
|
||
left_names = ["FP1","AF3","AF7","F3","F5","F7"]
|
||
right_names = ["FP2","AF4","AF8","F4","F6","F8"]
|
||
|
||
central = pick_indices_by_names(name_to_idx, central_names)
|
||
frontal = pick_indices_by_names(name_to_idx, frontal_names)
|
||
parietal = pick_indices_by_names(name_to_idx, parietal_names)
|
||
prefrontal = pick_indices_by_names(name_to_idx, prefrontal_names)
|
||
posterior = pick_indices_by_names(name_to_idx, posterior_names)
|
||
left = pick_indices_by_names(name_to_idx, left_names)
|
||
right = pick_indices_by_names(name_to_idx, right_names)
|
||
|
||
if not (central and frontal and parietal and prefrontal and posterior):
|
||
fb = _fallback_region_indices(n_channels)
|
||
frontal2, central2, parietal2, prefrontal2, posterior2, left2, right2 = fb
|
||
frontal = frontal if frontal else frontal2
|
||
central = central if central else central2
|
||
parietal = parietal if parietal else parietal2
|
||
prefrontal = prefrontal if prefrontal else prefrontal2
|
||
posterior = posterior if posterior else posterior2
|
||
left = left if left else left2
|
||
right = right if right else right2
|
||
|
||
return frontal, central, parietal, prefrontal, posterior, left, right
|
||
|
||
|
||
# ==========================
|
||
# PSD 和频段功率计算
|
||
# ==========================
|
||
def welch_psd(eeg_tc, fs):
|
||
"""计算 PSD"""
|
||
nperseg = min(PSD_NPERSEG, eeg_tc.shape[0])
|
||
noverlap = int(nperseg * 0.75)
|
||
freqs, pxx = signal.welch(
|
||
eeg_tc, fs=fs, nperseg=nperseg, noverlap=noverlap,
|
||
axis=0, scaling="density"
|
||
)
|
||
return freqs, pxx
|
||
|
||
|
||
def band_power_from_psd(freqs, pxx_fc, band):
|
||
lo, hi = band
|
||
m = (freqs >= lo) & (freqs < hi)
|
||
if not np.any(m):
|
||
return np.zeros((pxx_fc.shape[1],), dtype=np.float32)
|
||
|
||
return np.trapz(pxx_fc[m, :], freqs[m], axis=0).astype(np.float32)
|
||
|
||
|
||
def region_mean_power(freqs, pxx_fc, idx, band) -> float:
|
||
if not idx:
|
||
return 0.0
|
||
pw = band_power_from_psd(freqs, pxx_fc, band)
|
||
return float(np.mean(pw[idx]))
|
||
|
||
|
||
def compute_iaf(freqs, pxx_fc, posterior_idx):
|
||
lo, hi = BANDS_METRICS["Alpha"]
|
||
m = (freqs >= lo) & (freqs <= hi)
|
||
if not np.any(m) or not posterior_idx:
|
||
return 0.0
|
||
spec = np.mean(pxx_fc[:, posterior_idx], axis=1)
|
||
sub = spec[m]
|
||
fsub = freqs[m]
|
||
return float(fsub[int(np.argmax(sub))])
|
||
|
||
|
||
# ==========================
|
||
# 画图函数
|
||
# ==========================
|
||
def plot_eeg_waveforms(data_uv_tc, fs, ch_names, out_dir, seconds=10, t_start_sec=30.0):
|
||
"""画 EEG 波形图（固定通道）"""
|
||
T, C = data_uv_tc.shape
|
||
|
||
start_sample = int(t_start_sec * fs)
|
||
end_sample = int(min(T, start_sample + seconds * fs))
|
||
|
||
if start_sample >= T:
|
||
start_sample = max(0, T - int(seconds * fs))
|
||
end_sample = T
|
||
print(f"[WARN] t_start_sec={t_start_sec}s exceeds data, using last {seconds}s")
|
||
|
||
seg_samples = end_sample - start_sample
|
||
x = np.arange(seg_samples) / fs + t_start_sec
|
||
|
||
# 过滤有效索引
|
||
idxs = [i for i in FIXED_EEG_IDXS if 0 <= i < C]
|
||
if len(idxs) < len(FIXED_EEG_IDXS):
|
||
missing = [i for i in FIXED_EEG_IDXS if not (0 <= i < C)]
|
||
print(f"[WARN] Some indices out of range (C={C}): {missing}")
|
||
|
||
if len(idxs) == 0:
|
||
raise RuntimeError(f"No valid indices for data (C={C})")
|
||
|
||
picked_names = []
|
||
for idx in idxs:
|
||
pos = FIXED_EEG_IDXS.index(idx)
|
||
std_label = FIXED_EEG_LABELS[pos] if pos < len(FIXED_EEG_LABELS) else f"CH{idx}"
|
||
if ch_names and idx < len(ch_names):
|
||
picked_names.append(std_label)
|
||
else:
|
||
picked_names.append(std_label)
|
||
|
||
fig_h = 1.4 * len(idxs) + 1
|
||
fig, axes = plt.subplots(len(idxs), 1, figsize=(10, fig_h), sharex=True)
|
||
if len(idxs) == 1:
|
||
axes = [axes]
|
||
|
||
seg = data_uv_tc[start_sample:end_sample, idxs].T
|
||
lo = float(np.percentile(seg, 1))
|
||
hi = float(np.percentile(seg, 99))
|
||
m = max(abs(lo), abs(hi), 50.0)
|
||
|
||
for ax, ch_idx, nm in zip(axes, idxs, picked_names):
|
||
y = data_uv_tc[start_sample:end_sample, ch_idx]
|
||
ax.plot(x, y, linewidth=1.2)
|
||
ax.set_ylabel("uV")
|
||
ax.set_title(nm, loc="left", fontsize=10)
|
||
ax.grid(True, alpha=0.3)
|
||
ax.set_ylim(-m, m)
|
||
|
||
axes[-1].set_xlabel("Time (s)")
|
||
plt.tight_layout()
|
||
|
||
out_path = os.path.join(out_dir, "EEG.png")
|
||
plt.savefig(out_path, dpi=200)
|
||
plt.close(fig)
|
||
print(f"[OK] EEG waveform saved: {out_path}")
|
||
|
||
|
||
def plot_psd(eeg_uV_tc, fs, ch_names, out_dir):
|
||
"""画 PSD 图"""
|
||
C = eeg_uV_tc.shape[1]
|
||
chosen_idx = []
|
||
|
||
if ch_names:
|
||
mp = {n.upper(): i for i, n in enumerate(ch_names)}
|
||
for p in ["C3", "C4", "CZ"]:
|
||
if p in mp:
|
||
chosen_idx.append(mp[p])
|
||
if len(chosen_idx) < 3:
|
||
stds = [(i, float(np.std(eeg_uV_tc[:, i]))) for i in range(C)]
|
||
stds.sort(key=lambda x: x[1], reverse=True)
|
||
for i, _ in stds:
|
||
if i not in chosen_idx:
|
||
chosen_idx.append(i)
|
||
if len(chosen_idx) == 3:
|
||
break
|
||
chosen_name = [ch_names[i] for i in chosen_idx]
|
||
else:
|
||
stds = [(i, float(np.std(eeg_uV_tc[:, i]))) for i in range(C)]
|
||
stds.sort(key=lambda x: x[1], reverse=True)
|
||
chosen_idx = [i for i, _ in stds[:3]]
|
||
chosen_name = [f"CH{i}" for i in chosen_idx]
|
||
|
||
# 增大 nperseg 提高频率分辨率
|
||
nperseg = min(PSD_NPERSEG, eeg_uV_tc.shape[0])
|
||
noverlap = int(nperseg * 0.75)
|
||
|
||
fig = plt.figure(figsize=(7.5, 4.8))
|
||
for idx, nm in zip(chosen_idx, chosen_name):
|
||
f, pxx = signal.welch(eeg_uV_tc[:, idx], fs=fs, nperseg=nperseg, noverlap=noverlap)
|
||
mask = (f >= PSD_FMIN) & (f <= PSD_FMAX)
|
||
p_db = 10 * np.log10(pxx[mask] + 1e-20)
|
||
plt.plot(f[mask], p_db, linewidth=1.8, label=nm)
|
||
|
||
plt.xlabel("Hz")
|
||
plt.ylabel("Power (dB)")
|
||
plt.title("PSD")
|
||
plt.grid(True, alpha=0.3)
|
||
plt.legend()
|
||
plt.tight_layout()
|
||
|
||
out_path = os.path.join(out_dir, "psd.png")
|
||
plt.savefig(out_path, dpi=200)
|
||
plt.close(fig)
|
||
print(f"[OK] PSD saved: {out_path}")
|
||
|
||
|
||
def _get_standard_1020_channel_indices(raw):
|
||
"""获取符合 standard_1020 montage 的通道索引和名称"""
|
||
try:
|
||
standard_montage = mne.channels.make_standard_montage("standard_1020")
|
||
standard_names_upper = {ch.upper() for ch in standard_montage.ch_names}
|
||
standard_name_map = {ch.upper(): ch for ch in standard_montage.ch_names}
|
||
data_ch_names = raw.ch_names
|
||
exclude_names = {"A1", "A2", "M1", "M2", "LE", "RE", "LM", "RM"}
|
||
|
||
valid_indices = []
|
||
valid_names = []
|
||
for i, name in enumerate(data_ch_names):
|
||
name_upper = name.upper()
|
||
if name_upper in standard_names_upper and name_upper not in exclude_names:
|
||
valid_indices.append(i)
|
||
valid_names.append(standard_name_map[name_upper])
|
||
|
||
print(f"[INFO] Found {len(valid_indices)}/{len(data_ch_names)} channels matching standard_1020")
|
||
return valid_indices, valid_names
|
||
except Exception as e:
|
||
print(f"[WARN] Failed to get standard_1020 channels: {e}")
|
||
return None, None
|
||
|
||
|
||
def compute_band_powers_for_topomap(raw, bands):
|
||
"""计算各频段功率，只使用 standard_1020 中有位置的通道"""
|
||
# 获取 standard_1020 montage 和位置信息
|
||
standard_montage = mne.channels.make_standard_montage("standard_1020")
|
||
std_names_upper = {ch.upper() for ch in standard_montage.ch_names}
|
||
ch_pos_map = standard_montage.get_positions()['ch_pos']
|
||
|
||
data_ch_names = raw.ch_names
|
||
exclude = {"A1", "A2", "M1", "M2", "LE", "RE", "LM", "RM"}
|
||
|
||
# 只保留有位置信息的通道
|
||
valid_indices = []
|
||
valid_names = []
|
||
for i, name in enumerate(data_ch_names):
|
||
name_upper = name.upper()
|
||
if name_upper in std_names_upper and name_upper not in exclude:
|
||
if name_upper in ch_pos_map: # 必须有位置
|
||
valid_indices.append(i)
|
||
valid_names.append(name)
|
||
|
||
if len(valid_indices) < 8:
|
||
return None
|
||
|
||
data = raw.get_data()
|
||
data_standard = data[valid_indices, :]
|
||
|
||
fs = raw.info["sfreq"]
|
||
n_fft = min(PSD_NPERSEG, data_standard.shape[1])
|
||
n_overlap = int(n_fft * 0.75)
|
||
|
||
psds, freqs = mne.time_frequency.psd_array_welch(
|
||
data_standard, sfreq=fs,
|
||
fmin=min(v[0] for v in bands.values()),
|
||
fmax=max(v[1] for v in bands.values()),
|
||
n_fft=n_fft, n_overlap=n_overlap,
|
||
average="mean", verbose=False
|
||
)
|
||
|
||
out = {"_valid_names": valid_names}
|
||
print(f"[DEBUG] PSD: fs={fs}Hz, n_fft={n_fft}, freq_res={fs/n_fft:.3f}Hz/bin")
|
||
for k, (fmin, fmax) in bands.items():
|
||
idx = np.where((freqs >= fmin) & (freqs < fmax))[0]
|
||
if len(idx) == 0:
|
||
out[k] = np.zeros(len(valid_indices), dtype=np.float32)
|
||
print(f"[DEBUG] {k.upper()}: NO freq bins in [{fmin}-{fmax}]Hz")
|
||
continue
|
||
print(f"[DEBUG] {k.upper()}: freq bins {freqs[idx[0]]:.2f}-{freqs[idx[-1]]:.2f}Hz (bins {idx[0]}-{idx[-1]}, count={len(idx)})")
|
||
# 使用线性功率值 (V^2 -> uV^2: * 1e12)
|
||
bp = np.trapz(psds[:, idx], freqs[idx], axis=1) * 1e12
|
||
out[k] = bp
|
||
print(f"[DEBUG] {k.upper()}: power range [{bp.min():.4f}, {bp.max():.4f}] uV^2, mean={bp.mean():.4f}")
|
||
|
||
print(f"[INFO] Band powers computed for {len(valid_names)} channels with positions")
|
||
return out
|
||
|
||
|
||
def _create_topomap_raw(ch_names):
|
||
"""创建只有 standard_1020 通道位置信息的临时 Raw 对象"""
|
||
standard_montage = mne.channels.make_standard_montage("standard_1020")
|
||
ch_pos_map = standard_montage.get_positions()['ch_pos']
|
||
|
||
valid_ch_names = []
|
||
valid_positions = []
|
||
for name in ch_names:
|
||
name_upper = name.upper()
|
||
if name_upper in ch_pos_map:
|
||
valid_ch_names.append(name)
|
||
valid_positions.append(ch_pos_map[name_upper])
|
||
|
||
if len(valid_ch_names) < 8:
|
||
return None
|
||
|
||
ch_pos = {name: pos for name, pos in zip(valid_ch_names, valid_positions)}
|
||
montage = mne.channels.make_dig_montage(ch_pos=ch_pos, coord_frame='head')
|
||
|
||
info = mne.create_info(ch_names=valid_ch_names, sfreq=250.0, ch_types=["eeg"] * len(valid_ch_names))
|
||
info.set_montage(montage)
|
||
|
||
dummy_data = np.zeros((len(valid_ch_names), 1))
|
||
return mne.io.RawArray(dummy_data, info, verbose=False)
|
||
|
||
|
||
def plot_average_topomap(band_values, out_dir):
|
||
"""绘制平均拓扑图"""
|
||
valid_names = band_values.get("_valid_names", [])
|
||
if not valid_names:
|
||
return
|
||
|
||
values = band_values["broad"]
|
||
temp_raw = _create_topomap_raw(valid_names)
|
||
if temp_raw is None:
|
||
return
|
||
|
||
vmin, vmax = _compute_topomap_vlim([values])
|
||
|
||
fig, ax = plt.subplots(1, 1, figsize=(6.5, 4.6))
|
||
im, _ = mne.viz.plot_topomap(
|
||
values, temp_raw.info, axes=ax, show=False, contours=0,
|
||
sphere=(0, 0, 0, TOPOMAP_SPHERE_RADIUS), extrapolate='head', border='mean',
|
||
cmap='turbo'
|
||
)
|
||
im.set_clim(vmin=vmin, vmax=vmax)
|
||
ax.set_title("0.8-30 Hz", fontsize=12)
|
||
plt.colorbar(im, ax=ax, shrink=0.85)
|
||
plt.tight_layout()
|
||
plt.savefig(os.path.join(out_dir, "average_topomap.png"), dpi=200)
|
||
plt.close(fig)
|
||
print(f"[OK] average_topomap saved")
|
||
|
||
|
||
def plot_band_topomaps(band_values, out_dir):
|
||
"""绘制分频段拓扑图"""
|
||
valid_names = band_values.get("_valid_names", [])
|
||
if not valid_names:
|
||
return
|
||
|
||
order = [
|
||
("delta", "δ (1-4Hz)"),
|
||
("theta", "θ (4-8Hz)"),
|
||
("alpha", "α (8-13Hz)"),
|
||
("beta", "β (13-30Hz)"),
|
||
("broad", "1-30 Hz"),
|
||
]
|
||
|
||
temp_raw = _create_topomap_raw(valid_names)
|
||
if temp_raw is None:
|
||
return
|
||
|
||
all_values = [band_values[k] for k, _ in order]
|
||
vmin, vmax = _compute_topomap_vlim(all_values)
|
||
|
||
fig, axes = plt.subplots(1, 5, figsize=(16, 4.2))
|
||
ims = []
|
||
for ax, (k, title) in zip(axes, order):
|
||
im, _ = mne.viz.plot_topomap(
|
||
band_values[k], temp_raw.info, axes=ax, show=False, contours=0,
|
||
sphere=(0, 0, 0, TOPOMAP_SPHERE_RADIUS), extrapolate='head', border='mean',
|
||
cmap='turbo'
|
||
)
|
||
im.set_clim(vmin=vmin, vmax=vmax)
|
||
ax.set_title(title, fontsize=11)
|
||
ims.append(im)
|
||
fig.subplots_adjust(left=0.02, right=0.85, top=0.88, bottom=0.05, wspace=0.35)
|
||
cax = fig.add_axes([0.87, 0.15, 0.015, 0.7])
|
||
fig.colorbar(ims[-1], cax=cax)
|
||
plt.savefig(os.path.join(out_dir, "topomaps.png"), dpi=200)
|
||
plt.close(fig)
|
||
print(f"[OK] topomaps saved")
|
||
|
||
|
||
def _compute_topomap_vlim(values):
|
||
"""计算脑地形图颜色范围"""
|
||
v_all = np.concatenate(values) if isinstance(values, list) else np.array(values)
|
||
if TOPOMAP_VMIN is not None and TOPOMAP_VMAX is not None:
|
||
return TOPOMAP_VMAX - 60, TOPOMAP_VMAX # 保持 50 的范围
|
||
if TOPOMAP_SYM_SCALE is not None and TOPOMAP_SYM_SCALE > 0:
|
||
mean_val = np.mean(v_all)
|
||
std_val = np.std(v_all)
|
||
return mean_val - std_val * TOPOMAP_SYM_SCALE, mean_val + std_val * TOPOMAP_SYM_SCALE
|
||
# 统一 vmax：使用所有频段中的最大值
|
||
# vmin = 0：这样低功率频段会接近 0（白色/冷色），高功率频段突出
|
||
vmax = np.max(v_all)
|
||
vmin = 0
|
||
return vmin, vmax
|
||
|
||
|
||
# ==========================
|
||
# 预测接口
|
||
# ==========================
|
||
def _predict_label_by_model(model_path: str, data_path: str) -> dict:
|
||
"""调用 infer_pth.py 进行预测"""
|
||
try:
|
||
from infer_pth import predict_hc_mdd
|
||
except Exception as e:
|
||
raise RuntimeError(f"无法导入 predict_hc_mdd: {e}")
|
||
|
||
import tempfile
|
||
import scipy.io
|
||
|
||
ext = os.path.splitext(data_path)[1].lower()
|
||
|
||
if ext == ".mat":
|
||
# 直接使用 mat 文件
|
||
result = predict_hc_mdd(os.path.dirname(data_path), model_path)
|
||
elif ext == ".bdf":
|
||
# 转换为 mat 格式
|
||
raw = mne.io.read_raw_bdf(data_path, preload=True, verbose=False)
|
||
data, times = raw[:]
|
||
sfreq = raw.info['sfreq']
|
||
ch_names = raw.ch_names
|
||
|
||
with tempfile.TemporaryDirectory() as temp_dir:
|
||
mat_path = os.path.join(temp_dir, "preprocessed_eeg.mat")
|
||
scipy.io.savemat(mat_path, {
|
||
'eeg': {
|
||
'data': (data * 1e6).T,
|
||
'sample_rate': sfreq,
|
||
'electrode_name': ch_names
|
||
}
|
||
})
|
||
result = predict_hc_mdd(temp_dir, model_path)
|
||
else:
|
||
raise ValueError(f"不支持的文件格式: {ext}")
|
||
|
||
return result
|
||
|
||
|
||
# ==========================
|
||
# 生成 ResultData.txt
|
||
# ==========================
|
||
def compute_and_save_txt(model_path, bdf_path, out_dir, eeg_uV_tc, fs, ch_names):
|
||
"""计算特征指标并保存 ResultData.txt"""
|
||
# 获取预测结果
|
||
pred_result = _predict_label_by_model(model_path, bdf_path)
|
||
pred_label = pred_result.get("pred_label", "UNKNOWN")
|
||
recommend = "是" if pred_label == "MDD" else "否"
|
||
|
||
T, C = eeg_uV_tc.shape
|
||
mp = build_channel_index_map(ch_names, C)
|
||
frontal_idx, central_idx, parietal_idx, prefrontal_idx, posterior_idx, left_idx, right_idx = \
|
||
get_region_indices(mp, C)
|
||
|
||
freqs, pxx = welch_psd(eeg_uV_tc, fs)
|
||
|
||
# 计算各频段功率比
|
||
central_alpha = region_mean_power(freqs, pxx, central_idx, BANDS_METRICS["Alpha"])
|
||
central_beta = region_mean_power(freqs, pxx, central_idx, BANDS_METRICS["Beta"])
|
||
frontal_alpha = region_mean_power(freqs, pxx, frontal_idx, BANDS_METRICS["Alpha"])
|
||
frontal_beta = region_mean_power(freqs, pxx, frontal_idx, BANDS_METRICS["Beta"])
|
||
par_alpha = region_mean_power(freqs, pxx, parietal_idx, BANDS_METRICS["Alpha"])
|
||
par_beta = region_mean_power(freqs, pxx, parietal_idx, BANDS_METRICS["Beta"])
|
||
|
||
central_ab = (central_alpha / (central_beta + EPS)) if central_beta > 0 else 0.0
|
||
frontal_ab = (frontal_alpha / (frontal_beta + EPS)) if frontal_beta > 0 else 0.0
|
||
par_ab = (par_alpha / (par_beta + EPS)) if par_beta > 0 else 0.0
|
||
|
||
central_theta = region_mean_power(freqs, pxx, central_idx, BANDS_METRICS["Theta"])
|
||
par_theta = region_mean_power(freqs, pxx, parietal_idx, BANDS_METRICS["Theta"])
|
||
central_tb = (central_theta / (central_beta + EPS)) if central_beta > 0 else 0.0
|
||
par_tb = (par_theta / (par_beta + EPS)) if par_beta > 0 else 0.0
|
||
|
||
if not left_idx or not right_idx:
|
||
left_idx = [i for i in prefrontal_idx if (i % 2 == 0)]
|
||
right_idx = [i for i in prefrontal_idx if (i % 2 == 1)]
|
||
left_alpha = region_mean_power(freqs, pxx, left_idx, BANDS_METRICS["Alpha"])
|
||
right_alpha = region_mean_power(freqs, pxx, right_idx, BANDS_METRICS["Alpha"])
|
||
prefrontal_alpha_asym = float(np.log(right_alpha + EPS) - np.log(left_alpha + EPS))
|
||
|
||
iaf = compute_iaf(freqs, pxx, posterior_idx)
|
||
|
||
pre_td = region_mean_power(freqs, pxx, prefrontal_idx,
|
||
(BANDS_METRICS["Delta"][0], BANDS_METRICS["Theta"][1]))
|
||
pre_total = region_mean_power(freqs, pxx, prefrontal_idx, TOTAL_POWER_BAND)
|
||
pre_td_rel = (pre_td / (pre_total + EPS)) * 100.0 if pre_total > 0 else 0.0
|
||
|
||
def f1(x): return f"{x:.1f}"
|
||
|
||
txt = (
|
||
f"中央区α/β波比值:{f1(central_ab)}\n"
|
||
f"额区α/β波比值:{f1(frontal_ab)}\n"
|
||
f"顶区α/β波比值:{f1(par_ab)}\n"
|
||
f"中央区θ/β波比值:{f1(central_tb)}\n"
|
||
f"顶区θ/β波比值:{f1(par_tb)}\n"
|
||
f"前额叶α波不对称性:{f1(prefrontal_alpha_asym)}\n"
|
||
f"个体化α峰值频率:{f1(iaf)}\n"
|
||
f"前额叶θ+δ波功率:{f1(pre_td_rel)}\n"
|
||
f"是否推荐治疗:{recommend}\n"
|
||
)
|
||
|
||
out_path = os.path.join(out_dir, "ResultData.txt")
|
||
with open(out_path, "w", encoding="utf-8") as f:
|
||
f.write(txt)
|
||
print(f"[OK] ResultData.txt saved: {out_path}")
|
||
|
||
# 打印预测结果
|
||
print(f"\n========== 预测结果 ==========")
|
||
print(f"预测标签: {pred_label}")
|
||
print(f"p(MDD)均值: {pred_result.get('p_mdd_mean', 'N/A'):.4f}")
|
||
print(f"切片数量: {pred_result.get('n_slices', 'N/A')}")
|
||
print(f"==============================\n")
|
||
|
||
|
||
# ==========================
|
||
# 主函数
|
||
# ==========================
|
||
def run_all(model_path: str, bdf_dir: str, out_root: str, seconds: int = EEG_PLOT_SECONDS):
|
||
"""主流程"""
|
||
if not os.path.exists(bdf_dir):
|
||
raise RuntimeError(f"输入目录不存在: {bdf_dir}")
|
||
|
||
# 支持 .bdf 和 .mat 文件
|
||
data_files = [f for f in os.listdir(bdf_dir) if f.lower().endswith((".bdf", ".mat"))]
|
||
if not data_files:
|
||
raise RuntimeError(f"目录中找不到 .bdf 或 .mat 文件: {bdf_dir}")
|
||
|
||
data_files.sort()
|
||
data_path = os.path.join(bdf_dir, data_files[0])
|
||
print(f"[INFO] Processing file: {data_path}")
|
||
|
||
out_dir = ensure_outdir(out_root)
|
||
print(f"[INFO] Output directory: {out_dir}")
|
||
|
||
raw, sfreq, ch_names = load_data_file(data_path)
|
||
|
||
raw_clean = preprocess_bdf(raw)
|
||
|
||
try:
|
||
raw_clean.set_montage("standard_1020", on_missing="ignore")
|
||
except Exception as e:
|
||
print(f"[WARN] Failed to re-apply montage: {e}")
|
||
|
||
raw_data = raw_clean.get_data()
|
||
eeg_uV_tc = (raw_data * 1e6).T.astype(np.float32)
|
||
print(f"[INFO] Preprocessed EEG shape: {eeg_uV_tc.shape}")
|
||
|
||
print("[INFO] Generating figures...")
|
||
plot_psd(eeg_uV_tc, sfreq, ch_names, out_dir)
|
||
plot_eeg_waveforms(eeg_uV_tc, sfreq, ch_names, out_dir, seconds=seconds)
|
||
|
||
print("[INFO] Generating topomaps...")
|
||
try:
|
||
band_vals = compute_band_powers_for_topomap(raw_clean, BANDS_TOPOMAP)
|
||
if band_vals is not None:
|
||
plot_average_topomap(band_vals, out_dir)
|
||
plot_band_topomaps(band_vals, out_dir)
|
||
except Exception as e:
|
||
print(f"[WARN] Topomap generation failed: {e}")
|
||
|
||
print("[INFO] Running prediction...")
|
||
compute_and_save_txt(model_path, data_path, out_dir, eeg_uV_tc, sfreq, ch_names)
|
||
|
||
print("[DONE] All tasks completed.")
|
||
return out_dir
|
||
|
||
|
||
# ==========================
|
||
# 命令行入口
|
||
# ==========================
|
||
if __name__ == "__main__":
|
||
import multiprocessing
|
||
multiprocessing.freeze_support()
|
||
import argparse
|
||
import sys
|
||
|
||
def get_resource_path(relative_path):
|
||
"""获取资源绝对路径"""
|
||
if getattr(sys, 'frozen', False):
|
||
base_path = os.path.dirname(sys.executable)
|
||
else:
|
||
base_path = os.path.dirname(os.path.abspath(__file__))
|
||
return os.path.join(base_path, relative_path)
|
||
|
||
# 默认路径
|
||
DEFAULT_MODEL = get_resource_path(os.path.join("model", "Model_1.pth"))
|
||
if getattr(sys, 'frozen', False):
|
||
EXE_DIR = os.path.dirname(sys.executable)
|
||
else:
|
||
EXE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||
|
||
DEFAULT_BDF_DIR = os.path.join(EXE_DIR, "raw_data")
|
||
DEFAULT_OUT = os.path.join(EXE_DIR, "out")
|
||
|
||
parser = argparse.ArgumentParser(description="EEG Depression Assessment")
|
||
parser.add_argument("--model_path", type=str, default=DEFAULT_MODEL, help="模型文件路径 (.pth)")
|
||
parser.add_argument("--bdf_dir", type=str, default=DEFAULT_BDF_DIR, help="输入文件夹路径 (包含 .bdf 或 .mat 文件)")
|
||
parser.add_argument("--out_root", type=str, default=DEFAULT_OUT, help="结果输出目录")
|
||
parser.add_argument("--seconds", type=int, default=10, help="画波形图的秒数")
|
||
|
||
args = parser.parse_args()
|
||
|
||
print(f"[*] 运行配置:")
|
||
print(f" - Model : {args.model_path}")
|
||
print(f" - Input : {args.bdf_dir}")
|
||
print(f" - Output: {args.out_root}")
|
||
|
||
if not os.path.exists(args.bdf_dir):
|
||
print(f"[ERROR] 输入目录不存在: {args.bdf_dir}")
|
||
if not os.path.exists(args.model_path):
|
||
print(f"[ERROR] 模型文件不存在: {args.model_path}")
|
||
|
||
run_all(args.model_path, args.bdf_dir, args.out_root, seconds=args.seconds)
|