Seizure probability (iEEG)

Seizure probability from intracranial EEG with a CNN + bidirectional LSTM applied to spectrograms (Sladky et al. 2022, Brain Communications, doi:10.1093/braincomms/fcac115).

This is the brainmaze-torch 0.2.0 seizure pipeline with the model running on ONNX Runtime (PyTorch is not needed). The public functions have the same names, arguments and outputs as brainmaze_torch.seizure_detection 0.2.0, so code moves over by changing the import:

# before: from brainmaze_torch.seizure_detection import predict_channel_seizure_probability
from brainmaze_eeg_models.seizure import predict_channel_seizure_probability

Differences from brainmaze-torch 0.2.0:

• load_trained_model() returns an OnnxModel (not a torch.nn.Module); passing a PyTorch model raises TypeError.

• Device: the default is the CPU, as in 0.2.0 (use_cuda=None and device=None -> CPU; load_trained_model() defaults to device='cpu'). The GPU is opt-in: use_cuda=True (GPU cuda_number; an error if CUDA is unusable, never a silent CPU fallback), or the keyword device='cuda' / 'auto' (GPU if usable, else CPU). Giving both use_cuda and device is allowed only when they agree (device='auto' with an explicit use_cuda raises). A model from load_trained_model() runs where it was loaded (device and GPU); a contradicting use_cuda / device / cuda_number raises. cuda_number defaults to None (= GPU 0 for a model loaded by name). On the GPU, TF32 is disabled; GPU results are expected to match the CPU to float32 rounding but have not been measured on a GPU yet.

• Numbers: the ONNX models (exported from the 0.2.0 weights, tools/export_seizure_onnx.py in the repository) give the same probabilities as PyTorch within float32 rounding: typically ~1e-6, at most ~2e-5 (largest near p = 0.5; the float32 BiLSTM accumulates rounding differently in the two runtimes); NaN masks identical. The parity tests use a tolerance of 1e-4, like the brainmaze-torch golden tests, which pass unchanged.

Bundled models (load_trained_model()):

• 'modelA': the model from the published work.

• 'modelB': the same architecture trained on an extended data set.

The model was designed for 300 s inputs; its outputs near the ends of an input have little LSTM context and are less reliable, which is why the high-level function drops discard_edges_s at each window edge.

Two ways to use it:

1. High level (recommended): predict_channel_seizure_probability() turns one long channel into a continuous probability trace on a 0.5 s grid. It handles the windowing, gaps, flat segments and the recording edges.

2. Low level: preprocess_input() + infer_seizure_probability() for batches of windows you cut yourself.

Conventions (both levels)

• fs must be a whole, even number of Hz, at least 200 Hz. A float such as 500.0 (typical of .mat / MEF headers) is accepted. Odd or fractional rates raise ValueError: resample first (anti-aliasing is the caller’s job). These rules come from the spectrogram grid (whole samples per 1 s segment, an integer 0.5 s hop, 100 bins below Nyquist), not from the training data. At 200-256 Hz the anti-alias filter of the acquisition or resampling usually attenuates the upper bins (from ~0.4 fs = 80-100 Hz), which the model did not see in training: prefer higher rates (see preprocess_input()).

• Spectrogram: 1 s segments (nperseg = fs), 0.5 s hop, bins 0-99 Hz at 1 Hz. Column j of a window describes the 1 s of signal centred (j + 1) * 0.5 s after the window’s first sample.

• In the output of predict_channel_seizure_probability(), NaN means “not evaluated”, never “no seizure”: t = 0 and every 1 s segment containing NaN/inf or a flat signal are NaN; every other segment has a value. Combine channels or time with NaN-aware functions (np.nanmax) and never replace NaN by 0.

• The low-level preprocess_input() zero-fills NaN samples and does not mark them; masking is done by the high-level function.

Example (high level, one channel)

import numpy as np
from brainmaze_eeg_models.seizure import predict_channel_seizure_probability

