| """Serializable preprocessing for trained CTNet checkpoints.""" |
|
|
| from __future__ import annotations |
|
|
| from typing import Any, Optional |
|
|
| import numpy as np |
| import torch |
| from transformers.feature_extraction_utils import BatchFeature, FeatureExtractionMixin |
|
|
|
|
| class CtnetPreprocessor(FeatureExtractionMixin): |
| """Apply the training-set normalization and validate the EEG layout.""" |
|
|
| model_input_names = ["input_values"] |
|
|
| def __init__( |
| self, |
| *, |
| n_channels: int, |
| n_times: int, |
| sampling_rate: Optional[int] = None, |
| mean: Any, |
| std: Any, |
| channel_names: Optional[list[str]] = None, |
| dataset: Optional[str] = None, |
| subjects: Optional[list[int]] = None, |
| unit: str = "microvolts", |
| train_session: Optional[str] = None, |
| selection_session: Optional[str] = None, |
| **kwargs: Any, |
| ) -> None: |
| if n_channels < 1 or n_times < 1: |
| raise ValueError("n_channels and n_times must be positive.") |
| if channel_names is not None and len(channel_names) != n_channels: |
| raise ValueError( |
| f"Expected {n_channels} channel names, got {len(channel_names)}." |
| ) |
|
|
| self.n_channels = int(n_channels) |
| self.n_times = int(n_times) |
| self.sampling_rate = sampling_rate |
| self.mean = np.asarray(mean, dtype=np.float32) |
| self.std = np.asarray(std, dtype=np.float32) |
| if np.any(self.std <= 0): |
| raise ValueError("std must be positive.") |
| self.channel_names = channel_names |
| self.dataset = dataset |
| self.subjects = subjects |
| self.unit = unit |
| self.train_session = train_session |
| self.selection_session = selection_session |
| super().__init__(**kwargs) |
|
|
| def __call__( |
| self, |
| input_values: torch.Tensor | np.ndarray, |
| *, |
| return_tensors: str = "pt", |
| ) -> BatchFeature: |
| """Normalize one trial or a batch in ``(channels, time)`` layout.""" |
| if isinstance(input_values, torch.Tensor): |
| values = input_values.detach().cpu().numpy() |
| else: |
| values = np.asarray(input_values) |
| if values.ndim == 2: |
| values = values[np.newaxis, ...] |
| if values.ndim != 3: |
| raise ValueError( |
| "Expected EEG input with shape (channels, time) or " |
| "(batch_size, channels, time)." |
| ) |
| if values.shape[1] != self.n_channels: |
| raise ValueError( |
| f"Expected {self.n_channels} EEG channels but received " |
| f"{values.shape[1]}." |
| ) |
| if values.shape[2] != self.n_times: |
| raise ValueError( |
| f"Expected {self.n_times} time samples but received " |
| f"{values.shape[2]}." |
| ) |
|
|
| values = values.astype(np.float32, copy=True) |
| values = (values - self.mean) / self.std |
| if return_tensors not in {"np", "pt"}: |
| raise ValueError("return_tensors must be 'pt' or 'np'.") |
| return BatchFeature({"input_values": values}, tensor_type=return_tensors) |
|
|
| def to_dict(self) -> dict[str, Any]: |
| payload = super().to_dict() |
| payload.update( |
| { |
| "feature_extractor_type": self.__class__.__name__, |
| "n_channels": self.n_channels, |
| "n_times": self.n_times, |
| "sampling_rate": self.sampling_rate, |
| "mean": self.mean.tolist(), |
| "std": self.std.tolist(), |
| "channel_names": self.channel_names, |
| "dataset": self.dataset, |
| "subjects": self.subjects, |
| "unit": self.unit, |
| "train_session": self.train_session, |
| "selection_session": self.selection_session, |
| } |
| ) |
| return payload |
|
|
