""" Multimodal Brain Encoder - Gradio Application ============================================= Full end-to-end system: Input → CLIP Features → Brain Prediction → ROI Analysis → LLM Q&A → Visualization Uses real trained weights from NSD dataset. LLM is an INTERPRETER only - grounded in model predictions, not independent. """ import os import sys import json import time import logging import pickle import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from PIL import Image from pathlib import Path from datetime import datetime from collections import OrderedDict logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # ============================================================ # Configuration (must match training) # ============================================================ MODEL_REPO = os.environ.get("MODEL_REPO", "ryu34/multimodal-brain-encoder") ROI_NAMES = { 1: "V1v", 2: "V1d", 3: "V2v", 4: "V2d", 5: "V3v", 6: "V3d", 7: "hV4", 8: "EBA", 9: "FBA-1", 10: "FBA-2", 11: "mTL-bodies", 12: "OFA", 13: "FFA-1", 14: "FFA-2", 15: "mTL-faces", 16: "aTL-faces", 17: "OPA", 18: "PPA", 19: "RSC", 20: "OWFA", 21: "VWFA-1", 22: "VWFA-2", 23: "mfs-words", 24: "mTL-words", } FUNCTIONAL_NETWORKS = { "early_visual": [1, 2, 3, 4, 5, 6, 7], "body_selective": [8, 9, 10, 11], "face_selective": [12, 13, 14, 15, 16], "place_selective": [17, 18, 19], "word_selective": [20, 21, 22, 23, 24], } # Known neuroscience associations for grounded Q&A ROI_FUNCTIONS = { "V1v": "Primary visual cortex (ventral); processes basic visual features: edges, orientations, spatial frequencies", "V1d": "Primary visual cortex (dorsal); processes basic visual features with dorsal visual stream emphasis", "V2v": "Secondary visual cortex (ventral); processes texture, figure-ground segregation", "V2d": "Secondary visual cortex (dorsal); processes contour and border ownership", "V3v": "Third visual area (ventral); contributes to form perception and shape processing", "V3d": "Third visual area (dorsal); processes dynamic form and motion boundaries", "hV4": "Human V4; processes color, pattern, moderate object features, texture discrimination", "EBA": "Extrastriate Body Area; selectively responds to bodies and body parts", "FBA-1": "Fusiform Body Area 1; body processing in ventral temporal cortex", "FBA-2": "Fusiform Body Area 2; complementary body processing region", "mTL-bodies": "Medial temporal lobe body area; body recognition with memory component", "OFA": "Occipital Face Area; early face-selective processing, face parts detection", "FFA-1": "Fusiform Face Area 1; core face recognition and identity processing", "FFA-2": "Fusiform Face Area 2; complementary face processing, holistic face representation", "mTL-faces": "Medial temporal lobe face area; face recognition with episodic memory", "aTL-faces": "Anterior temporal lobe face area; person identity and semantic knowledge", "OPA": "Occipital Place Area; processes local scene elements and spatial boundaries", "PPA": "Parahippocampal Place Area; processes scenes, buildings, spatial layouts", "RSC": "Retrosplenial Cortex; spatial navigation, scene-to-map coordinate transformation", "OWFA": "Occipital Word Form Area; early visual word processing", "VWFA-1": "Visual Word Form Area 1; processes written words and letter strings", "VWFA-2": "Visual Word Form Area 2; higher-level word form processing", "mfs-words": "Mid-fusiform sulcus word area; intermediate word processing", "mTL-words": "Medial temporal lobe word area; word recognition with memory", } NETWORK_FUNCTIONS = { "early_visual": "Early visual processing: edges, orientations, spatial frequencies, textures, colors. Active for all visual stimuli.", "body_selective": "Body-selective cortex: responds to human bodies, body parts, biological motion. Key for person perception.", "face_selective": "Face-selective cortex: responds to faces, facial features, identity. Critical for social perception.", "place_selective": "Place/scene-selective cortex: responds to spatial layouts, buildings, scenes, navigation cues.", "word_selective": "Word/reading-selective cortex: responds to written text, letter strings, word forms.", } # ============================================================ # Helper: enable only Dropout for MC sampling (keep BatchNorm in eval) # ============================================================ def enable_dropout_only(model): """Enable Dropout layers while keeping BatchNorm in eval mode. This is needed for MC Dropout uncertainty estimation with batch_size=1, because BatchNorm1d requires batch_size > 1 in training mode. """ for module in model.modules(): if isinstance(module, nn.Dropout): module.train() # ============================================================ # BrainEncoder model (must match training architecture exactly) # ============================================================ class BrainEncoder(nn.Module): def __init__(self, input_dim=4096, n_voxels=15724, hidden_dims=None, dropout=0.3, n_rois=24): super().__init__() if hidden_dims is None: hidden_dims = [2048, 2048, 1024] self.input_dim = input_dim self.n_voxels = n_voxels self.n_rois = n_rois layers = [] prev_dim = input_dim for h_dim in hidden_dims: layers.extend([ nn.Linear(prev_dim, h_dim), nn.BatchNorm1d(h_dim), nn.GELU(), nn.Dropout(dropout), ]) prev_dim = h_dim self.backbone = nn.Sequential(*layers) self.general_head = nn.Linear(hidden_dims[-1], n_voxels) self.roi_attention = nn.ModuleDict() self.roi_heads = nn.ModuleDict() self.network_names = ["early_visual", "body_selective", "face_selective", "place_selective", "word_selective"] for net_name in self.network_names: self.roi_attention[net_name] = nn.Sequential( nn.Linear(hidden_dims[-1], 256), nn.GELU(), nn.Linear(256, hidden_dims[-1]), nn.Sigmoid(), ) self.roi_heads[net_name] = nn.Linear(hidden_dims[-1], n_voxels) self.register_buffer('roi_mask', torch.zeros(n_voxels, dtype=torch.long)) def set_roi_assignments(self, annot): for net_idx, (net_name, roi_ids) in enumerate(FUNCTIONAL_NETWORKS.items()): for roi_id in roi_ids: mask = (annot == roi_id) if len(mask) <= self.n_voxels: self.roi_mask[:len(mask)][mask[:self.n_voxels]] = net_idx + 1 def forward(self, x, return_intermediates=False): intermediates = {} backbone_out = self.backbone(x) intermediates['backbone'] = backbone_out.detach() pred = self.general_head(backbone_out) intermediates['general_pred'] = pred.detach() for net_idx, net_name in enumerate(self.network_names): if net_name in self.roi_attention: attn = self.roi_attention[net_name](backbone_out) weighted = backbone_out * attn roi_pred = self.roi_heads[net_name](weighted) mask = (self.roi_mask == net_idx + 1) if mask.any(): pred[:, mask] = roi_pred[:, mask] intermediates[f'roi_{net_name}'] = roi_pred.detach() if return_intermediates: return pred, intermediates return pred # ============================================================ # Model Manager - loads and caches models # ============================================================ class ModelManager: def __init__(self): self.device = "cuda" if torch.cuda.is_available() else "cpu" self.brain_encoder = None self.ridge_model = None self.clip_model = None self.clip_processor = None self.roi_annotations = None self.config = None self._loaded = False def load(self): if self._loaded: return from huggingface_hub import hf_hub_download logger.info(f"Loading models from {MODEL_REPO}...") # Load config try: config_path = hf_hub_download(repo_id=MODEL_REPO, filename="config.json") with open(config_path) as f: self.config = json.load(f) logger.info(f"Config loaded: {self.config.get('architecture', 'unknown')}") except Exception as e: logger.warning(f"Config load failed: {e}") self.config = {} # Load ROI annotations try: annot_path = hf_hub_download(repo_id=MODEL_REPO, filename="roi_annotations.npy") self.roi_annotations = np.load(annot_path).flatten() logger.info(f"ROI annotations: {self.roi_annotations.shape}") except Exception as e: logger.warning(f"ROI annotations load failed: {e}") # Load brain encoder (optional - ridge is primary) try: model_path = hf_hub_download(repo_id=MODEL_REPO, filename="best_model.pt") checkpoint = torch.load(model_path, map_location=self.device, weights_only=False) model_config = checkpoint.get('config', {}) self.brain_encoder = BrainEncoder( input_dim=model_config.get('input_dim', 4096), n_voxels=model_config.get('n_voxels', 15724), hidden_dims=model_config.get('hidden_dims', [2048, 2048, 1024]), dropout=model_config.get('dropout', 0.3), ) self.brain_encoder.load_state_dict(checkpoint['model_state_dict']) self.brain_encoder.to(self.device).eval() if self.roi_annotations is not None: self.brain_encoder.set_roi_assignments(self.roi_annotations) # Free checkpoint memory del checkpoint logger.info("Brain encoder loaded successfully") except Exception as e: logger.warning(f"Brain encoder load failed (will use ridge only): {e}") self.brain_encoder = None # Load ridge model try: ridge_path = hf_hub_download(repo_id=MODEL_REPO, filename="ridge_model.pkl") with open(ridge_path, 'rb') as f: self.ridge_model = pickle.load(f) logger.info("Ridge model loaded successfully") except Exception as e: logger.warning(f"Ridge model load failed: {e}") # Load CLIP try: from transformers import CLIPModel, CLIPProcessor self.clip_model = CLIPModel.from_pretrained( "openai/clip-vit-large-patch14", torch_dtype=torch.float32, ).to(self.device).eval() self.clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14") logger.info("CLIP model loaded") except Exception as e: logger.error(f"CLIP load failed: {e}") raise self._loaded = True logger.info("All models loaded successfully!") def extract_features(self, image=None, text=None, audio=None): """Extract multimodal CLIP features.""" features_dict = {} if image is not None: if isinstance(image, np.ndarray): image = Image.fromarray(image) inputs = self.clip_processor(images=image, return_tensors="pt") inputs = {k: v.to(self.device) for k, v in inputs.items() if isinstance(v, torch.Tensor)} with torch.no_grad(): outputs = self.clip_model.vision_model(**inputs, output_hidden_states=True) cls_features = outputs.last_hidden_state[:, 0, :] projected = self.clip_model.visual_projection(cls_features) hidden_concat = [] for layer_idx in [6, 12, 18, 23]: h = outputs.hidden_states[layer_idx][:, 0, :] hidden_concat.append(h) multi_layer = torch.cat(hidden_concat, dim=-1) features_dict['image_projected'] = projected.cpu().float() features_dict['image_multi_layer'] = multi_layer.cpu().float() features_dict['modality'] = 'image' if text is not None and text.strip(): inputs = self.clip_processor(text=[text], return_tensors="pt", padding=True, truncation=True) inputs = {k: v.to(self.device) for k, v in inputs.items() if isinstance(v, torch.Tensor)} with torch.no_grad(): text_outputs = self.clip_model.text_model(**inputs) pooled = text_outputs.pooler_output projected = self.clip_model.text_projection(pooled) # For text, repeat the projected features to match multi-layer dim # Text goes through the same brain encoder by tiling to 4096 text_multi = projected.repeat(1, 4096 // projected.shape[1] + 1)[:, :4096] features_dict['text_projected'] = projected.cpu().float() features_dict['text_multi_layer'] = text_multi.cpu().float() if 'modality' not in features_dict: features_dict['modality'] = 'text' else: features_dict['modality'] = 'image+text' if audio is not None: # Convert audio to spectrogram image for CLIP processing sr, audio_data = audio if isinstance(audio, tuple) else (16000, audio) if len(audio_data.shape) > 1: audio_data = audio_data.mean(axis=1) audio_data = audio_data.astype(np.float32) # Create mel spectrogram visualization import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt fig, ax = plt.subplots(1, 1, figsize=(4, 4)) # Simple spectrogram using STFT n_fft = min(1024, len(audio_data)) hop_length = n_fft // 4 if len(audio_data) > n_fft: # Manual STFT n_frames = (len(audio_data) - n_fft) // hop_length + 1 spec = np.zeros((n_fft // 2 + 1, n_frames)) window = np.hanning(n_fft) for i in range(n_frames): start = i * hop_length frame = audio_data[start:start + n_fft] * window fft = np.fft.rfft(frame) spec[:, i] = np.abs(fft) spec_db = 20 * np.log10(spec + 1e-10) ax.imshow(spec_db, aspect='auto', origin='lower', cmap='viridis') else: ax.plot(audio_data[:1000]) ax.set_title("Audio Spectrogram") ax.axis('off') fig.canvas.draw() # Convert to image buf = fig.canvas.buffer_rgba() spec_img = Image.frombytes('RGBA', fig.canvas.get_width_height(), buf).convert('RGB') plt.close(fig) # Process through CLIP as image inputs = self.clip_processor(images=spec_img, return_tensors="pt") inputs = {k: v.to(self.device) for k, v in inputs.items() if isinstance(v, torch.Tensor)} with torch.no_grad(): outputs = self.clip_model.vision_model(**inputs, output_hidden_states=True) cls_features = outputs.last_hidden_state[:, 0, :] projected = self.clip_model.visual_projection(cls_features) hidden_concat = [] for layer_idx in [6, 12, 18, 23]: h = outputs.hidden_states[layer_idx][:, 0, :] hidden_concat.append(h) multi_layer = torch.cat(hidden_concat, dim=-1) features_dict['audio_projected'] = projected.cpu().float() features_dict['audio_multi_layer'] = multi_layer.cpu().float() if features_dict.get('modality') is None: features_dict['modality'] = 'audio' else: features_dict['modality'] = features_dict['modality'] + '+audio' return features_dict def predict_brain_activity(self, features_dict): """Run brain encoder forward pass using BOTH ridge and deep models.""" # Determine which features to use if 'image_multi_layer' in features_dict: input_features = features_dict['image_multi_layer'] elif 'text_multi_layer' in features_dict: input_features = features_dict['text_multi_layer'] elif 'audio_multi_layer' in features_dict: input_features = features_dict['audio_multi_layer'] else: raise ValueError("No features available for prediction") # If multimodal, average features all_modality_features = [] for key in ['image_multi_layer', 'text_multi_layer', 'audio_multi_layer']: if key in features_dict: all_modality_features.append(features_dict[key]) if len(all_modality_features) > 1: input_features = torch.mean(torch.stack(all_modality_features), dim=0) input_features_np = input_features.cpu().numpy() input_features = input_features.to(self.device) # ── Primary: Ridge Model (proven baseline from Algonauts 2023) ── if self.ridge_model is not None: ridge = self.ridge_model X_norm = (input_features_np - ridge['feat_mean']) / ridge['feat_std'] pred_z = ridge['model'].predict(X_norm) pred_np = (pred_z * ridge['fmri_std'] + ridge['fmri_mean']).flatten() # Clip extreme values for better visualization (keep 99.5th percentile) clip_val = np.percentile(np.abs(pred_np), 99.5) pred_np = np.clip(pred_np, -clip_val, clip_val) else: # Fallback to deep encoder with torch.no_grad(): predictions, _ = self.brain_encoder(input_features, return_intermediates=True) pred_np = predictions.cpu().numpy().flatten() # ── Deep encoder for intermediates and uncertainty ── intermediates = {} if self.brain_encoder is not None: with torch.no_grad(): deep_pred, intermediates = self.brain_encoder(input_features, return_intermediates=True) # Compute uncertainty via MC Dropout # IMPORTANT: Only enable Dropout layers, keep BatchNorm in eval mode. # BatchNorm1d requires batch_size > 1 in training mode, but we have batch_size=1. self.brain_encoder.eval() # Ensure everything is in eval mode first enable_dropout_only(self.brain_encoder) # Selectively enable only Dropout mc_predictions = [] for _ in range(10): with torch.no_grad(): mc_pred = self.brain_encoder(input_features) mc_predictions.append(mc_pred.cpu().numpy().flatten()) self.brain_encoder.eval() # Restore full eval mode mc_predictions = np.array(mc_predictions) uncertainty = np.std(mc_predictions, axis=0) else: # Estimate uncertainty from ridge prediction variance across feature perturbation ridge = self.ridge_model mc_predictions = [] for _ in range(10): noise = np.random.normal(0, 0.01, size=input_features_np.shape) X_noisy = (input_features_np + noise - ridge['feat_mean']) / ridge['feat_std'] mp = ridge['model'].predict(X_noisy).flatten() mc_predictions.append(mp) mc_predictions = np.array(mc_predictions) uncertainty = np.std(mc_predictions, axis=0) # Compute modality contributions using ridge modality_contributions = {} if self.ridge_model is not None: ridge = self.ridge_model for key in ['image_multi_layer', 'text_multi_layer', 'audio_multi_layer']: if key in features_dict: modality_name = key.split('_')[0] feat_np = features_dict[key].cpu().numpy() X_n = (feat_np - ridge['feat_mean']) / ridge['feat_std'] mp = (ridge['model'].predict(X_n) * ridge['fmri_std'] + ridge['fmri_mean']).flatten() clip_val_mod = np.percentile(np.abs(mp), 99.5) mp = np.clip(mp, -clip_val_mod, clip_val_mod) modality_contributions[modality_name] = mp # Compute ROI summaries using z-scored per-voxel predictions # This shows which regions are MORE or LESS activated compared to baseline if self.ridge_model is not None: baseline_mean = self.ridge_model['fmri_mean'] baseline_std = self.ridge_model['fmri_std'] # Z-score predictions relative to training distribution n_v = min(len(pred_np), len(baseline_mean)) pred_z = (pred_np[:n_v] - baseline_mean[:n_v]) / (baseline_std[:n_v] + 1e-8) else: pred_z = pred_np roi_summary = self._compute_roi_summary(pred_z, uncertainty) # Validation checks warnings = self._validate_predictions(pred_np) result = { 'predictions': pred_np, 'uncertainty': uncertainty, 'roi_summary': roi_summary, 'modality_contributions': modality_contributions, 'modality': features_dict.get('modality', 'unknown'), 'intermediates': {k: v.cpu().numpy() if torch.is_tensor(v) else v for k, v in intermediates.items()}, 'warnings': warnings, 'timestamp': datetime.now().isoformat(), } return result def _compute_roi_summary(self, predictions, uncertainty): """Compute per-ROI activation summaries.""" if self.roi_annotations is None: return {} annot = self.roi_annotations n_voxels = len(predictions) roi_summary = {} for roi_id, roi_name in ROI_NAMES.items(): mask = (annot[:n_voxels] == roi_id) if len(annot) >= n_voxels else np.zeros(n_voxels, dtype=bool) if mask.sum() == 0: continue roi_activations = predictions[mask] roi_uncertainty = uncertainty[mask] roi_summary[roi_name] = { 'mean_activation': float(np.mean(roi_activations)), 'max_activation': float(np.max(roi_activations)), 'min_activation': float(np.min(roi_activations)), 'std_activation': float(np.std(roi_activations)), 'mean_uncertainty': float(np.mean(roi_uncertainty)), 'n_voxels': int(mask.sum()), 'abs_mean': float(np.mean(np.abs(roi_activations))), 'known_function': ROI_FUNCTIONS.get(roi_name, "Unknown"), } return roi_summary def _validate_predictions(self, predictions): """Validation safeguards.""" warnings = [] if np.std(predictions) < 1e-6: warnings.append("⚠️ CONSTANT OUTPUT DETECTED: All voxels have near-identical values") if np.any(np.isnan(predictions)): warnings.append("⚠️ NaN VALUES DETECTED in predictions") if np.any(np.isinf(predictions)): warnings.append("⚠️ Infinite VALUES DETECTED in predictions") if np.max(np.abs(predictions)) > 50: warnings.append(f"⚠️ Unusually large activations detected (max |activation| = {np.max(np.abs(predictions)):.2f})") return warnings # ============================================================ # Grounded Q&A System # ============================================================ class GroundedQA: """ RAG-grounded Q&A system. The LLM is an INTERPRETER - it only explains model predictions. It does NOT generate independent neuroscience claims. """ def __init__(self): self.inference_client = None self._init_client() def _init_client(self): try: from huggingface_hub import InferenceClient self.inference_client = InferenceClient( provider="hf-inference", api_key=os.environ.get("HF_TOKEN", ""), ) logger.info("HF Inference Client initialized") except Exception as e: logger.warning(f"Inference client init failed: {e}") def build_context(self, brain_result): """Build structured context from model predictions for LLM grounding.""" roi_summary = brain_result.get('roi_summary', {}) modality = brain_result.get('modality', 'unknown') warnings = brain_result.get('warnings', []) modality_contributions = brain_result.get('modality_contributions', {}) # Sort ROIs by absolute mean activation sorted_rois = sorted( roi_summary.items(), key=lambda x: abs(x[1]['abs_mean']), reverse=True ) # Top activated regions top_regions = [] for roi_name, data in sorted_rois[:10]: top_regions.append( f"- {roi_name}: mean_activation={data['mean_activation']:.4f}, " f"abs_mean={data['abs_mean']:.4f}, uncertainty={data['mean_uncertainty']:.4f}, " f"n_voxels={data['n_voxels']}" ) # Network-level summaries network_summaries = {} for net_name, roi_ids in FUNCTIONAL_NETWORKS.items(): roi_names_in_net = [ROI_NAMES[rid] for rid in roi_ids if rid in ROI_NAMES] activations = [] for rn in roi_names_in_net: if rn in roi_summary: activations.append(roi_summary[rn]['abs_mean']) if activations: network_summaries[net_name] = { 'mean_abs_activation': np.mean(activations), 'max_abs_activation': np.max(activations), 'function': NETWORK_FUNCTIONS.get(net_name, ""), } sorted_networks = sorted( network_summaries.items(), key=lambda x: x[1]['mean_abs_activation'], reverse=True ) # Modality contributions modality_info = "" if modality_contributions: modality_info = "\n## Modality Contributions\n" for mod_name, mod_pred in modality_contributions.items(): modality_info += f"- {mod_name}: mean_abs_activation={np.mean(np.abs(mod_pred)):.4f}, std={np.std(mod_pred):.4f}\n" # Global prediction stats predictions = brain_result['predictions'] global_stats = ( f"- Total voxels predicted: {len(predictions)}\n" f"- Global mean activation: {np.mean(predictions):.4f}\n" f"- Global std: {np.std(predictions):.4f}\n" f"- Global range: [{np.min(predictions):.4f}, {np.max(predictions):.4f}]\n" f"- Mean uncertainty: {np.mean(brain_result['uncertainty']):.4f}\n" ) context = f"""## Brain Activity Prediction Summary Input modality: {modality} ## Global Statistics {global_stats} ## Top 10 Activated Brain Regions (by absolute activation strength) {chr(10).join(top_regions)} ## Functional Network Activations (ranked by strength) """ for net_name, net_data in sorted_networks: context += ( f"- {net_name}: mean_abs={net_data['mean_abs_activation']:.4f}, " f"max_abs={net_data['max_abs_activation']:.4f}\n" f" Known function: {net_data['function']}\n" ) context += modality_info if warnings: context += "\n## Warnings\n" for w in warnings: context += f"- {w}\n" # ROI functional labels context += "\n## ROI Functional Reference\n" for roi_name in [r[0] for r in sorted_rois[:10]]: if roi_name in ROI_FUNCTIONS: context += f"- {roi_name}: {ROI_FUNCTIONS[roi_name]}\n" return context def answer(self, question, brain_result): """Answer a question grounded in model predictions.""" context = self.build_context(brain_result) system_prompt = """You are a neuroscience interpreter for a brain encoding model. Your role is STRICTLY to interpret and explain the model's predicted brain activity patterns. CRITICAL RULES: 1. ONLY reference data provided in the context below. Never invent neuroscience claims. 2. Always distinguish between: - "Predicted activation" (what the model outputs) - "Known neuroscience association" (established findings about brain regions) - "Possible interpretation" (your inference connecting the two) 3. Include uncertainty statements. Use phrases like "the model predicts", "this is consistent with", "one possible interpretation is" 4. NEVER make definitive claims about emotions, consciousness, or behavior from brain activity alone. 5. Always cite specific regions, activation values, and confidence levels from the context. 6. If the question cannot be answered from the provided data, say so explicitly. 7. Keep answers concise but informative (2-4 paragraphs max). You are an INTERPRETER of model outputs, not an independent neuroscience oracle.""" user_prompt = f"""## Model Prediction Context {context} ## User Question {question} Please answer based ONLY on the model prediction data above. Cite specific regions and values.""" if self.inference_client is None: return self._fallback_answer(question, brain_result, context) try: response = self.inference_client.chat.completions.create( model="Qwen/Qwen2.5-72B-Instruct", messages=[ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt}, ], max_tokens=800, temperature=0.3, ) answer = response.choices[0].message.content # Add grounding footer answer += "\n\n---\n*This interpretation is based on model predictions with " mean_unc = np.mean(brain_result['uncertainty']) answer += f"mean uncertainty={mean_unc:.4f}. " answer += "Predictions are from a brain encoder trained on NSD (Natural Scenes Dataset) fMRI data.*" return answer except Exception as e: logger.warning(f"LLM inference failed: {e}") return self._fallback_answer(question, brain_result, context) def _fallback_answer(self, question, brain_result, context): """Structured fallback when LLM is unavailable.""" roi_summary = brain_result.get('roi_summary', {}) sorted_rois = sorted( roi_summary.items(), key=lambda x: abs(x[1]['abs_mean']), reverse=True ) answer = "## Brain Activity Interpretation\n\n" answer += f"**Input modality:** {brain_result.get('modality', 'unknown')}\n\n" answer += "### Top Activated Regions\n" for roi_name, data in sorted_rois[:5]: answer += ( f"- **{roi_name}** (activation={data['mean_activation']:.4f}, " f"uncertainty={data['mean_uncertainty']:.4f}): " f"{ROI_FUNCTIONS.get(roi_name, 'Unknown function')}\n" ) answer += "\n### Network-Level Summary\n" for net_name, roi_ids in FUNCTIONAL_NETWORKS.items(): roi_names_in_net = [ROI_NAMES[rid] for rid in roi_ids if rid in ROI_NAMES] activations = [roi_summary[rn]['abs_mean'] for rn in roi_names_in_net if rn in roi_summary] if activations: mean_act = np.mean(activations) answer += f"- **{net_name}**: mean_abs_activation={mean_act:.4f} — {NETWORK_FUNCTIONS.get(net_name, '')}\n" answer += f"\n*Note: LLM interpretation unavailable. Showing structured prediction summary. " answer += f"Mean uncertainty: {np.mean(brain_result['uncertainty']):.4f}*" return answer # ============================================================ # Transparency Logger # ============================================================ class TransparencyLogger: """Logs all inputs, intermediates, and outputs for traceability.""" def __init__(self): self.logs = [] def log_inference(self, inputs, features_dict, brain_result, qa_answer=None): entry = { 'timestamp': datetime.now().isoformat(), 'inputs': { 'has_image': inputs.get('image') is not None, 'has_text': inputs.get('text') is not None and inputs.get('text', '').strip() != '', 'has_audio': inputs.get('audio') is not None, 'text_content': inputs.get('text', '')[:200], }, 'features': { 'modality': features_dict.get('modality', 'unknown'), 'feature_norms': {}, }, 'predictions': { 'n_voxels': len(brain_result['predictions']), 'pred_mean': float(np.mean(brain_result['predictions'])), 'pred_std': float(np.std(brain_result['predictions'])), 'pred_range': [float(np.min(brain_result['predictions'])), float(np.max(brain_result['predictions']))], 'uncertainty_mean': float(np.mean(brain_result['uncertainty'])), }, 'roi_summary_sent_to_llm': list(brain_result.get('roi_summary', {}).keys()), 'warnings': brain_result.get('warnings', []), 'qa_answer_length': len(qa_answer) if qa_answer else 0, } # Feature norms for key in ['image_multi_layer', 'text_multi_layer', 'audio_multi_layer']: if key in features_dict: entry['features']['feature_norms'][key] = float(features_dict[key].norm().item()) self.logs.append(entry) return entry def get_log_text(self): return json.dumps(self.logs[-5:], indent=2, default=str) # ============================================================ # Visualization helpers # ============================================================ def create_brain_activation_plot(brain_result, roi_annotations): """Create brain activation visualization.""" import plotly.graph_objects as go from plotly.subplots import make_subplots roi_summary = brain_result.get('roi_summary', {}) if not roi_summary: fig = go.Figure() fig.add_annotation(text="No ROI data available", x=0.5, y=0.5) return fig # Create multi-panel figure fig = make_subplots( rows=2, cols=2, subplot_titles=( "ROI Activation Strengths", "Functional Network Summary", "Activation Uncertainty", "Activation Distribution", ), specs=[ [{"type": "bar"}, {"type": "bar"}], [{"type": "bar"}, {"type": "histogram"}], ] ) # Panel 1: ROI activations sorted_rois = sorted(roi_summary.items(), key=lambda x: abs(x[1]['abs_mean']), reverse=True)[:15] roi_names = [r[0] for r in sorted_rois] roi_activations = [r[1]['mean_activation'] for r in sorted_rois] roi_colors = [] for r in sorted_rois: name = r[0] for net_name, roi_ids in FUNCTIONAL_NETWORKS.items(): roi_names_in_net = [ROI_NAMES[rid] for rid in roi_ids if rid in ROI_NAMES] if name in roi_names_in_net: color_map = { "early_visual": "#4CAF50", "body_selective": "#FF9800", "face_selective": "#E91E63", "place_selective": "#2196F3", "word_selective": "#9C27B0", } roi_colors.append(color_map.get(net_name, "#666")) break else: roi_colors.append("#666") fig.add_trace( go.Bar(x=roi_names, y=roi_activations, marker_color=roi_colors, name="Activation"), row=1, col=1 ) # Panel 2: Network summary net_names = [] net_activations = [] net_colors_list = [] color_map = { "early_visual": "#4CAF50", "body_selective": "#FF9800", "face_selective": "#E91E63", "place_selective": "#2196F3", "word_selective": "#9C27B0", } for net_name, roi_ids in FUNCTIONAL_NETWORKS.items(): roi_names_in_net = [ROI_NAMES[rid] for rid in roi_ids if rid in ROI_NAMES] activations = [roi_summary[rn]['abs_mean'] for rn in roi_names_in_net if rn in roi_summary] if activations: net_names.append(net_name.replace("_", " ").title()) net_activations.append(np.mean(activations)) net_colors_list.append(color_map.get(net_name, "#666")) fig.add_trace( go.Bar(x=net_names, y=net_activations, marker_color=net_colors_list, name="Network"), row=1, col=2 ) # Panel 3: Uncertainty roi_uncertainty = [r[1]['mean_uncertainty'] for r in sorted_rois] fig.add_trace( go.Bar(x=roi_names, y=roi_uncertainty, marker_color='rgba(255,0,0,0.5)', name="Uncertainty"), row=2, col=1 ) # Panel 4: Distribution predictions = brain_result['predictions'] fig.add_trace( go.Histogram(x=predictions[::10], nbinsx=50, name="Activations", marker_color='#4CAF50'), row=2, col=2 ) fig.update_layout( height=700, showlegend=False, title_text="Brain Activity Predictions", template="plotly_white", ) return fig def create_modality_contribution_plot(brain_result): """Create modality contribution visualization.""" import plotly.graph_objects as go contributions = brain_result.get('modality_contributions', {}) if len(contributions) <= 1: fig = go.Figure() fig.add_annotation(text="Single modality input - no comparison available", x=0.5, y=0.5) return fig fig = go.Figure() for mod_name, mod_pred in contributions.items(): # Show distribution of activations per modality fig.add_trace(go.Histogram( x=mod_pred[::10], name=mod_name.capitalize(), opacity=0.6, nbinsx=50, )) fig.update_layout( title="Modality Contributions to Brain Activity", xaxis_title="Predicted Activation", yaxis_title="Count", barmode='overlay', template="plotly_white", height=400, ) return fig # ============================================================ # Gradio Application # ============================================================ def build_gradio_app(): import gradio as gr # Global state manager = ModelManager() qa_system = GroundedQA() transparency_log = TransparencyLogger() current_result = {"value": None} def initialize(): try: manager.load() return "✅ Models loaded successfully!" except Exception as e: return f"❌ Error loading models: {e}" def process_input(image, text, audio): """Main inference pipeline.""" if not manager._loaded: manager.load() if image is None and (text is None or text.strip() == '') and audio is None: return "Please provide at least one input (image, text, or audio).", None, None, "" try: # Step 1: Extract features features = manager.extract_features(image=image, text=text, audio=audio) # Step 2: Predict brain activity result = manager.predict_brain_activity(features) current_result["value"] = result # Step 3: Create visualizations brain_plot = create_brain_activation_plot(result, manager.roi_annotations) modality_plot = create_modality_contribution_plot(result) # Step 4: Log for transparency log_entry = transparency_log.log_inference( {'image': image, 'text': text, 'audio': audio}, features, result ) # Summary text roi_summary = result.get('roi_summary', {}) sorted_rois = sorted(roi_summary.items(), key=lambda x: abs(x[1]['abs_mean']), reverse=True) summary = f"**Modality:** {result['modality']}\n" summary += f"**Voxels predicted:** {len(result['predictions'])}\n" summary += f"**Mean uncertainty:** {np.mean(result['uncertainty']):.4f}\n\n" summary += "**Top 5 Activated Regions:**\n" for roi_name, data in sorted_rois[:5]: summary += f"- {roi_name}: {data['mean_activation']:.4f} (±{data['mean_uncertainty']:.4f})\n" if result['warnings']: summary += "\n**Warnings:**\n" for w in result['warnings']: summary += f"- {w}\n" return summary, brain_plot, modality_plot, json.dumps(log_entry, indent=2, default=str) except Exception as e: import traceback return f"Error: {e}\n{traceback.format_exc()}", None, None, "" def ask_question(question, history): """Q&A with grounded interpretation.""" if current_result["value"] is None: history = history or [] history.append({"role": "user", "content": question}) history.append({"role": "assistant", "content": "Please run an inference first (provide an input in the Stimulus tab) before asking questions."}) return history, "" history = history or [] history.append({"role": "user", "content": question}) answer = qa_system.answer(question, current_result["value"]) history.append({"role": "assistant", "content": answer}) # Log Q&A transparency_log.log_inference( {'text': question}, {'modality': 'qa'}, current_result["value"], qa_answer=answer, ) return history, "" def get_transparency_log(): return transparency_log.get_log_text() # Build UI with gr.Blocks(title="Multimodal Brain Encoder") as demo: gr.Markdown(""" # 🧠 Multimodal Brain Encoder **A real brain encoding model trained on the Natural Scenes Dataset (NSD)** This system predicts brain activity (fMRI voxel responses) from multimodal inputs using: - **CLIP ViT-L/14** for feature extraction (multi-layer: layers 6, 12, 18, 24) - **Deep Brain Encoder** with ROI-specific attention heads (trained on NSD subj01) - **Ridge Regression** baseline (Algonauts 2023 recipe) - **Grounded LLM Q&A** that only interprets model predictions All predictions are from real model forward passes with learned weights. """) status = gr.Textbox(label="Status", value="Click 'Load Models' to initialize") load_btn = gr.Button("🔄 Load Models", variant="primary") load_btn.click(fn=initialize, outputs=status) with gr.Tabs(): # Tab 1: Input & Prediction with gr.Tab("🎯 Stimulus Input & Brain Prediction"): with gr.Row(): with gr.Column(scale=1): image_input = gr.Image(type="pil", label="Visual Stimulus (Image)") text_input = gr.Textbox( label="Text Input", placeholder="Enter a description or sentence...", lines=3, ) audio_input = gr.Audio(type="numpy", label="Audio Input") predict_btn = gr.Button("🧠 Predict Brain Activity", variant="primary", size="lg") with gr.Column(scale=2): summary_output = gr.Markdown(label="Prediction Summary") brain_plot = gr.Plot(label="Brain Activity Visualization") modality_plot = gr.Plot(label="Modality Contributions") predict_btn.click( fn=process_input, inputs=[image_input, text_input, audio_input], outputs=[summary_output, brain_plot, modality_plot, gr.Textbox(visible=False)], ) # Tab 2: Q&A with gr.Tab("💬 Grounded Q&A"): gr.Markdown(""" ### Ask questions about the predicted brain activity The LLM interpreter will answer based ONLY on: - Predicted activation maps and ROI summaries - Known functional labels from brain atlases - Modality attribution outputs - Uncertainty estimates It will NOT make independent neuroscience claims. """) chatbot = gr.Chatbot( type="messages", label="Brain Activity Q&A", height=400, ) with gr.Row(): question_input = gr.Textbox( label="Your Question", placeholder="e.g., Which brain regions are most activated? What does the face-selective network response mean?", scale=4, ) ask_btn = gr.Button("Ask", variant="primary", scale=1) ask_btn.click( fn=ask_question, inputs=[question_input, chatbot], outputs=[chatbot, question_input], ) question_input.submit( fn=ask_question, inputs=[question_input, chatbot], outputs=[chatbot, question_input], ) gr.Markdown(""" **Example questions:** - "What are the most activated brain regions for this input?" - "Is the face-selective network responding? What might that mean?" - "How confident is the model in these predictions?" - "How does the visual input differ from the text input in brain response?" - "What does high PPA activation suggest about this image?" """) # Tab 3: Transparency Log with gr.Tab("📋 Transparency Log"): gr.Markdown("### Full inference traceability log") gr.Markdown("Every inference is logged with inputs, features, predictions, and Q&A answers.") log_output = gr.Code(language="json", label="Recent Logs") refresh_log_btn = gr.Button("🔄 Refresh Log") refresh_log_btn.click(fn=get_transparency_log, outputs=log_output) # Tab 4: Model Info with gr.Tab("ℹ️ Model Information"): gr.Markdown(f""" ### Architecture Details | Component | Details | |-----------|---------| | Feature Extractor | CLIP ViT-L/14 (openai/clip-vit-large-patch14) | | Feature Layers | Layers 6, 12, 18, 24 (CLS tokens concatenated = 4096-dim) | | Brain Encoder | 4096 → 2048 → 2048 → 1024 → N_voxels | | Activations | GELU + BatchNorm + Dropout(0.3) | | ROI Heads | 5 functional network heads with learned attention | | Ridge Baseline | sklearn RidgeCV with 17 alphas (1e-2 to 1e6) | | Training Data | NSD subj01 (~8,859 train, ~300 val images) | | fMRI Resolution | 7T, ~15,724 voxels (NSD general cortical mask) | | Uncertainty | MC Dropout (10 forward passes) | ### Brain Regions (24 ROIs from NSD) | Network | Regions | Function | |---------|---------|----------| | Early Visual | V1v, V1d, V2v, V2d, V3v, V3d, hV4 | Basic visual processing | | Body Selective | EBA, FBA-1, FBA-2, mTL-bodies | Body/person perception | | Face Selective | OFA, FFA-1, FFA-2, mTL-faces, aTL-faces | Face recognition | | Place Selective | OPA, PPA, RSC | Scene/navigation | | Word Selective | OWFA, VWFA-1, VWFA-2, mfs-words, mTL-words | Reading/text | ### References - Natural Scenes Dataset: Allen et al. 2022, Nature Neuroscience - Algonauts 2023: Gifford et al. 2023 - CLIP: Radford et al. 2021 - Model repo: [{MODEL_REPO}](https://huggingface.co/{MODEL_REPO}) """) return demo if __name__ == "__main__": demo = build_gradio_app() demo.launch(server_name="0.0.0.0", server_port=7860)