fs = 500
x = np.random.randn(fs * 600)                 # 10 min, one channel
x[100 * fs:130 * fs] = np.nan                 # a 30 s gap
t, p = predict_channel_seizure_probability(x, fs, model='modelA')
# t = 0, 0.5, 1.0, ... s; p is NaN at t = 0 and over the gap

Example (low level, batch of 300 s windows)

import numpy as np
from brainmaze_eeg_models.seizure import (
    load_trained_model, preprocess_input, infer_seizure_probability)

model = load_trained_model('modelA')
fs = 500
x = np.random.randn(3, fs * 300)              # 3 windows (rows) of 300 s
t, f, sxx = preprocess_input(x, fs, return_axes=True)   # sxx: (3, 100, 599)
y = infer_seizure_probability(sxx, model)     # (3, 599); y[:, j] belongs to t[j]

Moving from brainmaze-torch

brainmaze-torch is retired; its last release will re-export these functions with a DeprecationWarning. The public names are the same:

brainmaze-torch 0.2.0 (brainmaze_torch.seizure_detection)	brainmaze-eeg-models (brainmaze_eeg_models.seizure)
predict_channel_seizure_probability(x, fs, model, use_cuda, cuda_number, ...)	same signature and output; default device the CPU (as 0.2.0); use_cuda=True or the keyword device='cuda' / 'auto' opt in to the GPU (both given: must agree)
preprocess_input(x, fs, return_axes=False)	identical code and output
infer_seizure_probability(x, model, use_cuda, cuda_number)	same; model may also be a name (default 'modelA')
load_trained_model('modelA' | 'modelB') -> torch.nn.Module	-> OnnxModel (device default 'cpu', threads, cuda_device_id)

One 0.2.0 pattern now raises (loudly, never silently): in brainmaze-torch 0.2.0, load_trained_model(name) had no device argument, and the GPU was requested per call, which moved the torch model to cuda:<cuda_number> temporarily:

m = load_trained_model('modelA')                                   # 0.2.0: always a CPU model
predict_channel_seizure_probability(x, fs, m, use_cuda=True)       # 0.2.0: ran on cuda:0

Here a loaded model is an ONNX Runtime session that cannot move between devices, so this raises ValueError: the loaded model runs on 'cpu' but 'cuda' was requested. Replacement: choose the device when loading, or pass the model name and let the call load it:

m = load_trained_model('modelA', device='cuda', cuda_device_id=0)  # GPU k: cuda_device_id=k
predict_channel_seizure_probability(x, fs, m)
# or
predict_channel_seizure_probability(x, fs, 'modelA', use_cuda=True, cuda_number=0)

A model loaded on GPU k together with a different cuda_number also raises (0.2.0 moved the model to the requested GPU). cuda_number with a CPU model and no GPU request is ignored, as in 0.2.0. The brainmaze-torch 0.3.0 compatibility release (re-exporting these functions) should name this change in its DeprecationWarning and release notes.

Validation

• The brainmaze-torch golden tests (outputs of the published brainmaze-torch 0.1.1 code on a 15 min iEEG seizure recording, incl. gaps and truncations) pass unchanged on the ONNX path (tolerance 1e-4; observed differences about 1e-6), as do its robustness tests.

• Parity with PyTorch (tests/test_seizure_parity.py, needs pip install -e ".[export]"): end-to-end probabilities within 1e-4 (the golden tolerance) and identical NaN masks for both models on the fixture (with and without a gap) and on synthetic recordings at 200-1000 Hz with bursts, gaps and flat segments; preprocessing bit-identical. Observed differences are ~1e-6 typically and up to 2.1e-5 (modelB, a 400 Hz synthetic recording, at p = 0.49: the float32 BiLSTM over 599 steps accumulates rounding differently in the two runtimes, and the softmax is steepest near 0.5). These tests need PyTorch and are not run in CI; run them by hand when torch, onnx or onnxruntime change.

• GPU: not yet measured on a GPU (TF32 is off, so agreement to float32 rounding is expected).

