File size: 4,422 Bytes
dc9d5f6 | 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 | """EEG Annotation Tool adapter for the SenuaLab EEGPT temporal-head model."""
from __future__ import annotations
from typing import Any
import mne
import torch
from torch import nn
try:
from braindecode.models import EEGPT
except ImportError as exc: # imported by EEG Annotation Tool during discovery
raise ImportError("This model requires braindecode[hub]==1.6.1") from exc
from .preprocessing import INTERNATIONAL_10_20_CHANNELS, preprocess_batch as _preprocess
EEGPT_CHANNELS = [
"Fp1", "Fp2", "F3", "F4", "C3", "C4", "P3", "P4", "O1", "O2",
"F7", "F8", "T7", "T8", "P7", "P8", "Fz", "Cz", "Pz",
]
MODEL_CHANNELS = INTERNATIONAL_10_20_CHANNELS
MODEL_CHANNEL_ALIASES = {"T3": "T7", "T4": "T8", "T5": "P7", "T6": "P8"}
MODEL_INPUT_SAMPLES = 1000
MODEL_SAMPLING_RATE_HZ = 250.0
MODEL_WINDOW_SECONDS = 4.0
MODEL_ENTRY_CLASS = "SenuaEEGPT"
MODEL_NUM_CLASSES = 2
MODEL_CLASS_LABELS = ["Non-IED", "IED"]
MODEL_NON_IED_CLASS_INDEX = 0
MODEL_IED_CLASS_INDICES = [1]
MODEL_DESCRIPTION = "SenuaLab EEGPT encoder with an IED-specific temporal head"
MODEL_BATCH_PREPROCESSOR = "preprocess_batch"
MODEL_REQUIRED_REFERENCE = "common average (applied by model adapter)"
MODEL_REQUIRED_FILTERS = ["1-45 Hz zero-phase Butterworth (applied by model adapter)"]
MODEL_REQUIRED_NORMALIZATION = "global four-second window z-score, clipped to [-8,8]"
MODEL_INPUT_UNIT = "scale-invariant after window z-score"
MODEL_SOURCE_SIGNAL_POLICY = "raw"
MODEL_REQUIRES_FULL_WINDOW = True
MODEL_REQUIRES_ALL_CHANNELS = True
MODEL_DEFAULT_THRESHOLD = 0.8916015625
MODEL_DEFAULT_STEP_MS = 500.0
MODEL_DEFAULT_PAD_POLICY = "skip"
MODEL_DEFAULT_BATCH_SIZE = 16
MODEL_DEFAULT_BATCH_MEMORY_MB = 256.0
MODEL_VALIDATION_NOTE = "Threshold selected on vEpiSet validation subjects; requires Braindecode 1.6.1 and is not externally validated."
MODEL_PREPROCESSING_NOTE = "The adapter applies the released 1-45 Hz, common-average, resampling, and global-window z-score pipeline."
def preprocess_batch(batch, source_sfreq=None, channel_names=None):
return _preprocess(
batch,
source_sfreq=source_sfreq,
target_sfreq=250,
target_samples=1000,
channel_names=channel_names,
)
def _chs_info() -> list[dict[str, Any]]:
info = mne.create_info(EEGPT_CHANNELS, sfreq=250.0, ch_types="eeg")
info.set_montage("standard_1020")
return info["chs"]
class EEGPTIEDHead(nn.Module):
def __init__(self, hidden: int = 128, n_outputs: int = 2):
super().__init__()
flattened_embedding = 4 * 512
self.input_norm = nn.LayerNorm(flattened_embedding)
self.patch_projection = nn.Sequential(
nn.Linear(flattened_embedding, hidden), nn.GELU(), nn.Dropout(0.20)
)
self.temporal = nn.Sequential(
nn.Conv1d(hidden, hidden, 5, padding=2, groups=hidden, bias=False),
nn.BatchNorm1d(hidden),
nn.Conv1d(hidden, hidden, 1, bias=False),
nn.GELU(),
nn.Dropout(0.20),
nn.Conv1d(hidden, hidden, 3, padding=2, dilation=2, groups=hidden, bias=False),
nn.BatchNorm1d(hidden),
nn.Conv1d(hidden, hidden, 1, bias=False),
nn.GELU(),
)
self.attention = nn.Conv1d(hidden, 1, 1)
self.classifier = nn.Sequential(
nn.LayerNorm(hidden * 3), nn.Dropout(0.35), nn.Linear(hidden * 3, n_outputs)
)
def forward(self, z: torch.Tensor) -> torch.Tensor:
patches = self.input_norm(z.flatten(2))
patches = self.patch_projection(patches).transpose(1, 2)
patches = patches + self.temporal(patches)
weights = self.attention(patches).softmax(dim=-1)
pooled = torch.cat(
[(patches * weights).sum(-1), patches.mean(-1), patches.amax(-1)], dim=1
)
return self.classifier(pooled)
class SenuaEEGPT(EEGPT):
def __init__(self, num_channels: int = 19, num_classes: int = 2, input_length: int = 1000):
if num_channels != 19 or input_length != 1000:
raise ValueError("SenuaEEGPT requires 19 channels and 1,000 samples")
super().__init__(
n_outputs=num_classes,
n_chans=19,
chs_info=_chs_info(),
n_times=1000,
sfreq=250.0,
chan_proj_type="none",
return_encoder_output=False,
)
self.final_layer = EEGPTIEDHead(n_outputs=num_classes)
|