Add attention visualization to Gradio demo with multi-modal display

app.py

@@ -4,6 +4,7 @@ BrainAnytime Hugging Face Space Demo
 
 Interactive demo for brain image analysis with multi-modal support.
 Supports 4 tasks and 5 modality combinations.
+Now with attention visualization!
 """
 
 import os
@@ -12,6 +13,7 @@ import json
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple
 import warnings
+import tempfile
 warnings.filterwarnings("ignore")
 
 # 添加当前目录到路径（用于导入 inference_engine）
@@ -32,10 +34,23 @@ except Exception as e:
     print(f"Warning: Inference engine not available: {e}")
     INFERENCE_AVAILABLE = False
 
+# 导入注意力可视化
+try:
+    from brain_attention_viz import (
+        visualize_sample_attention_precise,
+        build_voxel_attention_map,
+        extract_last_layer_attention,
+    )
+    ATTENTION_AVAILABLE = True
+except Exception as e:
+    print(f"Warning: Attention visualization not available: {e}")
+    ATTENTION_AVAILABLE = False
+
 # 延迟导入 Gradio（避免启动时的循环导入问题）
 try:
     import gradio as gr
     import numpy as np
+    import nibabel as nib
     from PIL import Image
     GRADIO_AVAILABLE = True
 except Exception as e:
@@ -105,6 +120,9 @@ TASK_CONFIG = {
 SAMPLES_DIR = BASE_DIR / "demo_samples"
 SAMPLES_INDEX = SAMPLES_DIR / "samples.json"
 
+# Atlas 路径
+ATLAS_PATH = BASE_DIR / "BrainAnytime" / "altas" / "AAL116_standard.nii.gz"
+
 
 # =============================================================================
 # 全局状态
@@ -112,6 +130,7 @@ SAMPLES_INDEX = SAMPLES_DIR / "samples.json"
 
 _inference_engine: Optional[BrainAnytimeInference] = None
 _samples_cache: Optional[Dict] = None
+_atlas_cache: Optional[np.ndarray] = None
 
 
 def get_inference_engine() -> Optional[BrainAnytimeInference]:
@@ -172,78 +191,104 @@ def get_sample_info(task: str, combo: str) -> Optional[Dict]:
     return None
 
 
-def get_preview_images(sample_info: Dict) -> List[Tuple[str, str]]:
-    """获取样本的预览图路径和标题"""
+def load_atlas_mask() -> Optional[np.ndarray]:
+    """加载 AAL116 atlas 作为 brain mask"""
+    global _atlas_cache
+    
+    if _atlas_cache is not None:
+        return _atlas_cache
+    
+    if not ATLAS_PATH.exists():
+        print(f"Warning: Atlas not found at {ATLAS_PATH}")
+        return None
+    
+    try:
+        atlas_nii = nib.load(str(ATLAS_PATH))
+        atlas_data = atlas_nii.get_fdata().astype(np.int32)
+        _atlas_cache = (atlas_data > 0).astype(np.float32)
+        return _atlas_cache
+    except Exception as e:
+        print(f"Error loading atlas: {e}")
+        return None
+
+
+def get_sample_preview_paths(sample_info: Dict) -> Dict[str, str]:
+    """获取样本所有模态的预览图路径"""
     if not sample_info:
-        return []
+        return {}
     
-    images = []
     sample_dir = SAMPLES_DIR / sample_info["task"] / sample_info["modality_combination"] / sample_info["sample_name"]
     
-    for mod in MODALITY_ORDER:
+    preview_paths = {}
+    for mod in sample_info.get("modalities", []):
         preview_file = sample_info.get("files", {}).get("previews", {}).get(mod)
         if preview_file:
             img_path = sample_dir / preview_file
             if img_path.exists():
-                images.append((str(img_path), f"{mod} - Axial Slice"))
+                preview_paths[mod] = str(img_path)
     
-    return images
+    return preview_paths
 
 
 # =============================================================================
-# Gradio 回调函数
+# 推理 + 注意力可视化
 # =============================================================================
 
-def update_sample_gallery(task: str, combo: str):
-    """更新样本画廊"""
-    sample_info = get_sample_info(task, combo)
-    
-    if not sample_info:
-        return None, "Sample not found"
-    
-    # 获取预览图
-    images = get_preview_images(sample_info)
-    
-    if not images:
-        return None, "No preview images available"
-    
-    # 返回第一张图作为代表
-    return images[0][0], f"Sample: {sample_info['subject_id']} | Label: {sample_info.get('diag_group', 'N/A')}"
-
-
 @spaces.GPU
-def run_inference(task: str, combo: str):
-    """执行推理 (GPU accelerated)"""
+def run_inference_with_attention(task: str, combo: str):
+    """
+    执行推理并生成注意力可视化
+    
+    Returns:
+        (prediction_markdown, modality_images_dict)
+        - prediction_markdown: 预测结果的 markdown 文本
+        - modality_images_dict: {模态名: 注意力可视化图像路径}
+    """
     engine = get_inference_engine()
     
     if not engine:
-        return {
-            "error": "Inference engine not available. Please check model checkpoints."
-        }
+        return "❌ **Error**: Inference engine not available", {}
     
     # 获取样本信息
     sample_info = get_sample_info(task, combo)
     if not sample_info:
-        return {"error": f"No sample found for {task}/{combo}"}
+        return f"❌ **Error**: No sample found for {task}/{combo}", {}
     
-    # 构建��本目录路径
     sample_dir = SAMPLES_DIR / task / combo / sample_info["sample_name"]
     
-    # 执行推理
     try:
+        # 执行推理
         result = engine.predict_from_sample(str(sample_dir), task, combo)
         
         if result is None:
-            return {"error": "Inference failed"}
+            return "❌ **Error**: Inference failed", {}
         
-        return result
+        # 格式化预测结果
+        prediction_md = format_prediction_with_confidence(result)
+        
+        # 生成注意力可视化（如果可用）
+        attention_images = {}
+        if ATTENTION_AVAILABLE and engine:
+            try:
+                attention_images = generate_attention_visualization(
+                    engine, sample_dir, sample_info, task, combo
+                )
+            except Exception as e:
+                print(f"Attention visualization failed: {e}")
+                # 如果注意力可视化失败，返回原始预览图
+                attention_images = get_sample_preview_paths(sample_info)
+        else:
+            # 注意力不可用，返回原始预览图
+            attention_images = get_sample_preview_paths(sample_info)
+        
+        return prediction_md, attention_images
         
     except Exception as e:
-        return {"error": f"Inference error: {str(e)}"}
+        return f"❌ **Error**: {str(e)}", {}
 
 
-def format_prediction(result: Dict) -> str:
-    """格式化预测结果为 Markdown"""
+def format_prediction_with_confidence(result: Dict) -> str:
+    """格式化预测结果为 Markdown（带置信度颜色）"""
     if "error" in result:
         return f"❌ **Error**: {result['error']}"
     
@@ -251,7 +296,7 @@ def format_prediction(result: Dict) -> str:
     task_type = result.get("task_type", "unknown")
     
     lines = [
-        f"## Prediction Result: {TASK_CONFIG.get(task, {}).get('name', task)}",
+        f"## 🎯 Prediction Result: {TASK_CONFIG.get(task, {}).get('name', task)}",
         "",
     ]
     
@@ -261,6 +306,14 @@ def format_prediction(result: Dict) -> str:
         conf = result.get("confidence", 0)
         classes = result.get("classes", [])
         
+        # 置信度颜色
+        if conf >= 0.7:
+            conf_color = "🟢"
+        elif conf >= 0.5:
+            conf_color = "🟡"
+        else:
+            conf_color = "🔴"
+        
         lines.extend([
             f"**Predicted Class**: {pred}",
             "",
@@ -268,7 +321,7 @@ def format_prediction(result: Dict) -> str:
             f"- {classes[0]}: {1-prob:.3f}",
             f"- {classes[1]}: {prob:.3f}",
             "",
-            f"**Confidence**: {conf:.1%}",
+            f"**Confidence**: {conf_color} {conf:.1%}",
         ])
     else:
         pred = result.get("prediction", 0)
@@ -289,6 +342,153 @@ def format_prediction(result: Dict) -> str:
     return "\n".join(lines)
 
 
+def generate_attention_visualization(
+    engine: BrainAnytimeInference,
+    sample_dir: Path,
+    sample_info: Dict,
+    task: str,
+    combo: str
+) -> Dict[str, str]:
+    """
+    生成注意力可视化图像
+    
+    Returns:
+        {模态名: 可视化图像路径}
+    """
+    from scipy.ndimage import gaussian_filter
+    import matplotlib
+    matplotlib.use('Agg')
+    import matplotlib.pyplot as plt
+    from matplotlib import cm
+    
+    # 加载模型
+    model = engine.load_model(task)
+    
+    # 加载脑影像
+    brain_volumes = {}
+    for mod in sample_info.get("modalities", []):
+        nii_filename = sample_info.get("files", {}).get("nifti", {}).get(mod)
+        if nii_filename:
+            nii_path = sample_dir / nii_filename
+            if nii_path.exists():
+                nii = nib.load(str(nii_path))
+                data = nii.get_fdata().astype(np.float32)
+                # Min-Max 归一化
+                data_min, data_max = data.min(), data.max()
+                if data_max > data_min:
+                    data = (data - data_min) / (data_max - data_min)
+                brain_volumes[mod] = data
+    
+    # 准备模型输入
+    images_list = []
+    for mod in MODALITY_ORDER:
+        if mod in brain_volumes:
+            images_list.append(brain_volumes[mod])
+        else:
+            images_list.append(np.zeros((128, 128, 128), dtype=np.float32))
+    
+    images = np.stack(images_list, axis=0)
+    images = torch.from_numpy(images).unsqueeze(0).to(engine.device)
+    
+    # observed mask
+    observed = torch.zeros(1, 4)
+    for i, mod in enumerate(MODALITY_ORDER):
+        if mod in sample_info.get("modalities", []):
+            observed[0, i] = 1.0
+    
+    # 提取注意力
+    attention = extract_last_layer_attention(model.encoder, images, observed)
+    
+    # 处理注意力
+    attn_batch = attention[0].mean(dim=0)  # [N, N]
+    num_global = model.encoder.num_global_tokens
+    cls_to_all = attn_batch[0, num_global:]  # 跳过 CLS
+    
+    total_patches = cls_to_all.shape[0]
+    N_p_actual = total_patches // 4
+    
+    # 加载 atlas mask
+    brain_mask = load_atlas_mask()
+    
+    # 为每个模态生成可视化
+    attention_images = {}
+    mid_slice = 70
+    
+    for mod in sample_info.get("modalities", []):
+        mod_idx = MODALITY_ORDER.index(mod)
+        start_idx = mod_idx * N_p_actual
+        end_idx = (mod_idx + 1) * N_p_actual
+        mod_attn = cls_to_all[start_idx:end_idx].cpu().numpy()
+        
+        # 构建注意力图
+        atlas_data = np.ones((128, 128, 128), dtype=np.int32)
+        attention_map = build_voxel_attention_map(
+            mod_attn, atlas_data, (128, 128, 128), (16, 16, 16)
+        )
+        attention_map = gaussian_filter(attention_map, sigma=2.0)
+        
+        brain_vol = brain_volumes[mod]
+        
+        # 获取 mask 切片
+        if brain_mask is not None:
+            mask_slice = brain_mask[:, :, mid_slice]
+        else:
+            mask_slice = (brain_vol[:, :, mid_slice] > 0.01).astype(np.float32)
+        
+        # 计算相对注意力
+        masked_attn = attention_map[brain_mask > 0] if brain_mask is not None else attention_map.flatten()
+        if len(masked_attn) > 0:
+            mean_attn = masked_attn.mean()
+        else:
+            mean_attn = attention_map.mean()
+        
+        attn_slice = attention_map[:, :, mid_slice]
+        brain_slice = brain_vol[:, :, mid_slice]
+        attn_relative = attn_slice - mean_attn
+        max_abs = np.abs(attn_relative).max()
+        if max_abs > 0:
+            attn_norm = (attn_relative / max_abs) * 0.5 + 0.5
+        else:
+            attn_norm = np.ones_like(attn_relative) * 0.5
+        
+        mask_bool = mask_slice > 0
+        attn_norm = np.where(mask_bool, attn_norm, 0.5)
+        
+        # 生成图像（黑背景线上版本）
+        fig, ax = plt.subplots(1, 1, figsize=(6, 6), facecolor='black')
+        ax.set_facecolor('black')
+        
+        # 应用 RdBu_r colormap
+        cmap_obj = cm.get_cmap('RdBu_r')
+        attn_colored = cmap_obj(attn_norm)
+        attn_colored = (attn_colored[:, :, :3] * 255).astype(np.uint8)
+        
+        # 黑背景
+        black_bg = np.zeros_like(brain_slice)
+        display_img = np.stack([black_bg, black_bg, black_bg], axis=-1)
+        
+        brain_norm = (brain_slice * 255).astype(np.uint8)
+        alpha = 0.6
+        for c in range(3):
+            blended = (1 - alpha * 0.5) * brain_norm + alpha * 1.5 * attn_colored[:, :, c]
+            blended = np.clip(blended, 0, 255)
+            display_img[:, :, c] = np.where(mask_bool, blended, display_img[:, :, c])
+        
+        ax.imshow(display_img, origin='lower')
+        ax.set_title(f'{mod} Attention', fontsize=14, color='white')
+        ax.axis('off')
+        
+        # 保存到临时文件
+        temp_file = tempfile.NamedTemporaryFile(suffix=f'_{mod}_attention.png', delete=False)
+        plt.savefig(temp_file.name, dpi=100, bbox_inches='tight', 
+                   facecolor='black', edgecolor='none')
+        plt.close()
+        
+        attention_images[mod] = temp_file.name
+    
+    return attention_images
+
+
 # =============================================================================
 # 创建 Gradio 界面
 # =============================================================================
@@ -302,8 +502,9 @@ def create_demo():
     with gr.Blocks(
         title="BrainAnytime Demo",
         css="""
