| """EEG Annotation Tool adapter for SenuaLab IEDNet-Lite.""" |
|
|
| from __future__ import annotations |
|
|
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
|
|
| from .preprocessing import INTERNATIONAL_10_20_CHANNELS, preprocess_batch as _preprocess |
|
|
|
|
| MODEL_CHANNELS = INTERNATIONAL_10_20_CHANNELS |
| MODEL_INPUT_SAMPLES = 1000 |
| MODEL_SAMPLING_RATE_HZ = 250.0 |
| MODEL_WINDOW_SECONDS = 4.0 |
| MODEL_ENTRY_CLASS = "IEDNetLite" |
| MODEL_NUM_CLASSES = 2 |
| MODEL_CLASS_LABELS = ["Non-IED", "IED"] |
| MODEL_NON_IED_CLASS_INDEX = 0 |
| MODEL_IED_CLASS_INDICES = [1] |
| MODEL_DESCRIPTION = "SenuaLab patient-independent IEDNet-Lite binary IED detector" |
| 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.6004377007484436 |
| MODEL_DEFAULT_STEP_MS = 500.0 |
| MODEL_DEFAULT_PAD_POLICY = "skip" |
| MODEL_DEFAULT_BATCH_SIZE = 64 |
| MODEL_DEFAULT_BATCH_MEMORY_MB = 64.0 |
| MODEL_VALIDATION_NOTE = "Threshold selected on the vEpiSet validation subjects; not externally validated." |
| MODEL_PREPROCESSING_NOTE = "The adapter applies the exact 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, |
| ) |
|
|
|
|
| class SqueezeExcitation(nn.Module): |
| def __init__(self, channels: int, reduction: int = 8): |
| super().__init__() |
| hidden = max(8, channels // reduction) |
| self.net = nn.Sequential( |
| nn.AdaptiveAvgPool1d(1), |
| nn.Conv1d(channels, hidden, 1), |
| nn.GELU(), |
| nn.Conv1d(hidden, channels, 1), |
| nn.Sigmoid(), |
| ) |
|
|
| def forward(self, x): |
| return x * self.net(x) |
|
|
|
|
| class DSResidual(nn.Module): |
| def __init__(self, in_ch: int, out_ch: int, kernel: int, stride: int = 1, dilation: int = 1): |
| super().__init__() |
| pad = dilation * (kernel - 1) // 2 |
| self.body = nn.Sequential( |
| nn.Conv1d(in_ch, in_ch, kernel, stride=stride, padding=pad, dilation=dilation, groups=in_ch, bias=False), |
| nn.BatchNorm1d(in_ch), |
| nn.GELU(), |
| nn.Conv1d(in_ch, out_ch, 1, bias=False), |
| nn.BatchNorm1d(out_ch), |
| SqueezeExcitation(out_ch), |
| ) |
| self.skip = ( |
| nn.Identity() |
| if in_ch == out_ch and stride == 1 |
| else nn.Sequential(nn.Conv1d(in_ch, out_ch, 1, stride=stride, bias=False), nn.BatchNorm1d(out_ch)) |
| ) |
|
|
| def forward(self, x): |
| return F.gelu(self.body(x) + self.skip(x)) |
|
|
|
|
| class IEDNetLite(nn.Module): |
| def __init__(self, num_channels: int = 19, num_classes: int = 2, input_length: int = 1000, subtypes: int = 5): |
| super().__init__() |
| del input_length |
| branches = [] |
| for kernel in (7, 15, 31): |
| branches.append( |
| nn.Sequential( |
| nn.Conv1d(num_channels, 24, kernel, stride=2, padding=kernel // 2, bias=False), |
| nn.BatchNorm1d(24), |
| nn.GELU(), |
| ) |
| ) |
| self.stem = nn.ModuleList(branches) |
| self.backbone = nn.Sequential( |
| DSResidual(72, 72, 9, stride=2), |
| DSResidual(72, 96, 7, stride=2), |
| DSResidual(96, 96, 7, dilation=2), |
| DSResidual(96, 144, 5, stride=2), |
| DSResidual(144, 144, 5, dilation=2), |
| DSResidual(144, 192, 3, stride=2), |
| DSResidual(192, 192, 3, dilation=2), |
| ) |
| self.dropout = nn.Dropout(0.35) |
| self.binary_head = nn.Linear(384, num_classes) |
| self.subtype_head = nn.Linear(384, subtypes) |
|
|
| def forward(self, x): |
| x = torch.cat([branch(x) for branch in self.stem], dim=1) |
| x = self.backbone(x) |
| features = self.dropout(torch.cat([x.mean(-1), x.amax(-1)], dim=1)) |
| return self.binary_head(features), self.subtype_head(features) |
|
|
|
|
