# -*- coding: utf-8 -*- """ infer_pth.py 用途： - 从一个文件夹中自动读取第一个 .mat EEG 文件（64通道或32通道） - 若为64通道，则按 idx64_to_32 映射选出32通道 - 提取切片特征（DE + PSD(var近似)，不含Asym） - 加载你训练好的 .pth 模型（FusionNet结构） - 输出该受试者的 HC / MDD 判断结果 运行方式（命令行）： python infer_pth_from64_to32.py --eeg_dir "D:\\xxx\\folder" --model_path "C:\\xxx\\model.pth" 也可在其他py里import： from infer_pth_from64_to32 import predict_hc_mdd res = predict_hc_mdd(eeg_dir, model_path) print(res) """ from __future__ import annotations import os import argparse import numpy as np import scipy.io import scipy.signal as signal import torch import torch.nn as nn import torch.nn.functional as F # ========================================================= # 0) 配置区（按需改这里） # ========================================================= # 采样率（必须与训练时一致） SAMPLING_RATE = 250 # 滑窗参数（必须与训练时一致） WINDOW_SIZE = 500 STRIDE = 250 # 频段（必须与训练时一致） BAND_NAMES = ["Delta", "Theta", "Alpha", "Beta", "Gamma"] BANDS = { "Delta": (1, 4), "Theta": (4, 8), "Alpha": (8, 13), "Beta": (13, 30), "Gamma": (30, 50), } # 是否使用扩展特征（DE+PSD） USE_EXTENDED_FEATURES = True # 数值稳定项 EPS = 1e-12 # 设备 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 通道映射 IDX64_TO_32 = [ 23, # C5 47, # O1 39, # TP7 6, # FPZ 2, # PO6 21, # P4 35, # AF7 57, # AF3 1, # FP2 37, # T7 63, # F1 36, # A1 18, # FC4 31, # FC5 14, # FC2 48, # T8 60, # P2 41, # AF8 11, # CP1 0, # FP1 55, # PO7 59, # C1 22, # F5 10, # CP2 16, # C3 61, # P1 27, # CP5 17, # C4 26, # CP6 62, # F2 3, # POZ 13, # PO5 ] # 推理阈值：如果模型checkpoint里有 subject_threshold，会优先用它；否则用这个 DEFAULT_SUBJECT_THRESHOLD = 0.5 # ========================================================= # 1) 模型结构 # ========================================================= class SEBlock(nn.Module): def __init__(self, channels: int, reduction: int = 4) -> None: super().__init__() hidden = max(1, channels // reduction) self.fc1 = nn.Linear(channels, hidden) self.fc2 = nn.Linear(hidden, channels) def forward(self, x: torch.Tensor) -> torch.Tensor: se = F.relu(self.fc1(x)) se = torch.sigmoid(self.fc2(se)) return x * se class ResidualBlock(nn.Module): def __init__(self, in_features: int, out_features: int, dropout: float = 0.3) -> None: super().__init__() self.fc1 = nn.Linear(in_features, out_features) self.bn1 = nn.BatchNorm1d(out_features) self.fc2 = nn.Linear(out_features, out_features) self.bn2 = nn.BatchNorm1d(out_features) self.dropout = nn.Dropout(dropout) self.shortcut = nn.Identity() if in_features != out_features: self.shortcut = nn.Sequential( nn.Linear(in_features, out_features), nn.BatchNorm1d(out_features), ) def forward(self, x: torch.Tensor) -> torch.Tensor: identity = self.shortcut(x) out = F.relu(self.bn1(self.fc1(x))) out = self.dropout(out) out = self.bn2(self.fc2(out)) out = F.relu(out + identity) return out class FusionNet(nn.Module): def __init__(self, num_classes: int = 2, num_eeg_features: int = 320, num_scales: int = 6) -> None: super().__init__() self.input_norm = nn.BatchNorm1d(num_eeg_features) self.block1 = ResidualBlock(num_eeg_features, 512, dropout=0.4) self.block2 = ResidualBlock(512, 256, dropout=0.3) self.block3 = ResidualBlock(256, 128, dropout=0.2) self.attention = SEBlock(128, reduction=4) self.final_fc = nn.Sequential( nn.Linear(128, 64), nn.BatchNorm1d(64), nn.ReLU(), nn.Dropout(0.2), ) self.cls_head = nn.Linear(64, num_classes) # 训练时有回归头也没关系（推理只用cls） self.reg_head = nn.Sequential( nn.Linear(64, 32), nn.ReLU(), nn.Dropout(0.1), nn.Linear(32, num_scales), ) self._init_weights() def _init_weights(self) -> None: for m in self.modules(): if isinstance(m, nn.Linear): nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") if m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.BatchNorm1d): nn.init.constant_(m.weight, 1) nn.init.constant_(m.bias, 0) def forward(self, x: torch.Tensor): x = self.input_norm(x) x = self.block1(x) x = self.block2(x) x = self.block3(x) x = self.attention(x) features = self.final_fc(x) cls_out = self.cls_head(features) reg_out = self.reg_head(features) return cls_out, reg_out # ========================================================= # 2) mat读取 + 通道裁剪 # ========================================================= def _find_first_mat_file(folder: str) -> str: if not os.path.isdir(folder): raise RuntimeError(f"eeg_dir 不是文件夹: {folder}") mats = sorted([f for f in os.listdir(folder) if f.lower().endswith(".mat")]) if not mats: raise RuntimeError(f"文件夹内没有 .mat 文件: {folder}") return os.path.join(folder, mats[0]) import numpy as np import scipy.io def _unwrap_singleton(x): """ 把 (1,1) / (1,) 这种包裹层一直剥掉，直到不是 singleton。 也处理 object array 的情况。 """ while True: if isinstance(x, np.ndarray): if x.dtype == object and x.size == 1: x = x.item() continue if x.size == 1 and x.ndim >= 1: # 例如 (1,1) 或 (1,) 的数值/对象数组 try: x = x.reshape(-1)[0] continue except Exception: pass break return x def _try_get_struct_field(v, field_name="data"): """ 尝试从以下几种结构中提取字段： 1) scipy 读出的 mat_struct（有 _fieldnames） 2) numpy structured/record array（dtype.names） """ # case 1: mat_struct（推荐 loadmat(..., struct_as_record=False, squeeze_me=True)） if hasattr(v, "_fieldnames") and (field_name in getattr(v, "_fieldnames", [])): return getattr(v, field_name) # case 2: structured array if isinstance(v, np.ndarray) and v.dtype.names and (field_name in v.dtype.names): # 常见是 v[field] 仍然是 ndarray / object，需要 unwrap try: return v[field_name] except Exception: return None return None def load_eeg_from_mat_any_channels(mat_path: str) -> np.ndarray: """ 读取 .mat 中 EEG 数据，支持： - 直接二维矩阵 (T,C) 或 (C,T) - struct 里有字段 data 返回统一为 float32 的 (T, C) """ # 用这两个参数会让 struct 更容易处理：字段变成属性，且自动 squeeze mat = scipy.io.loadmat(mat_path, struct_as_record=False, squeeze_me=True) candidates = [] for k, v in mat.items(): if k.startswith("__"): continue # --- 1) 直接二维数值矩阵 --- if isinstance(v, np.ndarray) and v.ndim == 2 and np.issubdtype(v.dtype, np.number): candidates.append((k, v)) continue # --- 2) struct/record：优先提取 data 字段 --- data_field = _try_get_struct_field(v, "data") if data_field is not None: data_field = _unwrap_singleton(data_field) # data_field 可能仍然被 object 包一层 if isinstance(data_field, np.ndarray) and data_field.dtype == object: data_field = _unwrap_singleton(data_field) if isinstance(data_field, np.ndarray) and data_field.ndim == 2: # 只收数值矩阵 if np.issubdtype(data_field.dtype, np.number) or data_field.dtype == object: candidates.append((f"{k}.data", data_field)) continue # --- 3) object array：尝试 item() 解包后再看是不是二维数值矩阵/struct --- if isinstance(v, np.ndarray) and v.dtype == object: vv = _unwrap_singleton(v) # 解包后若是二维数值矩阵 if isinstance(vv, np.ndarray) and vv.ndim == 2 and np.issubdtype(vv.dtype, np.number): candidates.append((k, vv)) continue # 解包后若是 struct，再取 data data2 = _try_get_struct_field(vv, "data") if data2 is not None: data2 = _unwrap_singleton(data2) if isinstance(data2, np.ndarray) and data2.ndim == 2: candidates.append((f"{k}.data", data2)) continue if not candidates: raise RuntimeError(f"mat里没找到可用EEG二维矩阵或struct.data：{mat_path}") # 选一个最像 EEG 的（优先含32/64通道维度的） def score(arr: np.ndarray) -> int: s = 0 if 64 in arr.shape: s += 10 if 32 in arr.shape: s += 9 if 128 in arr.shape: s += 8 if 129 in arr.shape: s += 7 s += int(np.prod(arr.shape) // 100000) # 大一些更像EEG return s candidates.sort(key=lambda kv: score(kv[1]), reverse=True) key, eeg = candidates[0] eeg = _unwrap_singleton(eeg) # 如果还是 object dtype，尽力转成 float if isinstance(eeg, np.ndarray) and eeg.dtype == object: # 有时 object 里其实是数值 eeg = np.array(eeg, dtype=np.float32) else: eeg = np.asarray(eeg, dtype=np.float32) if eeg.ndim != 2: raise RuntimeError(f"解析结果不是二维矩阵: key={key}, shape={eeg.shape}, file={mat_path}") # 统一为 (T, C) # 常见 (C,T) 或 (T,C)，我们用“通道维通常较小”+ “32/64/128/129”判断 if eeg.shape[0] in (32, 64, 128, 129) and eeg.shape[1] not in (32, 64, 128, 129): eeg = eeg.T elif eeg.shape[1] in (32, 64, 128, 129): # 如果第一维也是这些数且更小，可能是(C,T) if eeg.shape[0] in (32, 64, 128, 129) and eeg.shape[0] < eeg.shape[1]: eeg = eeg.T return eeg def ensure_32_channels(eeg: np.ndarray) -> np.ndarray: """ 输入 (T, C)，输出 (T, 32) - 若C=64：按 IDX64_TO_32 选32通道 - 若C=32：直接返回 """ if eeg.ndim != 2: raise RuntimeError(f"EEG必须是二维(T,C)，但得到: {eeg.shape}") C = eeg.shape[1] if C == 64: idx = np.asarray(IDX64_TO_32, dtype=np.int64) if idx.min() < 0 or idx.max() >= 64: raise RuntimeError(f"IDX64_TO_32 越界：min={idx.min()}, max={idx.max()} (要求0~63)") return eeg[:, idx] if C == 32: return eeg raise RuntimeError(f"不支持的通道数C={C}，当前只支持 64->32 或 32 直推。") # ========================================================= # 3) 特征提取（DE + PSD(var近似)） # ========================================================= class FeatureExtractor32: """ 只针对32通道，输出维度： - USE_EXTENDED_FEATURES=True：DE(32*5) + PSD(32*5) = 320 - 否则：DE(32*5) = 160 """ def __init__( self, fs: int = SAMPLING_RATE, window_size: int = WINDOW_SIZE, stride: int = STRIDE, filter_order: int = 4, zero_phase: bool = False, ) -> None: self.fs = fs self.window_size = window_size self.stride = stride self.filter_order = filter_order self.zero_phase = zero_phase self._sos = {} for bn in BAND_NAMES: low, high = BANDS[bn] self._sos[bn] = signal.butter( self.filter_order, [low, high], btype="band", fs=self.fs, output="sos" ) def _filter_bands(self, eeg: np.ndarray) -> dict[str, np.ndarray]: out = {} for bn in BAND_NAMES: sos = self._sos[bn] if self.zero_phase: out[bn] = signal.sosfiltfilt(sos, eeg, axis=0).astype(np.float32) else: out[bn] = signal.sosfilt(sos, eeg, axis=0).astype(np.float32) return out def extract(self, eeg32: np.ndarray) -> np.ndarray: """ eeg32: (T, 32) return: feats (N_slices, feat_dim) """ if eeg32.ndim != 2 or eeg32.shape[1] != 32: raise RuntimeError(f"extract需要 (T,32)，得到 {eeg32.shape}") bands_data = self._filter_bands(eeg32) T = eeg32.shape[0] feats = [] for start in range(0, T - self.window_size, self.stride): end = start + self.window_size de_list = [] psd_list = [] for bn in BAND_NAMES: seg = bands_data[bn][start:end, :] # (W, 32) var = np.var(seg, axis=0, ddof=1) # (32,) # DE de = 0.5 * np.log(2 * np.pi * np.e * (var + EPS)) de_list.append(de) if USE_EXTENDED_FEATURES: # PSD近似：log(var) psd_list.append(np.log(var + EPS)) de_feat = np.stack(de_list, axis=0).T.reshape(-1) # (32*5,) if USE_EXTENDED_FEATURES: psd_feat = np.stack(psd_list, axis=0).T.reshape(-1) # (32*5,) f = np.concatenate([de_feat, psd_feat], axis=0).astype(np.float32) else: f = de_feat.astype(np.float32) feats.append(f) if not feats: raise RuntimeError("EEG长度不足以切片（请检查T是否太短，或调整WINDOW_SIZE/STRIDE）") return np.stack(feats, axis=0).astype(np.float32) # ========================================================= # 4) 模型加载 + 推理接口 # ========================================================= def _safe_torch_load(path: str): try: return torch.load(path, map_location=DEVICE, weights_only=False) except TypeError: return torch.load(path, map_location=DEVICE) def load_model(model_path: str) -> tuple[FusionNet, dict]: """ 返回: (model, ckpt_dict) """ obj = _safe_torch_load(model_path) if isinstance(obj, dict) and "model_state" in obj: ckpt = obj state = obj["model_state"] feat_dim = int(obj.get("feat_dim", 320)) else: ckpt = {} state = obj feat_dim = 320 model = FusionNet(num_classes=2, num_eeg_features=feat_dim).to(DEVICE) model.load_state_dict(state, strict=True) model.eval() return model, ckpt def predict_hc_mdd(eeg_dir: str, model_path: str) -> dict: """ 接口：传入 EEG文件夹 和 模型路径，返回判断结果 dict 返回字段： - mat_file: 使用的mat文件 - pred_label: "HC" or "MDD" - p_mdd_mean: 切片p(MDD)均值 - threshold: subject判定阈值 - n_slices: 切片数 """ mat_file = _find_first_mat_file(eeg_dir) # 1) 读EEG (T,C)，并保证变成32通道 eeg = load_eeg_from_mat_any_channels(mat_file) # (T,C) eeg32 = ensure_32_channels(eeg) # (T,32) # 2) 提特征 (N,feat_dim) extractor = FeatureExtractor32(fs=SAMPLING_RATE, window_size=WINDOW_SIZE, stride=STRIDE) feats = extractor.extract(eeg32) # (N, dim) # 3) 加载模型 model, ckpt = load_model(model_path) # 4) 可选：归一化（若ckpt里保存的mean/std维度刚好匹配） mean = ckpt.get("global_mean", None) if isinstance(ckpt, dict) else None std = ckpt.get("global_std", None) if isinstance(ckpt, dict) else None if mean is not None and std is not None: mean = np.asarray(mean, dtype=np.float32) std = np.asarray(std, dtype=np.float32) if mean.shape[0] == feats.shape[1] and std.shape[0] == feats.shape[1]: feats = (feats - mean) / (std + 1e-8) # 不匹配就跳过 # 你说“先不用其他步骤或信息”，所以这里按“尽量运行”处理 # 5) 推理：对所有切片算 p(MDD)，取均值做subject-level x = torch.from_numpy(feats).to(DEVICE) with torch.no_grad(): cls_out, _ = model(x) prob_mdd = torch.softmax(cls_out, dim=1)[:, 1].detach().cpu().numpy() p_mdd_mean = float(np.mean(prob_mdd)) thr = float(ckpt.get("subject_threshold", DEFAULT_SUBJECT_THRESHOLD)) if isinstance(ckpt, dict) else float(DEFAULT_SUBJECT_THRESHOLD) pred_is_mdd = (p_mdd_mean >= thr) pred_label = "MDD" if pred_is_mdd else "HC" return { "mat_file": mat_file, "pred_label": pred_label, "p_mdd_mean": p_mdd_mean, "threshold": thr, "n_slices": int(feats.shape[0]), } # ========================================================= # 5) CLI：命令行运行入口 # ========================================================= def main(): parser = argparse.ArgumentParser(description="Infer HC/MDD from 64ch->32ch EEG mat using a .pth FusionNet model (no Asym).") parser.add_argument("--eeg_dir", type=str, required=True, help="包含.mat EEG文件的文件夹（自动读取第一个.mat）") parser.add_argument("--model_path", type=str, required=True, help="训练好的.pth模型路径") args = parser.parse_args() res = predict_hc_mdd(args.eeg_dir, args.model_path) print("\n========== 推理结果 ==========") print(f"MAT文件: {res['mat_file']}") print(f"切片数量: {res['n_slices']}") print(f"p(MDD)_mean: {res['p_mdd_mean']:.4f}") print(f"阈值thr: {res['threshold']:.4f}") print(f"预测结果: {res['pred_label']}") print("==============================\n") if __name__ == "__main__": main()