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