• The ONNX files are produced by tools/export_seizure_onnx.py from the brainmaze-torch 0.2.0 weights (reproducible byte for byte with the same PyTorch version) and their SHA-256 is verified when they are loaded.

API

brainmaze_eeg_models.seizure.predict_channel_seizure_probability(x, fs, model='modelA', use_cuda=None, cuda_number=None, n_batch=128, window_s=300, step_s=20, discard_edges_s=10, min_valid_fraction=0.0, fill_recording_edges=True, *, device=None)

Continuous seizure-probability trace for one (long) iEEG channel.

The signal is cut into overlapping windows of window_s seconds every step_s seconds (the “regular” windows of the published method). Each window is converted to a spectrogram and run through the model; discard_edges_s seconds are dropped at both window edges (the BiLSTM has little context there), and overlapping windows are combined with a NaN-ignoring maximum. If the recording end does not fall on the step grid, one extra window aligned to the recording end covers the tail; it only fills time that no regular window covers and never changes a value produced by the regular windows. With the default parameters the output is therefore identical to the published method (brainmaze-torch <= 0.1.1) wherever that method produced an estimate, except that invalid (gap / flat) segments are NaN instead of a number.

Parameters:

• x (array_like, shape (n_samples,)) – Single-channel raw signal (1D). Lists and integer arrays are accepted. NaN / +-inf mark missing data (gaps).

• fs (int or float) – Sampling rate in Hz. Must be a whole, even number >= 200 Hz; a float such as 500.0 is accepted. Resample other rates first.

• model (str or OnnxModel, optional) – 'modelA' (published model), 'modelB' (extended training set) or a model from load_trained_model(). Default 'modelA'.

• use_cuda (bool, optional) – True: GPU cuda_number (error if CUDA is unusable); False: CPU. Default None: device, else the CPU (as in brainmaze-torch 0.2.0; a loaded model runs where it was loaded).

• cuda_number (int, optional) – CUDA device index for a model loaded by name. Default None (= GPU 0; a loaded model runs on its own GPU, and a different cuda_number raises ValueError).

• n_batch (int, optional) – Windows per inference batch (memory/speed trade-off only; does not change the result). Default 128.

• window_s (float, optional) – Window length in seconds; multiple of 0.5, >= 2. Default 300 (the length the model was designed for). The recording must be at least this long.

• step_s (float, optional) – Step between window starts in seconds; multiple of 0.5. Must satisfy step_s <= window_s - 2 * discard_edges_s - 0.5 so that the kept (non-discarded) parts of consecutive windows leave no hole (a window has 2 * window_s - 1 half-second columns). Default 20.

• discard_edges_s (float, optional) – Seconds dropped at each window edge; multiple of 0.5, may be 0. Default 10.

• min_valid_fraction (float, optional) – Window preference near gaps. A window is preferred if at least this fraction of its 1 s segments (spectrogram columns) are valid (finite samples, not flat). A valid segment gets the maximum over the preferred windows covering it; only if none of them covers it, the maximum over all windows covering it is used. This never turns a valid segment into NaN; it only lets gap-heavy (mostly zero-filled) windows be ignored where a better window exists. Default 0.0: every window is preferred, i.e. the published method. Values > 0 change probabilities next to long gaps compared with the published method (not validated). The choice is made per time point, so the result is not monotone in this value: a point switches from “preferred windows” to the fallback (all windows) as soon as the last preferred window covering it drops below the threshold, and different points switch at different values. If the value exceeds the valid fraction of every window, no window is preferred and the result equals that of 0.0. Choose one value for a study and do not interpolate between values.

• fill_recording_edges (bool, optional) – The first and last discard_edges_s seconds of the recording lie in no window’s kept interior. If True (default), they are taken from the edge of the first/last window (shorter model context, hence less reliable). If False, they are NaN.

• device ({'auto', 'cpu', 'cuda'}, optional) – Alternative to use_cuda (keyword only); 'auto' opts in to the GPU when it is usable. Giving both device and use_cuda is allowed only when they agree.

Returns:

