| """Exact four-second preprocessing used by the released SenuaLab models.""" |
|
|
| from __future__ import annotations |
|
|
| import math |
|
|
| import numpy as np |
| from scipy.signal import butter, resample, resample_poly, sosfiltfilt |
|
|
|
|
| NUM_CHANNELS = 19 |
| INTERNATIONAL_10_20_CHANNELS = [ |
| "Fp1", "Fp2", "F3", "F4", "C3", "C4", "P3", "P4", "O1", "O2", |
| "F7", "F8", "T3", "T4", "T5", "T6", "Fz", "Cz", "Pz", |
| ] |
| DATASET_29_CHANNEL_ORDER = INTERNATIONAL_10_20_CHANNELS |
|
|
|
|
| def normalize_sample(sample): |
| x = np.asarray(sample, dtype=np.float32) |
| mean = x.mean() |
| std = x.std() |
| return np.clip((x - mean) / max(float(std), 1e-7), -8.0, 8.0).astype(np.float32) |
|
|
|
|
| def preprocess_batch( |
| batch, |
| *, |
| source_sfreq=None, |
| target_sfreq=250, |
| target_samples=1000, |
| channel_names=None, |
| ): |
| """Filter, average-reference, resample, and globally normalize `[B,C,T]`.""" |
|
|
| del channel_names |
| x = np.asarray(batch, dtype=np.float32) |
| if x.ndim == 2: |
| x = x[None] |
| if x.ndim != 3 or x.shape[1] != NUM_CHANNELS: |
| raise ValueError(f"Expected [batch,19,time], received {x.shape}") |
| sfreq = float(source_sfreq or target_sfreq) |
| if sfreq <= 90.0: |
| raise ValueError(f"Sampling rate must be above the 45 Hz cutoff; received {sfreq}") |
| sos = butter(4, [1.0, 45.0], btype="bandpass", fs=sfreq, output="sos") |
| x = sosfiltfilt(sos, x, axis=-1).astype(np.float32, copy=False) |
| x -= x.mean(axis=1, keepdims=True) |
| if not math.isclose(sfreq, float(target_sfreq)): |
| up = int(target_sfreq) |
| down = int(round(sfreq)) |
| common = math.gcd(up, down) |
| x = resample_poly(x, up // common, down // common, axis=-1).astype(np.float32) |
| if x.shape[-1] != int(target_samples): |
| x = resample(x, int(target_samples), axis=-1).astype(np.float32) |
| mean = x.mean(axis=(1, 2), keepdims=True) |
| std = x.std(axis=(1, 2), keepdims=True) |
| return np.clip((x - mean) / np.maximum(std, 1e-7), -8.0, 8.0).astype(np.float32) |
|
|
|
|
