562 lines
18 KiB
Python
562 lines
18 KiB
Python
# -*- 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()
|