Spaces:
Sleeping
Sleeping
| """ | |
| Training script for EEG-to-fMRI translation model. | |
| Uses publicly available paired EEG-fMRI datasets to train the MLP | |
| to map EEG features → fMRI BOLD signals across 96 Harvard-Oxford regions. | |
| Public datasets used: | |
| 1. Simultaneous EEG-fMRI dataset from OpenNeuro (ds003505) | |
| - Resting-state EEG-fMRI from multiple subjects | |
| - EEG: 64 channels, 5000 Hz | |
| - fMRI: 3T, TR=2s, Harvard-Oxford parcellation | |
| Usage: | |
| python train.py --data_dir ./data --epochs 100 --output model_weights.pt | |
| If no real data is available, generates structured synthetic paired data | |
| that preserves known EEG-fMRI correlations (alpha-occipital, beta-motor, etc.) | |
| """ | |
| import argparse | |
| import os | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| from torch.utils.data import Dataset, DataLoader | |
| from model import EEGtoFMRIModel | |
| from eeg_processor import ( | |
| extract_all_features, | |
| get_feature_dimension, | |
| BANDS, | |
| BAND_NAMES, | |
| ) | |
| class EEGfMRIDataset(Dataset): | |
| """Dataset of paired EEG features and fMRI BOLD values.""" | |
| def __init__(self, features: np.ndarray, bold: np.ndarray): | |
| self.features = torch.FloatTensor(features) | |
| self.bold = torch.FloatTensor(bold) | |
| def __len__(self): | |
| return len(self.features) | |
| def __getitem__(self, idx): | |
| return self.features[idx], self.bold[idx] | |
| def generate_synthetic_paired_data( | |
| n_samples: int = 5000, | |
| n_channels: int = 8, | |
| sfreq: float = 256.0, | |
| n_regions: int = 96, | |
| epoch_duration: float = 2.0, | |
| ) -> tuple[np.ndarray, np.ndarray]: | |
| """ | |
| Generate synthetic paired EEG-fMRI data that preserves known | |
| neurophysiological relationships between EEG bands and brain regions. | |
| Known correlations used: | |
| - Alpha power (8-13 Hz) negatively correlates with visual cortex BOLD | |
| - Beta power (13-30 Hz) correlates with motor cortex BOLD | |
| - Theta power (4-8 Hz) correlates with hippocampal/default mode BOLD | |
| - Gamma power (30-45 Hz) correlates with attention/executive BOLD | |
| - Delta power (1-4 Hz) correlates with deep structures during rest | |
| - Cross-hemisphere coherence correlates with bilateral activation | |
| This is NOT a substitute for real data training — it's a structured | |
| initialization that produces physiologically plausible outputs while | |
| we collect/process real paired datasets. | |
| """ | |
| print(f"Generating {n_samples} synthetic paired EEG-fMRI samples...") | |
| n_samples_per_epoch = int(sfreq * epoch_duration) | |
| all_features = [] | |
| all_bold = [] | |
| # Define which regions respond to which EEG features | |
| # Based on known EEG-fMRI correlations from literature | |
| # Region indices follow Harvard-Oxford atlas ordering | |
| # Visual regions (occipital) | |
| visual_regions = [21, 22, 23, 31, 35, 38, 39, 46, 47, 69, 70, 71, 79, 83, 86, 87, 94, 95] | |
| # Motor regions | |
| motor_regions = [6, 16, 17, 25, 54, 64, 65, 73] | |
| # Language regions (temporal) | |
| language_regions = [4, 5, 8, 9, 10, 11, 18, 19, 20, 44, 45, 52, 53, 56, 57, 58, 59, 66, 67, 68, 92, 93] | |
| # Default mode network | |
| dmn_regions = [0, 24, 27, 28, 29, 30, 33, 34, 48, 72, 75, 76, 77, 78, 81, 82] | |
| # Executive/attention | |
| exec_regions = [1, 2, 3, 26, 32, 40, 41, 42, 49, 50, 51, 74, 80, 88, 89, 90] | |
| for sample_idx in range(n_samples): | |
| # Generate different "brain states" for diversity | |
| state = np.random.choice(["rest", "visual", "motor", "cognitive", "mixed"]) | |
| # Generate EEG with state-specific spectral properties | |
| eeg = np.zeros((n_channels, n_samples_per_epoch)) | |
| t = np.arange(n_samples_per_epoch) / sfreq | |
| # Base noise | |
| eeg += np.random.randn(n_channels, n_samples_per_epoch) * 0.5 | |
| # State-dependent oscillations | |
| alpha_power = 0.3 | |
| beta_power = 0.3 | |
| theta_power = 0.3 | |
| gamma_power = 0.2 | |
| delta_power = 0.2 | |
| if state == "rest": | |
| alpha_power = 0.8 + np.random.rand() * 0.4 | |
| theta_power = 0.5 + np.random.rand() * 0.3 | |
| delta_power = 0.4 + np.random.rand() * 0.2 | |
| elif state == "visual": | |
| alpha_power = 0.1 + np.random.rand() * 0.2 # Alpha suppression | |
| gamma_power = 0.6 + np.random.rand() * 0.4 | |
| elif state == "motor": | |
| beta_power = 0.7 + np.random.rand() * 0.5 | |
| alpha_power = 0.2 + np.random.rand() * 0.2 | |
| elif state == "cognitive": | |
| theta_power = 0.7 + np.random.rand() * 0.4 | |
| gamma_power = 0.5 + np.random.rand() * 0.3 | |
| alpha_power = 0.2 + np.random.rand() * 0.2 | |
| else: # mixed | |
| alpha_power = np.random.rand() * 0.8 | |
| beta_power = np.random.rand() * 0.8 | |
| theta_power = np.random.rand() * 0.8 | |
| gamma_power = np.random.rand() * 0.5 | |
| # Add oscillations to EEG channels with spatial distribution | |
| for ch in range(n_channels): | |
| # Each channel gets slightly different weights | |
| ch_weight = 0.7 + 0.6 * np.random.rand() | |
| # Posterior channels (6,7) get more alpha | |
| alpha_w = alpha_power * (1.5 if ch >= 6 else 0.7) * ch_weight | |
| # Central channels (2,3,4,5) get more beta | |
| beta_w = beta_power * (1.5 if 2 <= ch <= 5 else 0.7) * ch_weight | |
| # Frontal channels (0,1) get more theta | |
| theta_w = theta_power * (1.5 if ch <= 1 else 0.7) * ch_weight | |
| freq_alpha = 10 + np.random.randn() * 0.5 | |
| freq_beta = 20 + np.random.randn() * 2 | |
| freq_theta = 6 + np.random.randn() * 0.5 | |
| freq_gamma = 40 + np.random.randn() * 3 | |
| freq_delta = 2 + np.random.randn() * 0.3 | |
| eeg[ch] += alpha_w * np.sin(2 * np.pi * freq_alpha * t + np.random.rand() * 2 * np.pi) | |
| eeg[ch] += beta_w * np.sin(2 * np.pi * freq_beta * t + np.random.rand() * 2 * np.pi) | |
| eeg[ch] += theta_w * np.sin(2 * np.pi * freq_theta * t + np.random.rand() * 2 * np.pi) | |
| eeg[ch] += gamma_power * ch_weight * np.sin(2 * np.pi * freq_gamma * t + np.random.rand() * 2 * np.pi) | |
| eeg[ch] += delta_power * ch_weight * np.sin(2 * np.pi * freq_delta * t + np.random.rand() * 2 * np.pi) | |
| # Extract real features from this synthetic EEG | |
| features = extract_all_features(eeg, sfreq, target_channels=n_channels) | |
| # Generate corresponding BOLD values based on known correlations | |
| bold = np.ones(n_regions) * 0.3 + np.random.randn(n_regions) * 0.05 | |
| # Visual cortex: negatively correlated with alpha, positively with gamma | |
| for r in visual_regions: | |
| if r < n_regions: | |
| bold[r] = 0.3 + (1.0 - alpha_power) * 0.4 + gamma_power * 0.2 + np.random.randn() * 0.05 | |
| # Motor cortex: correlated with beta | |
| for r in motor_regions: | |
| if r < n_regions: | |
| bold[r] = 0.3 + beta_power * 0.5 + np.random.randn() * 0.05 | |
| # Language regions: mix of theta and gamma | |
| for r in language_regions: | |
| if r < n_regions: | |
| bold[r] = 0.3 + theta_power * 0.2 + gamma_power * 0.3 + np.random.randn() * 0.05 | |
| # Default mode: correlated with alpha and theta, anti-correlated with task | |
| for r in dmn_regions: | |
| if r < n_regions: | |
| bold[r] = 0.3 + alpha_power * 0.3 + theta_power * 0.2 - gamma_power * 0.15 + np.random.randn() * 0.05 | |
| # Executive: correlated with gamma and theta, anti-correlated with alpha | |
| for r in exec_regions: | |
| if r < n_regions: | |
| bold[r] = 0.3 + gamma_power * 0.35 + theta_power * 0.15 - alpha_power * 0.1 + np.random.randn() * 0.05 | |
| # Clip to [0, 1] | |
| bold = np.clip(bold, 0.05, 0.95) | |
| all_features.append(features) | |
| all_bold.append(bold) | |
| features_array = np.array(all_features) | |
| bold_array = np.array(all_bold) | |
| print(f"Generated features shape: {features_array.shape}") | |
| print(f"Generated BOLD shape: {bold_array.shape}") | |
| return features_array, bold_array | |
| def load_real_data(data_dir: str): | |
| """ | |
| Load real paired EEG-fMRI data from preprocessed files. | |
| Expected format: | |
| data_dir/features.npy — (n_samples, 236) EEG feature vectors | |
| data_dir/bold.npy — (n_samples, 96) fMRI BOLD values | |
| To create these files from raw data: | |
| 1. Download dataset from OpenNeuro (e.g., ds003505) | |
| 2. Run preprocessing with MNE (EEG) and nilearn (fMRI) | |
| 3. Extract EEG features using eeg_processor.extract_all_features() | |
| 4. Extract fMRI ROI values using Harvard-Oxford atlas | |
| 5. Save as .npy files | |
| """ | |
| features_path = os.path.join(data_dir, "features.npy") | |
| bold_path = os.path.join(data_dir, "bold.npy") | |
| if os.path.exists(features_path) and os.path.exists(bold_path): | |
| features = np.load(features_path) | |
| bold = np.load(bold_path) | |
| print(f"Loaded real data: {features.shape[0]} samples") | |
| print(f" Features: {features.shape}") | |
| print(f" BOLD: {bold.shape}") | |
| return features, bold | |
| return None | |
| def train( | |
| features: np.ndarray, | |
| bold: np.ndarray, | |
| epochs: int = 100, | |
| batch_size: int = 64, | |
| lr: float = 1e-3, | |
| val_split: float = 0.15, | |
| output_path: str = "model_weights.pt", | |
| ): | |
| """Train the EEG-to-fMRI translation model.""" | |
| # Normalize features | |
| feature_mean = features.mean(axis=0) | |
| feature_std = features.std(axis=0) | |
| feature_std[feature_std < 1e-8] = 1.0 | |
| features_norm = (features - feature_mean) / feature_std | |
| # Train/val split | |
| n = len(features_norm) | |
| n_val = int(n * val_split) | |
| indices = np.random.permutation(n) | |
| train_idx = indices[n_val:] | |
| val_idx = indices[:n_val] | |
| train_dataset = EEGfMRIDataset(features_norm[train_idx], bold[train_idx]) | |
| val_dataset = EEGfMRIDataset(features_norm[val_idx], bold[val_idx]) | |
| train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) | |
| val_loader = DataLoader(val_dataset, batch_size=batch_size) | |
| print(f"Training: {len(train_dataset)} samples") | |
| print(f"Validation: {len(val_dataset)} samples") | |
| print(f"Feature dim: {features.shape[1]}") | |
| print(f"Output dim: {bold.shape[1]}") | |
| # Model | |
| model = EEGtoFMRIModel(input_dim=features.shape[1], output_dim=bold.shape[1]) | |
| device = torch.device("cpu") | |
| model.to(device) | |
| # Loss and optimizer | |
| criterion = nn.MSELoss() | |
| optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4) | |
| scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs) | |
| best_val_loss = float("inf") | |
| patience = 15 | |
| patience_counter = 0 | |
| for epoch in range(epochs): | |
| # Training | |
| model.train() | |
| train_loss = 0 | |
| for batch_features, batch_bold in train_loader: | |
| batch_features = batch_features.to(device) | |
| batch_bold = batch_bold.to(device) | |
| pred = model(batch_features) | |
| loss = criterion(pred, batch_bold) | |
| optimizer.zero_grad() | |
| loss.backward() | |
| torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) | |
| optimizer.step() | |
| train_loss += loss.item() | |
| train_loss /= len(train_loader) | |
| # Validation | |
| model.eval() | |
| val_loss = 0 | |
| with torch.no_grad(): | |
| for batch_features, batch_bold in val_loader: | |
| batch_features = batch_features.to(device) | |
| batch_bold = batch_bold.to(device) | |
| pred = model(batch_features) | |
| loss = criterion(pred, batch_bold) | |
| val_loss += loss.item() | |
| val_loss /= len(val_loader) | |
| scheduler.step() | |
| if (epoch + 1) % 10 == 0 or epoch == 0: | |
| print( | |
| f"Epoch {epoch+1:3d}/{epochs} | " | |
| f"Train Loss: {train_loss:.6f} | " | |
| f"Val Loss: {val_loss:.6f} | " | |
| f"LR: {scheduler.get_last_lr()[0]:.2e}" | |
| ) | |
| # Early stopping | |
| if val_loss < best_val_loss: | |
| best_val_loss = val_loss | |
| patience_counter = 0 | |
| # Save best model | |
| torch.save( | |
| { | |
| "model_state_dict": model.state_dict(), | |
| "feature_mean": feature_mean, | |
| "feature_std": feature_std, | |
| "val_loss": best_val_loss, | |
| "epoch": epoch, | |
| "input_dim": features.shape[1], | |
| "output_dim": bold.shape[1], | |
| }, | |
| output_path, | |
| ) | |
| else: | |
| patience_counter += 1 | |
| if patience_counter >= patience: | |
| print(f"Early stopping at epoch {epoch+1}") | |
| break | |
| print(f"\nBest validation loss: {best_val_loss:.6f}") | |
| print(f"Model saved to {output_path}") | |
| # Compute per-region correlation on validation set | |
| model.load_state_dict( | |
| torch.load(output_path, map_location=device, weights_only=False)["model_state_dict"] | |
| ) | |
| model.eval() | |
| all_preds = [] | |
| all_true = [] | |
| with torch.no_grad(): | |
| for batch_features, batch_bold in val_loader: | |
| pred = model(batch_features.to(device)).cpu().numpy() | |
| all_preds.append(pred) | |
| all_true.append(batch_bold.numpy()) | |
| all_preds = np.vstack(all_preds) | |
| all_true = np.vstack(all_true) | |
| # Per-region Pearson correlation | |
| from scipy.stats import pearsonr | |
| correlations = [] | |
| for r in range(bold.shape[1]): | |
| if np.std(all_preds[:, r]) > 0 and np.std(all_true[:, r]) > 0: | |
| corr, _ = pearsonr(all_preds[:, r], all_true[:, r]) | |
| correlations.append(corr) | |
| else: | |
| correlations.append(0.0) | |
| mean_corr = np.mean(correlations) | |
| print(f"Mean per-region correlation: {mean_corr:.4f}") | |
| print(f"Top 5 regions by correlation:") | |
| top_indices = np.argsort(correlations)[-5:][::-1] | |
| for idx in top_indices: | |
| from model import HARVARD_OXFORD_REGIONS | |
| print(f" {HARVARD_OXFORD_REGIONS[idx]}: r={correlations[idx]:.4f}") | |
| return model | |
| def main(): | |
| parser = argparse.ArgumentParser(description="Train EEG-to-fMRI translation model") | |
| parser.add_argument("--data_dir", type=str, default="./data", help="Directory with paired data") | |
| parser.add_argument("--epochs", type=int, default=100, help="Training epochs") | |
| parser.add_argument("--batch_size", type=int, default=64, help="Batch size") | |
| parser.add_argument("--lr", type=float, default=1e-3, help="Learning rate") | |
| parser.add_argument("--n_synthetic", type=int, default=5000, help="Number of synthetic samples if no real data") | |
| parser.add_argument("--output", type=str, default="model_weights.pt", help="Output model path") | |
| args = parser.parse_args() | |
| # Try to load real data first | |
| real_data = load_real_data(args.data_dir) | |
| if real_data is not None: | |
| features, bold = real_data | |
| print("Training on REAL paired EEG-fMRI data") | |
| else: | |
| print("No real data found — generating structured synthetic data") | |
| print("(To train on real data, place features.npy and bold.npy in data_dir)") | |
| features, bold = generate_synthetic_paired_data(n_samples=args.n_synthetic) | |
| train( | |
| features, | |
| bold, | |
| epochs=args.epochs, | |
| batch_size=args.batch_size, | |
| lr=args.lr, | |
| output_path=args.output, | |
| ) | |
| if __name__ == "__main__": | |
| main() | |