• t (numpy.ndarray, shape (n_out,)) – Time in seconds from the first sample: t[k] = k * 0.5 (so the value for time t is at index 2 * t), n_out = n_samples // (fs // 2). The last stamp is the centre of the last complete 1 s segment.

• prob (numpy.ndarray, shape (n_out,), float64) – Seizure probability in [0, 1] for the 1 s segment centred on t[k], or NaN where no estimate exists. NaN is returned for:

  • t = 0 (a 1 s segment centred at 0 s would start before the data),

  • every 1 s segment containing at least one NaN/inf sample (gaps),

  • every 1 s segment over which the signal is constant (flat line, e.g. disconnected or saturated channel),

  • the first/last discard_edges_s s if fill_recording_edges is False.

  NaN therefore always means “not evaluated”, never “no seizure”; a value is never 0 merely because a time point was not covered. Every other (valid) segment always gets a value, whatever min_valid_fraction is.

Raises:

ValueError – If x is not 1D, fs is invalid, the window parameters are inconsistent, or the recording is shorter than window_s.

Notes

• Gaps are not interpolated. Within an evaluated window the gap samples are zero-filled for the model (as in the original implementation), which can influence the probability of valid segments next to a gap via the BiLSTM context; only the gap segments themselves are masked to NaN. To fill short gaps instead (e.g. by interpolation), do so before calling this function; filled samples are then treated as valid data. See min_valid_fraction to ignore mostly-gap windows where possible.

• Probabilities of overlapping windows are combined with a NaN-ignoring maximum (np.fmax), as in the published method.

Examples

>>> import numpy as np
>>> from brainmaze_eeg_models.seizure import predict_channel_seizure_probability
>>> fs = 500
>>> x = np.random.randn(fs * 600)          # 10 min, single channel
>>> t, p = predict_channel_seizure_probability(x, fs, model='modelA')
>>> t.shape == p.shape, float(t[1] - t[0])
(True, 0.5)

brainmaze_eeg_models.seizure.preprocess_input(x, fs, return_axes=False)

Convert raw signal windows into normalised spectrograms for the model.

Each row of x is processed independently:

1. z-score the raw signal (NaN-aware mean/std over the row; a row with zero or undefined std is only mean-centred, never divided by 0),

2. replace NaN / +-inf samples by 0,

