File size: 4,388 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 118 119 120 | """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)
|