eegannotationtool/senua_iednet_lite/model.py

"""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)