-        .preview-image { max-height: 300px; }
+        .preview-image { max-height: 250px; }
         .result-box { font-size: 16px; }
+        .modality-grid { display: grid; grid-template-columns: repeat(auto-fit, minmax(200px, 1fr)); gap: 10px; }
         """
     ) as demo:
         
@@ -313,7 +514,7 @@ def create_demo():
         **BrainAnytime: Anatomy-Aware Cross-Modal Pretraining for Brain Image Analysis**
         
         This demo showcases the BrainAnytime model for multi-modal 3D brain image analysis.
-        The model supports arbitrary missing modality combinations at inference time.
+        Select a task and modality combination, then run inference to see predictions and attention visualization.
         
         ---
         """)
@@ -362,13 +563,38 @@ def create_demo():
                 
                 # 右侧：结果展示
                 with gr.Column(scale=2):
-                    gr.Markdown("### Sample Preview")
+                    gr.Markdown("### Sample Preview & Attention Visualization")
                     
-                    sample_preview = gr.Image(
-                        label="Brain MRI Preview (Axial Slice)",
-                        type="filepath",
-                        height=300,
-                    )
+                    # 动态创建图像组件 - 根据选择的模态组合显示对应数量的图像
+                    with gr.Row():
+                        # T1 图像（总是显示，因为所有组合都有 T1）
+                        t1_preview = gr.Image(
+                            label="T1",
+                            type="filepath",
+                            height=250,
+                            visible=True
+                        )
+                        # T2 图像
+                        t2_preview = gr.Image(
+                            label="T2",
+                            type="filepath",
+                            height=250,
+                            visible=False
+                        )
+                        # Flair 图像
+                        flair_preview = gr.Image(
+                            label="Flair",
+                            type="filepath",
+                            height=250,
+                            visible=False
+                        )
+                        # PET 图像
+                        pet_preview = gr.Image(
+                            label="PET",
+                            type="filepath",
+                            height=250,
+                            visible=False
+                        )
                     
                     sample_info = gr.Textbox(
                         label="Sample Info",
@@ -391,6 +617,48 @@ def create_demo():
                 mods = combo_info.get("modalities", [])
                 return f"**{combo}**: {', '.join(mods)}"
             
+            def update_preview_visibility(combo):
+                """根据选择的模态组合，更新各图像组件的可见性"""
+                combo_info = MODALITY_COMBOS.get(combo, {})
+                mods = combo_info.get("modalities", [])
+                
+                # 返回各组件的可见性
+                return {
+                    t1_preview: gr.update(visible="T1" in mods),
+                    t2_preview: gr.update(visible="T2" in mods),
+                    flair_preview: gr.update(visible="Flair" in mods),
+                    pet_preview: gr.update(visible="PET" in mods),
+                }
+            
+            def on_run_inference(task, combo):
+                """运行推理并更新所有输出"""
+                pred_md, attention_images = run_inference_with_attention(task, combo)
+                
+                # 获取样本信息
+                sample_info_data = get_sample_info(task, combo)
+                info_text = f"Sample: {sample_info_data.get('subject_id', 'N/A')} | Label: {sample_info_data.get('diag_group', 'N/A')}"
+                
+                # 构建输出字典
+                outputs = {
+                    result_display: pred_md,
+                    sample_info: info_text,
+                }
+                
+                # 更新各模态图像
+                combo_info = MODALITY_COMBOS.get(combo, {})
+                mods = combo_info.get("modalities", [])
+                
+                if "T1" in mods:
+                    outputs[t1_preview] = attention_images.get("T1", None)
+                if "T2" in mods:
+                    outputs[t2_preview] = attention_images.get("T2", None)
+                if "Flair" in mods:
+                    outputs[flair_preview] = attention_images.get("Flair", None)
+                if "PET" in mods:
+                    outputs[pet_preview] = attention_images.get("PET", None)
+                
+                return outputs
+            
             task_selector.change(
                 update_task_desc,
                 inputs=task_selector,
@@ -403,38 +671,24 @@ def create_demo():
                 outputs=modality_desc
             )
             
-            # 更新样本预览
-            def on_config_change(task, combo):
-                return update_sample_gallery(task, combo)
-            
-            task_selector.change(
-                on_config_change,
-                inputs=[task_selector, modality_selector],
-                outputs=[sample_preview, sample_info]
-            )
-            
             modality_selector.change(
-                on_config_change,
-                inputs=[task_selector, modality_selector],
-                outputs=[sample_preview, sample_info]
+                update_preview_visibility,
+                inputs=modality_selector,
+                outputs=[t1_preview, t2_preview, flair_preview, pet_preview]
             )
             
-            # 运行推理
-            def on_run_inference(task, combo):
-                result = run_inference(task, combo)
-                return format_prediction(result)
-            
+            # 运行推理 - 更新所有输出
             run_btn.click(
                 on_run_inference,
                 inputs=[task_selector, modality_selector],
-                outputs=result_display
-            )
-            
-            # 初始化
-            demo.load(
-                on_config_change,
-                inputs=[task_selector, modality_selector],
-                outputs=[sample_preview, sample_info]
+                outputs=[
+                    result_display,
+                    sample_info,
+                    t1_preview,
+                    t2_preview,
+                    flair_preview,
+                    pet_preview
+                ]
             )
         
         # ==================== Tab 2: 项目信息 ====================
@@ -450,9 +704,9 @@ def create_demo():
             ### Key Features
             
             - **Multi-modal Support**: T1, T2, Flair, PET
-            - **Missing Modality Robustness**: Handles arbitrary missing combinations
+            - **Missing Modality Robustness**: Handles arbitrary missing modality combinations
             - **Anatomy-Aware**: Uses AAL116 brain atlas for adaptive masking
-            - **Pretrained Model**: Self-supervised pretraining on large-scale datasets
+            - **Attention Visualization**: See which brain regions the model focuses on
             
             ### Supported Tasks
             
@@ -482,28 +736,6 @@ def create_demo():
             ```
             """)
         
-        # ==================== Tab 3: 样本库 ====================
-        with gr.Tab("🗂️ 样本库 (Sample Gallery)"):
-            gr.Markdown("""
-            Browse all 20 pre-selected samples (5 modality combinations × 4 tasks).
-            Each sample includes brain MRI preview images.
-            """)
-            
-            # 为每个任务创建样本展示
-            for task_key, task_info in TASK_CONFIG.items():
-                with gr.Accordion(f"{task_info['name']}", open=False):
-                    for combo_key, combo_info in MODALITY_COMBOS.items():
-                        sample = get_sample_info(task_key, combo_key)
-                        if sample:
-                            with gr.Row():
-                                images = get_preview_images(sample)
-                                for img_path, title in images[:2]:  # 最多显示2张
-                                    gr.Image(
-                                        value=img_path,
-                                        label=f"{combo_key} - {title}",
-                                        height=200,
-                                    )
-        
         gr.Markdown("""
         ---
         

brain_attention_viz.py

@@ -0,0 +1,556 @@
+#!/usr/bin/env python3
+"""
+BrainAnytime 精确注意力可视化
+
+利用 AAL116 atlas 的 patch-region mapping 实现精确的 3D 空间映射
+避免简单的插值造成的 artifact
+"""
+
+import sys
+import os
+from pathlib import Path
+
+import numpy as np
+import torch
+import nibabel as nib
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+from matplotlib.colors import LinearSegmentedColormap
+
+# 添加路径
+BASE_DIR = Path(__file__).parent
+sys.path.insert(0, str(BASE_DIR))
+sys.path.insert(0, str(BASE_DIR / "BrainAnytime"))
+
+from inference_engine import BrainAnytimeInference, MODALITY_ORDER, TASKS
+from models.multimae3d import MultiMAE3D
+from models.multimae3d_utils import patchify
+
+
+def extract_last_layer_attention(
+    model: MultiMAE3D,
+    images: torch.Tensor,
+    observed: torch.Tensor,
+) -> torch.Tensor:
+    """
+    提取最后一层编码器的 CLS-to-patch 注意力权重
+    
+    Args:
+        model: MultiMAE3D 模型
+        images: [B, 4, D, H, W] 输入图像
+        observed: [B, 4] 模态掩码
+    
+    Returns:
+        attention: [B, num_heads, N, N] 注意力权重
+    """
+    model.eval()
+    
+    with torch.no_grad():
+        B = images.shape[0]
+        device = images.device
+        
+        # 准备输入 tokens
+        all_patches = []
+        for i, modality in enumerate(model.MODALITY_NAMES):
+            x = images[:, i:i+1, ...]
+            patches = patchify(x, model.patch_size)
+            patches_emb = model.input_adapters[modality](patches)
+            all_patches.append(patches_emb)
+        
+        input_tokens = torch.cat(all_patches, dim=1)
+        
+        # 添加全局 token (CLS)
+        if model.num_global_tokens > 0:
+            cls = model.global_tokens.unsqueeze(0).expand(B, -1, -1)
+            input_tokens = torch.cat([cls, input_tokens], dim=1)
+        
+        # 构建 attention mask
+        N_p = model.num_patches
+        attn_mask = torch.zeros(B, 1, 1, input_tokens.shape[1], device=device)
+        mod_offset = model.num_global_tokens
+        
+        for i, mod in enumerate(model.MODALITY_NAMES):
+            mask_val = (1.0 - observed[:, i:i+1]) * -1e9
+            attn_mask[:, :, :, mod_offset:mod_offset + N_p] = mask_val.unsqueeze(-1)
+            mod_offset += N_p
+        
+        if (attn_mask == 0).all():
+            attn_mask = None
+        
+        # 手动执行前向传播直到最后一层
+        x = input_tokens
+        
+        if model.pos_embed is not None:
+            num_global = model.num_global_tokens
+            cls_tokens = x[:, :num_global, :]
+            patch_tokens = x[:, num_global:, :]
+            
+            pos_emb = model.pos_embed.expand(B, -1, -1)
+            for i in range(4):
+                start = i * N_p
+                end = (i + 1) * N_p
+                if end <= patch_tokens.shape[1]:
+                    patch_tokens[:, start:end, :] = patch_tokens[:, start:end, :] + pos_emb
+            
+            x = torch.cat([cls_tokens, patch_tokens], dim=1)
+        
+        # 通过所有层直到倒数第二层
+        for i, block in enumerate(model.encoder[:-1]):
+            x = block(x, attn_mask=attn_mask)
+        
+        # 最后一层：手动提取注意力
+        last_block = model.encoder[-1]
+        x_norm = last_block.norm1(x)
+        B_, N_, C_ = x_norm.shape
+        
+        qkv = last_block.attn.qkv(x_norm).reshape(
+            B_, N_, 3, last_block.attn.num_heads, C_ // last_block.attn.num_heads
+        ).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)
+        
+        attn = (q @ k.transpose(-2, -1)) * last_block.attn.scale
+        if attn_mask is not None:
+            attn = attn + attn_mask.unsqueeze(1)
+        attn = attn.softmax(dim=-1)
+        
+        return attn
+
+
+def build_voxel_attention_map(
+    patch_attention: np.ndarray,
+    atlas_data: np.ndarray,
+    img_size: tuple = (128, 128, 128),
+    patch_size: tuple = (16, 16, 16),
+) -> np.ndarray:
+    """
+    将 per-patch 注意力映射回 voxel 空间
+    
+    方法：每个 patch 内的所有 voxel 赋予该 patch 的注意力值
+    
+    Args:
+        patch_attention: [N_patches] 每个 patch 的注意力值
+        atlas_data: [D, H, W] atlas 数据（用于验证 patch 边界）
+        img_size: 目标图像尺寸
+        patch_size: patch 大小
+    
+    Returns:
+        attention_map: [D, H, W] voxel-level 注意力图
+    """
+    D, H, W = img_size
+    pd, ph, pw = patch_size
+    
+    grid_d = D // pd
+    grid_h = H // ph
+    grid_w = W // pw
+    
+    attention_map = np.zeros((D, H, W), dtype=np.float32)
+    
+    # 遍历所有 patches
+    idx = 0
+    for d in range(grid_d):
+        for h in range(grid_h):
+            for w in range(grid_w):
+                if idx < len(patch_attention):
+                    # 该 patch 对应的 voxel 区域
+                    d_start, d_end = d * pd, (d + 1) * pd
+                    h_start, h_end = h * ph, (h + 1) * ph
+                    w_start, w_end = w * pw, (w + 1) * pw
+                    
+                    # 赋予注意力值
+                    attention_map[d_start:d_end, h_start:h_end, w_start:w_end] = patch_attention[idx]
+                idx += 1
+    
+    return attention_map
+
+
+def visualize_attention_brain_regions(
+    brain_volume: np.ndarray,
+    attention_map: np.ndarray,
+    brain_mask: np.ndarray = None,
+    slice_indices: list = None,
+    plane: str = 'axial',
+    alpha: float = 0.6,
+    cmap: str = 'viridis',
+    output_path: str = None,
+    show_colorbar: bool = False,
+    relative_mode: bool = False,
+) -> np.ndarray:
+    """
+    可视化脑区注意力分布（带脑轮廓 mask）
+    
+    Args:
+        brain_volume: [D, H, W] 原始脑影像
+        attention_map: [D, H, W] 注意力图
+        brain_mask: [D, H, W] 脑轮廓 mask (非零=脑内, 零=背景)
+        slice_indices: 要显示的切片索引列表
+        plane: 'axial', 'coronal', 'sagittal'
+        alpha: 叠加透明度
+        cmap: 热力图颜色映射，默认 'viridis'
+        output_path: 保存路径
+        show_colorbar: 是否显示 colorbar，默认 False
+        relative_mode: 是否显示相对注意力（与平均值差异），默认 False
+    
+    Returns:
+        可视化图像数组
+    """
+    # 如果没有提供 mask，使用图像非零区域作为 mask
+    if brain_mask is None:
+        brain_mask = (brain_volume > 0.01).astype(np.float32)
+    
+    # 确保 mask 是二值的
+    brain_mask = (brain_mask > 0).astype(np.float32)
+    
+    if slice_indices is None:
+        # 默认显示中间附近的切片
+        if plane == 'axial':
+            max_idx = brain_volume.shape[2]
+        elif plane == 'coronal':
+            max_idx = brain_volume.shape[1]
+        else:
+            max_idx = brain_volume.shape[0]
+        
+        slice_indices = [
+            max_idx // 4,
+            max_idx * 3 // 8,
+            max_idx // 2,
+            max_idx * 5 // 8,
+            max_idx * 3 // 4,
+        ]
+    
+    n_slices = len(slice_indices)
+    
+    # 创建图形 - 纯白背景
+    fig, axes = plt.subplots(2, n_slices, figsize=(4*n_slices, 8), 
+                             facecolor='white', edgecolor='white')
+    
+    for i, idx in enumerate(slice_indices):
+        # 提取切片
+        if plane == 'axial':
+            brain_slice = brain_volume[:, :, idx]
+            attn_slice = attention_map[:, :, idx]
+            mask_slice = brain_mask[:, :, idx]
+        elif plane == 'coronal':
+            brain_slice = brain_volume[:, idx, :]
+            attn_slice = attention_map[:, idx, :]
+            mask_slice = brain_mask[:, idx, :]
+        else:  # sagittal
+            brain_slice = brain_volume[idx, :, :]
+            attn_slice = attention_map[idx, :, :]
+            mask_slice = brain_mask[idx, :, :]
+        
+        # 创建白色背景图像
+        white_bg = np.ones_like(brain_slice) * 255  # 纯白背景 (8-bit)
+        
+        # 上一行：原始脑影像（mask 外白色，mask 内显示影像）
+        ax1 = axes[0, i]
+        # 白色背景
+        display_img = np.stack([white_bg, white_bg, white_bg], axis=-1).astype(np.uint8)
+        # 在 mask 内填入脑影像（灰度转 RGB）
+        brain_norm = (brain_slice * 255).astype(np.uint8)
+        mask_bool = mask_slice > 0
+        for c in range(3):
+            display_img[:, :, c] = np.where(mask_bool, brain_norm, display_img[:, :, c])
+        ax1.imshow(display_img, origin='lower')
+        ax1.set_title(f'{plane.capitalize()} {idx}', fontsize=12)
+        ax1.axis('off')
+        ax1.set_facecolor('white')
+        
+        # 下一行：注意力叠加（mask 外白色，mask 内显示叠加）
+        ax2 = axes[1, i]
+        ax2.set_facecolor('white')
+        
+        # 获取 colormap
+        from matplotlib import cm
+        
+        # 根据模式选择 colormap 和归一化方式
+        if relative_mode:
+            # 相对模式：显示与平均值的差异，使用 diverging colormap
+            cmap_obj = cm.get_cmap('RdBu_r')  # 红-白-蓝，红=高，蓝=低
+            
+            # 计算 mask 内的平均值
+            brain_mask_global = brain_mask if brain_mask is not None else (brain_volume > 0.01)
+            masked_attn = attention_map[brain_mask_global > 0]
+            if len(masked_attn) > 0:
+                mean_attn = masked_attn.mean()
+            else:
+                mean_attn = attention_map.mean()
+            
+            # 计算相对值（差异）
+            attn_relative = attn_slice - mean_attn
+            
+            # 对称归一化：以 0 为中心，正负范围相同
+            max_abs = np.abs(attn_relative).max()
+            if max_abs > 0:
+                attn_norm = (attn_relative / max_abs) * 0.5 + 0.5  # 映射到 [0, 1]，0.5 为中心
+            else:
+                attn_norm = np.ones_like(attn_relative) * 0.5
+            
+            # 在 mask 外设为中性色（0.5 = 白色）
+            attn_norm = np.where(mask_bool, attn_norm, 0.5)
+        else:
+            # 绝对模式：显示原始注意力值
+            cmap_obj = cm.get_cmap(cmap)
+            
+            # 归一化注意力值
+            attn_min, attn_max = attention_map.min(), attention_map.max()
+            if attn_max > attn_min:
+                attn_norm = (attn_slice - attn_min) / (attn_max - attn_min)
+            else:
+                attn_norm = np.zeros_like(attn_slice)
+        
+        # 应用 colormap 到注意力图
+        attn_colored = cmap_obj(attn_norm)  # [H, W, 4] (RGBA)
+        attn_colored = (attn_colored[:, :, :3] * 255).astype(np.uint8)  # 转为 RGB
+        
+        # 混合：白色背景 + mask 内叠加
+        display_attn = np.stack([white_bg, white_bg, white_bg], axis=-1).astype(np.uint8)
+        
+        # 在 mask 内混合原始影像和注意力
+        for c in range(3):
+            if relative_mode:
+                # 相对模式：更多显示 colormap，较少原始影像
+                blended = (1 - alpha * 0.5) * brain_norm + alpha * 1.5 * attn_colored[:, :, c]
+                blended = np.clip(blended, 0, 255)
+            else:
+                # 绝对模式：标准混合
+                blended = (1 - alpha) * brain_norm + alpha * attn_colored[:, :, c]
+            display_attn[:, :, c] = np.where(mask_bool, blended.astype(np.uint8), 
+                                              display_attn[:, :, c])
+        
+        ax2.imshow(display_attn, origin='lower')
+        ax2.axis('off')
+        
+        # 可选 colorbar（默认不显示）
+        if show_colorbar and i == n_slices - 1:
+            from mpl_toolkits.axes_grid1 import make_axes_locatable
+            divider = make_axes_locatable(ax2)
+            cax = divider.append_axes('right', size='5%', pad=0.05)
+            sm = plt.cm.ScalarMappable(cmap=cmap, 
+                                        norm=plt.Normalize(vmin=attn_min, vmax=attn_max))
+            sm.set_array([])
+            plt.colorbar(sm, cax=cax)
+    
+    plt.tight_layout()
+    
+    # 转换图像
+    fig.canvas.draw()
+    img = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
+    img = img.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+    
+    # 保存 - 纯白背景
+    if output_path:
+        plt.savefig(output_path, dpi=150, bbox_inches='tight', 
+                   facecolor='white', edgecolor='white')
+        print(f"  Saved: {output_path}")
+    
+    plt.close()
+    
+    return img
+
+
+def visualize_sample_attention_precise(
+    engine: BrainAnytimeInference,
+    sample_dir: Path,
+    task: str,
+    combo: str,
+    output_dir: Path = None,
+):
+    """
+    精确可视化样本的注意力分布
+    
+    Args:
+        engine: 推理引擎
+        sample_dir: 样本目录
+        task: 任务名称
+        combo: 模态组合
+        output_dir: 输出目录
+    
+    Returns:
+        生成的图像路径列表
+    """
+    import json
+    
+    # 加载元数据
+    meta_path = list(sample_dir.glob("*_meta.json"))[0]
+    with open(meta_path) as f:
+        meta = json.load(f)
+    
+    print(f"\nProcessing: {meta['subject_id']} - {task}/{combo}")
+    print(f"  Modalities: {meta['modalities']}")
+    
+    # 加载模型
+    model = engine.load_model(task)
+    
+    # 准备输入
+    from inference_engine import SHORT_TO_FULL
+    
+    # 加载脑影像
+    brain_volumes = {}
+    nifti_files = {}
+    
+    for mod in meta['modalities']:
+        nii_filename = meta['files']['nifti'][mod]
+        nii_path = sample_dir / nii_filename
+        
+        nii = nib.load(str(nii_path))
+        data = nii.get_fdata().astype(np.float32)
+        
+        # Min-Max 归一化
+        data_min, data_max = data.min(), data.max()
+        if data_max > data_min:
+            data = (data - data_min) / (data_max - data_min)
+        
+        brain_volumes[mod] = data
+        nifti_files[mod] = str(nii_path)
+    
+    # 构建输入张量
+    images_list = []
+    for mod in MODALITY_ORDER:
+        if mod in brain_volumes:
+            images_list.append(brain_volumes[mod])
+        else:
+            images_list.append(np.zeros((128, 128, 128), dtype=np.float32))
+    
+    images = np.stack(images_list, axis=0)
+    images = torch.from_numpy(images).unsqueeze(0).to(engine.device)
+    
+    # observed mask
+    observed = torch.zeros(1, 4)
+    for i, mod in enumerate(MODALITY_ORDER):
+        if mod in meta['modalities']:
+            observed[0, i] = 1.0
+    
+    # 执行推理
+    with torch.no_grad():
+        logits = model(images, observed)
+        prob = torch.sigmoid(logits).item() if task in ['CN_vs_AD', 'CN_vs_MCI'] else logits.item()
+    
+    print(f"  Prediction: {prob:.4f}")
+    
+    # 提取注意力
+    attention = extract_last_layer_attention(model.encoder, images, observed)
+    
+    if attention is None:
+        print("  Failed to extract attention")
+        return []
+    
+    # 处理注意力
+    attn_batch = attention[0].mean(dim=0)  # [N, N]
+    cls_to_all = attn_batch[0, 1:]  # 跳过 CLS
+    
+    # 分配到模态
+    num_global = model.encoder.num_global_tokens
+    N_p = model.encoder.num_patches  # 512
+    
+    # 计算实际 patches 数
+    total_patches = cls_to_all.shape[0]
+    actual_patches = (total_patches // 4) * 4
+    N_p_actual = actual_patches // 4
+    
+    print(f"  Patches per modality: {N_p_actual}")
+    
+    # 输出路径
+    if output_dir is None:
+        output_dir = sample_dir
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    
+    generated_files = []
+    
+    # 为每个模态生成可视化
+    for idx, mod in enumerate(meta['modalities']):
+        if mod not in brain_volumes:
+            continue
+        
+        mod_idx = MODALITY_ORDER.index(mod)
+        start_idx = mod_idx * N_p_actual
+        end_idx = (mod_idx + 1) * N_p_actual
+        
+        if end_idx > len(cls_to_all):
+            print(f"  Warning: {mod} index out of range")
+            continue
+        
+        mod_attn = cls_to_all[start_idx:end_idx].cpu().numpy()
+        
+        # 使用精确映射构建 voxel-level 注意力图
+        # 创建虚拟 atlas（因为我们没有精确的 patch 位置映射）
+        atlas_data = np.ones((128, 128, 128), dtype=np.int32)
+        
+        attention_map = build_voxel_attention_map(
+            mod_attn,
+            atlas_data,
+            img_size=(128, 128, 128),
+            patch_size=(16, 16, 16),
+        )
+        
+        # 平滑处理（减少 blockiness）
+        from scipy.ndimage import gaussian_filter
+        attention_map = gaussian_filter(attention_map, sigma=2.0)
+        
+        print(f"  {mod}: attention range [{mod_attn.min():.6f}, {mod_attn.max():.6f}]")
+        
+        # 生成可视化
+        output_path = output_dir / f"attention_{task}_{combo}_{mod}_{meta['subject_id']}.png"
+        
+        visualize_attention_brain_regions(
+            brain_volumes[mod],
+            attention_map,
+            slice_indices=[40, 55, 70, 85, 100],
+            plane='axial',
+            alpha=0.5,
+            output_path=str(output_path),
+        )
+        
+        generated_files.append(output_path)
+    
+    return generated_files
+
+
+def main():
+    """主函数：测试可视化"""
+    import argparse
+    
+    parser = argparse.ArgumentParser(description='Brain Attention Visualization')
+    parser.add_argument('--sample_dir', type=str,
+                        default='demo_samples/CN_vs_AD/TMFP/sample_005',
+                        help='Sample directory')
+    parser.add_argument('--task', type=str, default='CN_vs_AD',
+                        choices=list(TASKS.keys()),
+                        help='Task name')
+    parser.add_argument('--combo', type=str, default='TMFP',
+                        help='Modality combination')
+    parser.add_argument('--output_dir', type=str, default=None,
+                        help='Output directory (default: sample_dir)')
+    
+    args = parser.parse_args()
+    
+    print("="*70)
+    print("Brain Attention Visualization")
+    print("="*70)
+    
+    # 初始化引擎
+    checkpoints_dir = "/home/23037125r/code/random/multimae_freeze_then_finetune/"
+    engine = BrainAnytimeInference(checkpoints_dir=checkpoints_dir)
+    
+    # 生成可视化
+    sample_dir = Path(args.sample_dir)
+    output_dir = Path(args.output_dir) if args.output_dir else None
+    
+    files = visualize_sample_attention_precise(
+        engine,
+        sample_dir,
+        args.task,
+        args.combo,
+        output_dir,
+    )
+    
+    print("\n" + "="*70)
+    print(f"Generated {len(files)} visualizations:")
+    for f in files:
+        print(f"  - {f}")
+    print("="*70)
+
+
+if __name__ == '__main__':
+    main()