3. spectrogram with nperseg = fs (1 s, 1 Hz bins) and a 0.5 s hop (noverlap = fs // 2), keeping the first 100 bins (0-99 Hz),

4. z-score every frequency bin over time, using only columns with non-zero total power. Columns with zero power (e.g. a flat or zero-filled gap) are left at 0. If a bin has zero variance, or fewer than two columns have power, that bin is only mean-centred (no division by 0).

Parameters:

• x (array_like, shape (n_samples,) or (batch_size, n_samples)) – Raw signal. A 1D input is treated as a single window. Rows are windows (or channels), columns are samples. Must contain at least fs samples (1 s).

• fs (int or float) – Sampling rate in Hz. Must be a whole, even number >= 200; a float such as 500.0 is accepted. Odd or fractional rates raise ValueError – resample first (see Notes).

• return_axes (bool, optional) – If True, also return the time and frequency axes. Default False.

Returns:

• t (numpy.ndarray, shape (n_times,)) – Only if return_axes. Time stamp in seconds of every column, relative to the first sample of the window: (j + 1) * 0.5 – the centre of the 1 s segment [j * 0.5, j * 0.5 + 1).

• f (numpy.ndarray, shape (100,)) – Only if return_axes. Frequencies in Hz (0, 1, …, 99).

• sxx (numpy.ndarray, shape (batch_size, 100, n_times)) – Normalised spectrograms, n_times = floor((n_samples - fs) / (fs // 2)) + 1 (= 2 * duration_s - 1 for whole-second inputs). Always finite.

Raises:

ValueError – If fs is invalid, x is not 1D/2D, or shorter than 1 s.

Notes

NaN samples are zero-filled before the spectrogram, so the model still sees a value there. This low-level function does not mark such columns; use predict_channel_seizure_probability(), which reports any 1 s segment containing NaN/inf (or a flat signal) as NaN in its output.

Where the fs rules come from: they follow from the spectrogram grid, not from the training data (whose sampling rate is not documented in this repository). nperseg = fs needs a whole number of samples per second; the 0.5 s hop (fs / 2 samples) needs an even rate, otherwise the output would drift off the 0.5 s grid; and 100 one-Hz bins (0-99 Hz) need a Nyquist frequency of at least 100 Hz, i.e. fs >= 200.

Anti-aliasing caveat at 200-256 Hz: acquisition (or resampling) anti-alias filters usually start attenuating at about 0.4 * fs, i.e. 80-100 Hz at these rates. The upper spectrogram bins then hold less power than in recordings with a higher native rate, which the model never saw in that form: a silent domain shift. Prefer a native rate >= 250-500 Hz, and when you resample, keep the new rate high enough that the anti-alias filter’s transition band lies above 100 Hz.

brainmaze_eeg_models.seizure.infer_seizure_probability(x, model='modelA', use_cuda=None, cuda_number=None, *, device=None)

Run the seizure model on a batch of preprocessed spectrograms.

Parameters:

• x (array_like, shape (batch_size, 100, n_times)) – Output of preprocess_input(). Must be finite and have n_times >= 2. Recommended window length is 300 s (599 columns).

• model (str or OnnxModel, optional) – 'modelA' (default), 'modelB', or a model from load_trained_model() (then the device arguments must agree with the device it was loaded on).

• use_cuda (bool, optional) – True: run on the GPU cuda_number (error if CUDA is unusable); False: run on the CPU. Default None: device, else the CPU (as in brainmaze-torch 0.2.0; a loaded model runs where it was loaded).

• cuda_number (int, optional) – CUDA device index for a model loaded by name. Default None (= GPU 0; a loaded model runs on its own GPU, and a different cuda_number raises ValueError).

• device ({'auto', 'cpu', 'cuda'}, optional) – Alternative to use_cuda (keyword only); 'auto' opts in to the GPU when it is usable. Giving both device and use_cuda is allowed only when they agree.

Returns:

Seizure probability (softmax probability of class index 3 of the model’s 4 output classes) for every spectrogram column. Column j corresponds to the time stamp t[j] returned by preprocess_input() ((j + 1) * 0.5 s from window start). Values are in [0, 1].

Return type:

numpy.ndarray, shape (batch_size, n_times), float32

Raises:

ValueError – If x has the wrong shape or contains non-finite values.

Notes

The ONNX model gives the same probabilities as the PyTorch model of brainmaze-torch 0.2.0 to within float32 rounding: typically ~1e-6, at most ~2e-5 in an independent 40-run check (largest near p = 0.5, where the softmax is steepest; the float32 BiLSTM over 599 steps accumulates rounding differently in the two runtimes).

brainmaze_eeg_models.seizure.load_trained_model(model_name, device='cpu', threads=None, cuda_device_id=0)

Load one of the bundled, pre-trained seizure-detection models.

Parameters:

• model_name ({'modelA', 'modelB'}) – 'modelA': the model from the published work (Sladky et al. 2022). 'modelB': the same architecture trained on an extended data set.

• device ({'cpu', 'auto', 'cuda'}) – Default 'cpu', as brainmaze-torch 0.2.0 (bit-reproducible across machines). 'auto' (GPU if usable, else CPU) and 'cuda' are opt-in; see brainmaze_eeg_models.runtime.OnnxModel ('cuda' never falls back to the CPU silently).

• threads (int, optional) – CPU threads per inference call (None: ONNX Runtime default).

• cuda_device_id (int) – GPU index for 'cuda' / 'auto'.

Returns:

Input 'x': spectrograms (batch, 100, n_times) from preprocess_input(); outputs 'logits' and 'probs', both (n_times, batch, 4) (class index 3 = seizure). Loaded models are cached (one session per name/device/threads/GPU).

Return type:

OnnxModel

Raises:

KeyError – If model_name is not one of the bundled models.