| """
|
| Calcium-Bridge EEG Constraint Viewer (V2.1 - Fixed)
|
| Visualizes how constraint satisfaction unfolds across four temporal windows up to 550ms.
|
|
|
| Shows:
|
| 1. Original COCO image
|
| 2. EEG heatmaps for each of the 4 time windows
|
| 3. Calcium "attention" evolution (what the model focuses on at each stage)
|
| 4. Top predictions crystallizing across the 4 windows
|
|
|
| V2.1 Fixes:
|
| - Corrected 'figsize' argument placement during figure creation.
|
| - Corrected colorbar creation to use the figure object directly, resolving warnings.
|
| """
|
|
|
| import os
|
| import tkinter as tk
|
| from tkinter import filedialog, messagebox, ttk
|
| import numpy as np
|
| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
| from PIL import Image, ImageTk
|
| import matplotlib.pyplot as plt
|
| from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
|
| import json
|
| from pathlib import Path
|
| from collections import defaultdict
|
| import random
|
|
|
| try:
|
| from datasets import load_dataset
|
| except ImportError:
|
| print("Missing datasets library.")
|
| exit()
|
|
|
| DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| EEG_SAMPLE_RATE = 512
|
|
|
| TIME_WINDOWS = [
|
| (50, 150, "EarlyVisual"),
|
| (150, 250, "MidFeature"),
|
| (250, 350, "LateSemantic"),
|
| (350, 550, "CognitiveEvaluation")
|
| ]
|
|
|
| TARGET_CATEGORIES = {
|
| 'elephant': 22, 'giraffe': 25, 'bear': 23, 'zebra': 24,
|
| 'cow': 21, 'sheep': 20, 'horse': 19, 'dog': 18, 'cat': 17, 'bird': 16,
|
| 'airplane': 5, 'train': 7, 'boat': 9, 'bus': 6, 'truck': 8,
|
| 'motorcycle': 4, 'bicycle': 2, 'car': 3,
|
| 'traffic light': 10, 'fire hydrant': 11, 'stop sign': 13,
|
| 'parking meter': 14, 'bench': 15,
|
| }
|
|
|
| CATEGORY_NAMES = {v: k for k, v in TARGET_CATEGORIES.items()}
|
| TARGET_IDS = set(TARGET_CATEGORIES.values())
|
| ALL_COCO_IDS = list(range(1, 91))
|
| EXCLUDED_IDS = set(ALL_COCO_IDS) - TARGET_IDS
|
|
|
|
|
|
|
|
|
| class CalciumAttentionModule(nn.Module):
|
| def __init__(self, n_features, d_model=256):
|
| super().__init__()
|
| self.n_features = n_features
|
| self.d_model = d_model
|
| self.phase_proj = nn.Linear(n_features, d_model)
|
| self.ca_gate = nn.Sequential(
|
| nn.Linear(d_model, d_model // 2),
|
| nn.Sigmoid()
|
| )
|
| self.W = nn.Parameter(torch.randn(d_model, d_model) * 0.01)
|
| self.norm = nn.LayerNorm(d_model)
|
|
|
| def forward(self, x, prev_ca=None, prev_W=None):
|
| batch_size = x.size(0)
|
| phi = self.phase_proj(x)
|
|
|
| if prev_ca is None:
|
| ca = torch.zeros(batch_size, self.d_model, device=x.device)
|
| else:
|
| ca = prev_ca.clone()
|
|
|
| W = self.W if prev_W is None else prev_W
|
|
|
| coherence = torch.abs(torch.cos(phi[:, :, None] - phi[:, None, :]))
|
| ca_update = torch.mean(coherence, dim=2)
|
| ca = ca * 0.95 + ca_update * 0.05
|
|
|
| ca_gate = self.ca_gate(ca)
|
| coupled = torch.matmul(phi, W)
|
| ca_gate_full = torch.cat([ca_gate, ca_gate], dim=1)
|
| features = coupled * ca_gate_full
|
| features = self.norm(features + phi)
|
|
|
| return features, ca, W
|
|
|
|
|
| class TemporalConstraintEEGModel(nn.Module):
|
| def __init__(self, n_channels=64, num_classes=len(TARGET_CATEGORIES)):
|
| super().__init__()
|
| self.n_channels = n_channels
|
|
|
| self.window_encoders = nn.ModuleList([
|
| self._build_cnn_encoder() for _ in TIME_WINDOWS
|
| ])
|
|
|
| self.ca_modules = nn.ModuleList([
|
| CalciumAttentionModule(256, d_model=256) for _ in TIME_WINDOWS
|
| ])
|
|
|
| self.classifier = nn.Sequential(
|
| nn.Linear(256 * len(TIME_WINDOWS), 512),
|
| nn.BatchNorm1d(512),
|
| nn.GELU(),
|
| nn.Dropout(0.3),
|
| nn.Linear(512, num_classes)
|
| )
|
|
|
| def _build_cnn_encoder(self):
|
| return nn.Sequential(
|
| nn.Conv1d(self.n_channels, 128, kernel_size=15, padding=7),
|
| nn.BatchNorm1d(128),
|
| nn.ELU(),
|
| nn.MaxPool1d(2),
|
| nn.Conv1d(128, 256, kernel_size=7, padding=3),
|
| nn.BatchNorm1d(256),
|
| nn.ELU(),
|
| nn.AdaptiveAvgPool1d(1)
|
| )
|
|
|
| def forward(self, eeg_windows, return_intermediates=False):
|
| batch_size = eeg_windows[0].size(0)
|
|
|
| window_features, ca_history, W_history, window_logits_list = [], [], [], []
|
| ca_state, W_state = None, None
|
|
|
| for i, (encoder, ca_module, eeg_window) in enumerate(
|
| zip(self.window_encoders, self.ca_modules, eeg_windows)
|
| ):
|
| cnn_features = encoder(eeg_window).squeeze(-1)
|
| features, ca_state, W_state = ca_module(cnn_features, ca_state, W_state)
|
|
|
| window_features.append(features)
|
| if return_intermediates:
|
| ca_history.append(ca_state.detach())
|
| W_history.append(W_state.detach())
|
|
|
| padded_features = window_features + [
|
| torch.zeros_like(features) for _ in range(len(TIME_WINDOWS) - len(window_features))
|
| ]
|
| intermediate_logits = self.classifier(torch.cat(padded_features, dim=1))
|
| window_logits_list.append(intermediate_logits.detach())
|
|
|
| combined = torch.cat(window_features, dim=1)
|
| logits = self.classifier(combined)
|
|
|
| if return_intermediates:
|
| return logits, ca_history, W_history, window_logits_list
|
| return logits, ca_history
|
|
|
|
|
|
|
| class FilteredTestDataset:
|
| def __init__(self, annotations_path, max_samples=1000):
|
| print("Loading and filtering test dataset...")
|
| self.eeg_dataset = load_dataset("Alljoined/05_125", split='test', streaming=False).select(range(max_samples))
|
| with open(annotations_path, 'r') as f:
|
| coco_data = json.load(f)
|
|
|
| image_annotations = defaultdict(set)
|
| for ann in coco_data['annotations']:
|
| image_annotations[ann['image_id']].add(ann['category_id'])
|
|
|
| self.filtered_samples = []
|
| for idx, sample in enumerate(self.eeg_dataset):
|
| ann_ids = image_annotations.get(sample['coco_id'], set())
|
| if not any(cat_id in EXCLUDED_IDS for cat_id in ann_ids) and any(cat_id in TARGET_IDS for cat_id in ann_ids):
|
| self.filtered_samples.append({
|
| 'coco_id': sample['coco_id'],
|
| 'eeg_data': np.array(sample['EEG'], dtype=np.float32)
|
| })
|
| print(f"Loaded {len(self.filtered_samples)} filtered test samples.")
|
| if not self.filtered_samples: raise RuntimeError("No suitable test samples found.")
|
|
|
| def get_eeg_windows(self, sample_info):
|
| eeg_data = sample_info['eeg_data']
|
| eeg_windows = []
|
| for start_ms, end_ms, _ in TIME_WINDOWS:
|
| start_idx, end_idx = int(start_ms / 1000 * EEG_SAMPLE_RATE), int(end_ms / 1000 * EEG_SAMPLE_RATE)
|
| n_timepoints = end_idx - start_idx
|
| window = eeg_data[:, start_idx:end_idx] if eeg_data.shape[1] >= end_idx else eeg_data[:, start_idx:]
|
|
|
| if window.shape[1] != n_timepoints:
|
| pad_width = n_timepoints - window.shape[1]
|
| window = np.pad(window, ((0,0), (0, pad_width)), 'edge') if pad_width > 0 else window[:, :n_timepoints]
|
|
|
| window = (window - window.mean(axis=1, keepdims=True)) / (window.std(axis=1, keepdims=True) + 1e-8)
|
| eeg_windows.append(window)
|
| return eeg_windows
|
|
|
| def get_random_sample_info(self):
|
| return random.choice(self.filtered_samples)
|
|
|
|
|
| class CalciumBridgeViewer(tk.Tk):
|
| def __init__(self):
|
| super().__init__()
|
| self.title("Calcium-Bridge EEG Constraint Viewer V2 (Extended Window)")
|
| self.geometry("2000x1000")
|
| self.model, self.test_data = None, None
|
| self.setup_gui()
|
|
|
| def setup_gui(self):
|
| control_frame = ttk.Frame(self); control_frame.pack(pady=10, padx=10, fill=tk.X)
|
| ttk.Label(control_frame, text="COCO Path:").pack(side=tk.LEFT, padx=5)
|
| self.coco_var = tk.StringVar(); ttk.Entry(control_frame, textvariable=self.coco_var, width=20).pack(side=tk.LEFT, padx=2)
|
| ttk.Button(control_frame, text="Browse", command=self.browse_coco).pack(side=tk.LEFT, padx=5)
|
| ttk.Label(control_frame, text="Annotations:").pack(side=tk.LEFT, padx=5)
|
| self.ann_var = tk.StringVar(); ttk.Entry(control_frame, textvariable=self.ann_var, width=20).pack(side=tk.LEFT, padx=2)
|
| ttk.Button(control_frame, text="Browse", command=self.browse_ann).pack(side=tk.LEFT, padx=5)
|
| ttk.Button(control_frame, text="Load V2 Model", command=self.load_model).pack(side=tk.LEFT, padx=20)
|
| self.test_btn = ttk.Button(control_frame, text="Test Random Sample", command=self.test_sample, state=tk.DISABLED); self.test_btn.pack(side=tk.LEFT, padx=5)
|
| self.status_label = tk.Label(control_frame, text="Model: Not loaded", fg="gray"); self.status_label.pack(side=tk.LEFT, padx=20)
|
|
|
| main_paned = ttk.PanedWindow(self, orient=tk.HORIZONTAL); main_paned.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
|
| image_frame = ttk.Frame(main_paned, width=400); main_paned.add(image_frame, weight=0)
|
| ttk.Label(image_frame, text="COCO Image", font=("Arial", 12, "bold")).pack(pady=5)
|
| self.image_canvas = tk.Canvas(image_frame, width=400, height=400, bg='lightgray'); self.image_canvas.pack()
|
| self.coco_id_label = ttk.Label(image_frame, text="COCO ID: N/A"); self.coco_id_label.pack(pady=5)
|
|
|
| self.notebook = ttk.Notebook(main_paned); main_paned.add(self.notebook, weight=1)
|
| self.create_tabs()
|
|
|
| def create_tabs(self):
|
| self.constraint_fig, self.constraint_canvas = self.create_tab("Constraint Satisfaction", "How predictions crystallize as constraints are satisfied")
|
| self.calcium_fig, self.calcium_canvas = self.create_tab("Calcium Attention", "Calcium state evolution: What the model 'focuses on' at each stage")
|
| self.eeg_fig, self.eeg_canvas = self.create_tab("EEG Heatmaps", "Raw EEG signals for each time window")
|
|
|
| def create_tab(self, title, description):
|
| tab = ttk.Frame(self.notebook); self.notebook.add(tab, text=title)
|
| ttk.Label(tab, text=description, font=("Arial", 11)).pack(pady=5)
|
| fig = plt.Figure()
|
| canvas = FigureCanvasTkAgg(fig, tab); canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True)
|
| return fig, canvas
|
|
|
| def browse_coco(self):
|
| path = filedialog.askdirectory(); self.coco_var.set(path); self.coco_path = path
|
|
|
| def browse_ann(self):
|
| path = filedialog.askopenfilename(filetypes=[("JSON", "*.json")]); self.ann_var.set(path); self.annotations_path = path
|
|
|
| def load_model(self):
|
| model_path = filedialog.askopenfilename(filetypes=[("PyTorch Model", "*.pth")], title="Select calcium_bridge_eeg_model_v2.pth")
|
| if not model_path or not self.annotations_path: return
|
| try:
|
| checkpoint = torch.load(model_path, map_location=DEVICE)
|
| self.model = TemporalConstraintEEGModel().to(DEVICE)
|
| self.model.load_state_dict(checkpoint['model_state_dict'])
|
| self.model.eval()
|
| self.test_data = FilteredTestDataset(self.annotations_path)
|
| self.status_label.config(text="Model: V2 Loaded ✓", fg="green")
|
| self.test_btn.config(state=tk.NORMAL)
|
| except Exception as e: messagebox.showerror("Error", f"Failed to load model:\n{e}"); print(traceback.format_exc())
|
|
|
| def _fetch_image(self, coco_id):
|
| formatted_id = f"{coco_id:012d}.jpg"
|
| for s in ["train2017", "val2017", "test2017"]:
|
| path = os.path.join(self.coco_path, s, formatted_id)
|
| if os.path.exists(path): return Image.open(path).convert("RGB")
|
| return None
|
|
|
| def test_sample(self):
|
| if not self.model: return
|
| try:
|
| sample_info = self.test_data.get_random_sample_info()
|
| image = self._fetch_image(sample_info['coco_id'])
|
| if image: self.display_image(image, sample_info['coco_id'])
|
|
|
| eeg_windows_np = self.test_data.get_eeg_windows(sample_info)
|
| eeg_windows = [torch.from_numpy(w).unsqueeze(0).to(DEVICE) for w in eeg_windows_np]
|
|
|
| with torch.no_grad():
|
| logits, ca_history, _, window_logits = self.model(eeg_windows, return_intermediates=True)
|
|
|
| self.visualize_constraint_satisfaction(window_logits, logits)
|
| self.visualize_calcium_evolution(ca_history)
|
| self.visualize_eeg_heatmaps(eeg_windows_np)
|
| except Exception as e: messagebox.showerror("Error", f"Failed to process sample:\n{e}"); print(traceback.format_exc())
|
|
|
| def display_image(self, image, coco_id):
|
| ratio = min(400/image.width, 400/image.height)
|
| resized = image.resize((int(image.width * ratio), int(image.height * ratio)), Image.LANCZOS)
|
| self.pil_image_tk = ImageTk.PhotoImage(resized)
|
| self.image_canvas.create_image(200, 200, image=self.pil_image_tk)
|
| self.coco_id_label.config(text=f"COCO ID: {coco_id}")
|
|
|
| def visualize_constraint_satisfaction(self, window_logits, final_logits):
|
| self.constraint_fig.clear()
|
| cat_list = list(TARGET_CATEGORIES.keys())
|
| n_windows = len(window_logits)
|
| final_probs = torch.sigmoid(final_logits).squeeze(0).cpu().numpy()
|
| top_indices = np.argsort(final_probs)[::-1][:10]
|
| axes = self.constraint_fig.subplots(1, n_windows + 1)
|
|
|
| for i, (ax, wl) in enumerate(zip(axes[:-1], window_logits)):
|
| probs = torch.sigmoid(wl).squeeze(0).cpu().numpy()[top_indices]
|
| ax.barh([cat_list[idx] for idx in top_indices], probs, color='steelblue')
|
| ax.set_title(f"{TIME_WINDOWS[i][2]}\n({TIME_WINDOWS[i][0]}-{TIME_WINDOWS[i][1]}ms)", fontsize=10)
|
| ax.set_xlim(0, 1); ax.invert_yaxis(); ax.tick_params(axis='y', labelsize=8)
|
|
|
| axes[-1].barh([cat_list[idx] for idx in top_indices], final_probs[top_indices], color='darkgreen')
|
| axes[-1].set_title("Final\n(Combined)", fontsize=10); axes[-1].set_xlim(0, 1); axes[-1].invert_yaxis(); axes[-1].tick_params(axis='y', labelsize=8)
|
| self.constraint_fig.suptitle("Constraint Satisfaction: Predictions Crystallizing Over Time", fontsize=14, fontweight='bold')
|
| self.constraint_fig.tight_layout(); self.constraint_canvas.draw()
|
|
|
| def visualize_calcium_evolution(self, ca_history):
|
| self.calcium_fig.clear()
|
| n_windows = len(ca_history)
|
| axes = self.calcium_fig.subplots(2, n_windows)
|
|
|
| for i, ca_state in enumerate(ca_history):
|
| ca_np = ca_state.squeeze(0).cpu().numpy()
|
| top_20_idx = np.argsort(ca_np)[::-1][:20]
|
| axes[0, i].plot(ca_np, 'r'); axes[0, i].fill_between(range(len(ca_np)), ca_np, color='r', alpha=0.3)
|
| axes[0, i].set_title(f"{TIME_WINDOWS[i][2]}\n({TIME_WINDOWS[i][0]}-{TIME_WINDOWS[i][1]}ms)", fontsize=10)
|
| axes[1, i].barh([f"F{idx}" for idx in top_20_idx], ca_np[top_20_idx], color='darkred')
|
| axes[1, i].invert_yaxis(); axes[1, i].tick_params(axis='y', labelsize=7)
|
| self.calcium_fig.suptitle("Calcium Attention: What the Model Focuses On", fontsize=14, fontweight='bold')
|
| self.calcium_fig.tight_layout(); self.calcium_canvas.draw()
|
|
|
| def visualize_eeg_heatmaps(self, eeg_windows_np):
|
| self.eeg_fig.clear()
|
| n_windows = len(eeg_windows_np)
|
| axes = self.eeg_fig.subplots(1, n_windows)
|
|
|
| for i, (ax, eeg_data) in enumerate(zip(axes, eeg_windows_np)):
|
| im = ax.imshow(eeg_data, aspect='auto', cmap='RdBu_r', vmin=-3, vmax=3)
|
| ax.set_title(f"{TIME_WINDOWS[i][2]}\n({TIME_WINDOWS[i][0]}-{TIME_WINDOWS[i][1]}ms)", fontsize=10)
|
| if i == 0: ax.set_ylabel("Channel")
|
| self.eeg_fig.colorbar(im, ax=ax)
|
| self.eeg_fig.suptitle("Raw EEG Signals by Time Window", fontsize=14, fontweight='bold')
|
| self.eeg_fig.tight_layout(); self.eeg_canvas.draw()
|
|
|
| if __name__ == "__main__":
|
| app = CalciumBridgeViewer()
|
| app.mainloop() |
