Spaces:
Sleeping
Sleeping
| """ | |
| EEG-to-fMRI Translation Model. | |
| MLP that maps EEG features (236-dim) to fMRI BOLD signals (96 regions). | |
| Architecture: 236 → 512 → 256 → 128 → 96 | |
| Harvard-Oxford Cortical Atlas: 96 regions covering the full cortex. | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| import numpy as np | |
| import os | |
| from typing import Optional | |
| # Harvard-Oxford atlas region names (96 regions) | |
| HARVARD_OXFORD_REGIONS = [ | |
| "Frontal Pole", | |
| "Insular Cortex", | |
| "Superior Frontal Gyrus", | |
| "Middle Frontal Gyrus", | |
| "Inferior Frontal Gyrus pars triangularis", | |
| "Inferior Frontal Gyrus pars opercularis", | |
| "Precentral Gyrus", | |
| "Temporal Pole", | |
| "Superior Temporal Gyrus anterior", | |
| "Superior Temporal Gyrus posterior", | |
| "Middle Temporal Gyrus anterior", | |
| "Middle Temporal Gyrus posterior", | |
| "Middle Temporal Gyrus temporooccipital", | |
| "Inferior Temporal Gyrus anterior", | |
| "Inferior Temporal Gyrus posterior", | |
| "Inferior Temporal Gyrus temporooccipital", | |
| "Postcentral Gyrus", | |
| "Superior Parietal Lobule", | |
| "Supramarginal Gyrus anterior", | |
| "Supramarginal Gyrus posterior", | |
| "Angular Gyrus", | |
| "Lateral Occipital Cortex superior", | |
| "Lateral Occipital Cortex inferior", | |
| "Intracalcarine Cortex", | |
| "Frontal Medial Cortex", | |
| "Juxtapositional Lobule Cortex", | |
| "Subcallosal Cortex", | |
| "Paracingulate Gyrus", | |
| "Cingulate Gyrus anterior", | |
| "Cingulate Gyrus posterior", | |
| "Precuneous Cortex", | |
| "Cuneal Cortex", | |
| "Frontal Orbital Cortex", | |
| "Parahippocampal Gyrus anterior", | |
| "Parahippocampal Gyrus posterior", | |
| "Lingual Gyrus", | |
| "Temporal Fusiform Cortex anterior", | |
| "Temporal Fusiform Cortex posterior", | |
| "Temporal Occipital Fusiform Cortex", | |
| "Occipital Fusiform Gyrus", | |
| "Frontal Operculum Cortex", | |
| "Central Opercular Cortex", | |
| "Parietal Operculum Cortex", | |
| "Planum Polare", | |
| "Heschl Gyrus", | |
| "Planum Temporale", | |
| "Supracalcarine Cortex", | |
| "Occipital Pole", | |
| # Right hemisphere (mirrored) | |
| "Frontal Pole R", | |
| "Insular Cortex R", | |
| "Superior Frontal Gyrus R", | |
| "Middle Frontal Gyrus R", | |
| "Inferior Frontal Gyrus pars triangularis R", | |
| "Inferior Frontal Gyrus pars opercularis R", | |
| "Precentral Gyrus R", | |
| "Temporal Pole R", | |
| "Superior Temporal Gyrus anterior R", | |
| "Superior Temporal Gyrus posterior R", | |
| "Middle Temporal Gyrus anterior R", | |
| "Middle Temporal Gyrus posterior R", | |
| "Middle Temporal Gyrus temporooccipital R", | |
| "Inferior Temporal Gyrus anterior R", | |
| "Inferior Temporal Gyrus posterior R", | |
| "Inferior Temporal Gyrus temporooccipital R", | |
| "Postcentral Gyrus R", | |
| "Superior Parietal Lobule R", | |
| "Supramarginal Gyrus anterior R", | |
| "Supramarginal Gyrus posterior R", | |
| "Angular Gyrus R", | |
| "Lateral Occipital Cortex superior R", | |
| "Lateral Occipital Cortex inferior R", | |
| "Intracalcarine Cortex R", | |
| "Frontal Medial Cortex R", | |
| "Juxtapositional Lobule Cortex R", | |
| "Subcallosal Cortex R", | |
| "Paracingulate Gyrus R", | |
| "Cingulate Gyrus anterior R", | |
| "Cingulate Gyrus posterior R", | |
| "Precuneous Cortex R", | |
| "Cuneal Cortex R", | |
| "Frontal Orbital Cortex R", | |
| "Parahippocampal Gyrus anterior R", | |
| "Parahippocampal Gyrus posterior R", | |
| "Lingual Gyrus R", | |
| "Temporal Fusiform Cortex anterior R", | |
| "Temporal Fusiform Cortex posterior R", | |
| "Temporal Occipital Fusiform Cortex R", | |
| "Occipital Fusiform Gyrus R", | |
| "Frontal Operculum Cortex R", | |
| "Central Opercular Cortex R", | |
| "Parietal Operculum Cortex R", | |
| "Planum Polare R", | |
| "Heschl Gyrus R", | |
| "Planum Temporale R", | |
| "Supracalcarine Cortex R", | |
| "Occipital Pole R", | |
| ] | |
| # Network assignments for each region | |
| NETWORK_ASSIGNMENTS = { | |
| "Visual": [21, 22, 23, 31, 35, 38, 39, 46, 47, 69, 70, 71, 79, 83, 86, 87, 94, 95], | |
| "Motor": [6, 16, 17, 25, 54, 64, 65, 73], | |
| "Language": [4, 5, 8, 9, 10, 11, 18, 19, 20, 44, 45, 52, 53, 56, 57, 58, 59, 66, 67, 68, 92, 93], | |
| "Default Mode": [0, 24, 27, 28, 29, 30, 33, 34, 48, 72, 75, 76, 77, 78, 81, 82], | |
| "Executive": [1, 2, 3, 26, 32, 40, 41, 42, 49, 50, 51, 74, 80, 88, 89, 90], | |
| } | |
| class EEGtoFMRIModel(nn.Module): | |
| """ | |
| MLP for translating EEG features to fMRI BOLD signals. | |
| Architecture: | |
| Input (236) → Linear(512) → BN → GELU → Dropout | |
| → Linear(256) → BN → GELU → Dropout | |
| → Linear(128) → BN → GELU → Dropout | |
| → Linear(96) → Sigmoid | |
| Output is in [0, 1] range representing normalized BOLD signal. | |
| """ | |
| def __init__( | |
| self, | |
| input_dim: int = 236, | |
| output_dim: int = 96, | |
| hidden_dims: list[int] = [512, 256, 128], | |
| dropout: float = 0.2, | |
| ): | |
| super().__init__() | |
| self.input_dim = input_dim | |
| self.output_dim = output_dim | |
| layers = [] | |
| prev_dim = input_dim | |
| for hidden_dim in hidden_dims: | |
| layers.extend([ | |
| nn.Linear(prev_dim, hidden_dim), | |
| nn.BatchNorm1d(hidden_dim), | |
| nn.GELU(), | |
| nn.Dropout(dropout), | |
| ]) | |
| prev_dim = hidden_dim | |
| layers.append(nn.Linear(prev_dim, output_dim)) | |
| layers.append(nn.Sigmoid()) | |
| self.network = nn.Sequential(*layers) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| return self.network(x) | |
| class TranslationModel: | |
| """ | |
| Wrapper for loading and running the EEG-to-fMRI translation model. | |
| """ | |
| def __init__(self, model_path: Optional[str] = None): | |
| self.device = torch.device("cpu") # CPU inference — fast enough for MLP | |
| self.model = EEGtoFMRIModel() | |
| if model_path and os.path.exists(model_path): | |
| checkpoint = torch.load(model_path, map_location=self.device, weights_only=False) | |
| self.model.load_state_dict(checkpoint["model_state_dict"]) | |
| self.feature_mean = checkpoint.get("feature_mean", None) | |
| self.feature_std = checkpoint.get("feature_std", None) | |
| print(f"Loaded trained model from {model_path}") | |
| else: | |
| self.feature_mean = None | |
| self.feature_std = None | |
| print("No trained model found — using untrained model") | |
| self.model.to(self.device) | |
| self.model.eval() | |
| def predict(self, features: np.ndarray) -> dict: | |
| """ | |
| Translate EEG features to fMRI BOLD estimates. | |
| Args: | |
| features: (236,) feature vector from eeg_processor | |
| Returns: | |
| dict with region activations and network summaries | |
| """ | |
| # Normalize features if we have training stats | |
| if self.feature_mean is not None and self.feature_std is not None: | |
| features = (features - self.feature_mean) / (self.feature_std + 1e-8) | |
| # Run model | |
| with torch.no_grad(): | |
| x = torch.FloatTensor(features).unsqueeze(0).to(self.device) | |
| bold = self.model(x).squeeze(0).cpu().numpy() | |
| # Build response | |
| regions = [] | |
| for i, (name, value) in enumerate(zip(HARVARD_OXFORD_REGIONS, bold)): | |
| # Find which network this region belongs to | |
| network = "Other" | |
| for net_name, indices in NETWORK_ASSIGNMENTS.items(): | |
| if i in indices: | |
| network = net_name | |
| break | |
| regions.append({ | |
| "index": i, | |
| "name": name, | |
| "bold_value": round(float(value), 4), | |
| "network": network, | |
| }) | |
| # Network summaries | |
| networks = {} | |
| for net_name, indices in NETWORK_ASSIGNMENTS.items(): | |
| values = [float(bold[i]) for i in indices if i < len(bold)] | |
| networks[net_name] = { | |
| "mean_activation": round(np.mean(values), 4) if values else 0, | |
| "max_activation": round(np.max(values), 4) if values else 0, | |
| "active_regions": sum(1 for v in values if v > 0.5), | |
| "total_regions": len(values), | |
| } | |
| # Top active regions | |
| sorted_regions = sorted(regions, key=lambda r: r["bold_value"], reverse=True) | |
| return { | |
| "bold_values": [round(float(v), 4) for v in bold], | |
| "regions": regions, | |
| "networks": networks, | |
| "top_regions": sorted_regions[:10], | |
| "global_mean_activation": round(float(np.mean(bold)), 4), | |
| "global_max_activation": round(float(np.max(bold)), 4), | |
| } | |