Add online inference demo with 20 samples

BrainAnytime/altas/AAL116_standard.nii.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1a1e8d16d6b7616f3904f4319683692dda934010a52ace85de082d12702d2732
+size 203563

BrainAnytime/anatomy_masking.py

@@ -0,0 +1,595 @@
+"""
+Anatomy-Aware Adaptive Masking for MultiMAE3D Pretraining.
+
+Four components:
+1. Patch-Region Mapping: Maps 512 patches (8x8x8 grid) to AAL116 brain atlas regions
+2. Region Importance Scoring: Static (AD prior) + Dynamic (EMA teacher attention)
+3. Mask Probability Generation: Softmax with temperature control
+4. Curriculum Scheduler: Three-phase training schedule
+
+Usage:
+    masking = AnatomyAwareMasking(
+        img_size=128, patch_size=16,
+        atlas_path='altas/AAL116_standard.nii.gz',
+    )
+
+    # In training loop:
+    mask_probs = masking.get_mask_probs(epoch, total_epochs)
+    output = model(images, observed, patch_mask_probs=mask_probs)
+
+    # EMA teacher attention update (every N iterations):
+    attn = extract_teacher_attention(ema_teacher, images, observed)
+    masking.update_dynamic_importance(attn)
+"""
+
+import os
+import math
+import copy
+import numpy as np
+import torch
+
+try:
+    import nibabel as nib
+    HAS_NIBABEL = True
+except ImportError:
+    HAS_NIBABEL = False
+
+
+# =============================================================================
+# AAL116 Atlas: Label-to-Name Mapping + AD Importance
+# =============================================================================
+
+AAL116_LABEL_NAMES = {
+    1: 'Precentral_L', 2: 'Precentral_R',
+    3: 'Frontal_Sup_L', 4: 'Frontal_Sup_R',
+    5: 'Frontal_Sup_Orb_L', 6: 'Frontal_Sup_Orb_R',
+    7: 'Frontal_Mid_L', 8: 'Frontal_Mid_R',
+    9: 'Frontal_Mid_Orb_L', 10: 'Frontal_Mid_Orb_R',
+    11: 'Frontal_Inf_Oper_L', 12: 'Frontal_Inf_Oper_R',
+    13: 'Frontal_Inf_Tri_L', 14: 'Frontal_Inf_Tri_R',
+    15: 'Frontal_Inf_Orb_L', 16: 'Frontal_Inf_Orb_R',
+    17: 'Rolandic_Oper_L', 18: 'Rolandic_Oper_R',
+    19: 'Supp_Motor_Area_L', 20: 'Supp_Motor_Area_R',
+    21: 'Olfactory_L', 22: 'Olfactory_R',
+    23: 'Frontal_Sup_Medial_L', 24: 'Frontal_Sup_Medial_R',
+    25: 'Frontal_Med_Orb_L', 26: 'Frontal_Med_Orb_R',
+    27: 'Rectus_L', 28: 'Rectus_R',
+    29: 'Insula_L', 30: 'Insula_R',
+    31: 'Cingulum_Ant_L', 32: 'Cingulum_Ant_R',
+    33: 'Cingulum_Mid_L', 34: 'Cingulum_Mid_R',
+    35: 'Cingulum_Post_L', 36: 'Cingulum_Post_R',
+    37: 'Hippocampus_L', 38: 'Hippocampus_R',
+    39: 'ParaHippocampal_L', 40: 'ParaHippocampal_R',
+    41: 'Amygdala_L', 42: 'Amygdala_R',
+    43: 'Calcarine_L', 44: 'Calcarine_R',
+    45: 'Cuneus_L', 46: 'Cuneus_R',
+    47: 'Lingual_L', 48: 'Lingual_R',
+    49: 'Occipital_Sup_L', 50: 'Occipital_Sup_R',
+    51: 'Occipital_Mid_L', 52: 'Occipital_Mid_R',
+    53: 'Occipital_Inf_L', 54: 'Occipital_Inf_R',
+    55: 'Fusiform_L', 56: 'Fusiform_R',
+    57: 'Postcentral_L', 58: 'Postcentral_R',
+    59: 'Parietal_Sup_L', 60: 'Parietal_Sup_R',
+    61: 'Parietal_Inf_L', 62: 'Parietal_Inf_R',
+    63: 'SupraMarginal_L', 64: 'SupraMarginal_R',
+    65: 'Angular_L', 66: 'Angular_R',
+    67: 'Precuneus_L', 68: 'Precuneus_R',
+    69: 'Paracentral_Lobule_L', 70: 'Paracentral_Lobule_R',
+    71: 'Caudate_L', 72: 'Caudate_R',
+    73: 'Putamen_L', 74: 'Putamen_R',
+    75: 'Pallidum_L', 76: 'Pallidum_R',
+    77: 'Thalamus_L', 78: 'Thalamus_R',
+    79: 'Heschl_L', 80: 'Heschl_R',
+    81: 'Temporal_Sup_L', 82: 'Temporal_Sup_R',
+    83: 'Temporal_Pole_Sup_L', 84: 'Temporal_Pole_Sup_R',
+    85: 'Temporal_Mid_L', 86: 'Temporal_Mid_R',
+    87: 'Temporal_Pole_Mid_L', 88: 'Temporal_Pole_Mid_R',
+    89: 'Temporal_Inf_L', 90: 'Temporal_Inf_R',
+    91: 'Cerebelum_Crus1_L', 92: 'Cerebelum_Crus1_R',
+    93: 'Cerebelum_Crus2_L', 94: 'Cerebelum_Crus2_R',
+    95: 'Cerebelum_3_L', 96: 'Cerebelum_3_R',
+    97: 'Cerebelum_4_5_L', 98: 'Cerebelum_4_5_R',
+    99: 'Cerebelum_6_L', 100: 'Cerebelum_6_R',
+    101: 'Cerebelum_7b_L', 102: 'Cerebelum_7b_R',
+    103: 'Cerebelum_8_L', 104: 'Cerebelum_8_R',
+    105: 'Cerebelum_9_L', 106: 'Cerebelum_9_R',
+    107: 'Cerebelum_10_L', 108: 'Cerebelum_10_R',
+    109: 'Vermis_1_2', 110: 'Vermis_3',
+    111: 'Vermis_4_5', 112: 'Vermis_6',
+    113: 'Vermis_7', 114: 'Vermis_8',
+    115: 'Vermis_9', 116: 'Vermis_10',
+}
+
+# AD-relevant regions: base name (without _L/_R) -> importance level
+# Based on Braak staging and AD pathology literature
+AD_REGION_IMPORTANCE = {
+    # Hippocampus (Braak III-IV)
+    'Hippocampus': 'high',
+    # Parahippocampal / Entorhinal cortex (Braak I-II, earliest involvement)
+    'ParaHippocampal': 'high',
+    # Amygdala (Braak III-IV)
+    'Amygdala': 'high',
+    # Posterior cingulate cortex (early metabolic changes in AD)
+    'Cingulum_Post': 'high',
+    # Precuneus (default mode network hub, early amyloid deposition)
+    'Precuneus': 'high',
+    # Inferior temporal (early cortical atrophy)
+    'Temporal_Inf': 'high',
+    # Middle temporal
+    'Temporal_Mid': 'high',
+    # Fusiform gyrus
+    'Fusiform': 'high',
+    # Angular gyrus (default mode network)
+    'Angular': 'high',
+    # Medial orbitofrontal (default mode network)
+    'Frontal_Med_Orb': 'high',
+    # Temporal poles
+    'Temporal_Pole_Sup': 'high',
+    'Temporal_Pole_Mid': 'high',
+    # Insula
+    'Insula': 'high',
+    # Thalamus (subcortical relay)
+    'Thalamus': 'high',
+    # Caudate (striatal amyloid)
+    'Caudate': 'high',
+}
+
+
+def _get_region_importance(region_name):
+    """Match an AAL116 region name to its AD importance level."""
+    base = region_name
+    if base.endswith('_L') or base.endswith('_R'):
+        base = base[:-2]
+    return AD_REGION_IMPORTANCE.get(base, 'mid')
+
+
+# =============================================================================
+# Patch-Region Mapping
+# =============================================================================
+
+def build_patch_region_mapping(atlas_data, img_size, patch_size):
+    """Build mapping from 3D patches to atlas regions.
+
+    For each patch, computes the fraction of voxels belonging to each region.
+    Patch ordering matches einops rearrange:
+      "b c (nd pd) (nh ph) (nw pw) -> b (nd nh nw) c pd ph pw"
+      patch_index = d_idx * (grid_h * grid_w) + h_idx * grid_w + w_idx
+
+    Args:
+        atlas_data: [D, H, W] integer numpy array (0 = background)
+        img_size: (D, H, W) tuple
+        patch_size: (pd, ph, pw) tuple
+
+    Returns:
+        membership: [N_patches, K] float32 tensor (region membership fractions)
+        region_labels: sorted list of unique non-zero integer labels
+    """
+    grid = tuple(img_size[i] // patch_size[i] for i in range(3))
+    N = grid[0] * grid[1] * grid[2]
+
+    labels = sorted([int(l) for l in np.unique(atlas_data) if l > 0])
+    K = len(labels)
+    label_to_idx = {l: i for i, l in enumerate(labels)}
+
+    membership = np.zeros((N, K), dtype=np.float32)
+    voxels_per_patch = patch_size[0] * patch_size[1] * patch_size[2]
+
+    patch_idx = 0
+    for d in range(grid[0]):
+        for h in range(grid[1]):
+            for w in range(grid[2]):
+                block = atlas_data[
+                    d * patch_size[0]:(d + 1) * patch_size[0],
+                    h * patch_size[1]:(h + 1) * patch_size[1],
+                    w * patch_size[2]:(w + 1) * patch_size[2],
+                ].flatten()
+
+                for label in labels:
+                    count = np.sum(block == label)
+                    if count > 0:
+                        membership[patch_idx, label_to_idx[label]] = count / voxels_per_patch
+
+                patch_idx += 1
+
+    return torch.from_numpy(membership), labels
+
+
+# =============================================================================
+# Main Class
+# =============================================================================
+
+class AnatomyAwareMasking:
+    """Anatomy-aware adaptive masking with curriculum learning.
+
+    Args:
+        img_size: Input volume size (default 128)
+        patch_size: Patch size (default 16)
+        atlas_path: Path to AAL116 atlas NIfTI (128x128x128, labels 1-116)
+        w_high / w_mid / w_low: Importance weights for AD-critical / gray matter / non-brain regions
+        temperature_target: Final temperature for softmax (lower = more focused masking)
+        temperature_start: Starting temperature at Phase 2 onset
+        phase1_end: End of Phase 1 (uniform masking) as fraction of total epochs
+        phase2_end: End of Phase 2 (transition) as fraction of total epochs
+        ema_momentum: EMA momentum for teacher model updates
+        attention_update_freq: Extract teacher attention every N training iterations
+        teacher_batch_size: Number of samples for teacher attention extraction
+        importance_mode: 'static', 'dynamic', or 'combined'
+        dynamic_weight: Weight of dynamic importance in combined mode [0, 1]
+    """
+
+    def __init__(
+        self,
+        img_size=128,
+        patch_size=16,
+        atlas_path=None,
+        w_high=3.0,
+        w_mid=1.5,
+        w_low=0.3,
+        temperature_target=1.0,
+        temperature_start=5.0,
+        phase1_end=0.2,
+        phase2_end=0.7,
+        ema_momentum=0.998,
+        attention_update_freq=200,
+        teacher_batch_size=2,
+        importance_mode='combined',
+        dynamic_weight=0.5,
+    ):
+        self.img_size = (img_size,) * 3 if isinstance(img_size, int) else tuple(img_size)
+        self.patch_size = (patch_size,) * 3 if isinstance(patch_size, int) else tuple(patch_size)
+        self.grid = tuple(self.img_size[i] // self.patch_size[i] for i in range(3))
+        self.num_patches = self.grid[0] * self.grid[1] * self.grid[2]
+
+        self.w_high = w_high
+        self.w_mid = w_mid
+        self.w_low = w_low
+
+        self.temperature_target = temperature_target
+        self.temperature_start = temperature_start
+        self.phase1_end = phase1_end
+        self.phase2_end = phase2_end
+
+        self.ema_momentum = ema_momentum
+        self.attention_update_freq = attention_update_freq
+        self.teacher_batch_size = teacher_batch_size
+
+        self.importance_mode = importance_mode
+        self.dynamic_weight = dynamic_weight
+
+        # Internal state
+        self.patch_region_membership = None   # [N, K]
+        self.region_labels = None             # list[int]
+        self.static_importance = None         # [N]
+        self.dynamic_region_importance = None  # [K] running average
+        self.dynamic_patch_importance = None   # [N] fallback without atlas
+
+        if atlas_path is not None:
+            self._load_atlas(atlas_path)
+            self._compute_static_importance()
+
+    # -----------------------------------------------------------------
+    # Atlas loading and static importance
+    # -----------------------------------------------------------------
+
+    def _load_atlas(self, atlas_path):
+        if not HAS_NIBABEL:
+            raise ImportError("nibabel required for atlas loading: pip install nibabel")
+        if not os.path.exists(atlas_path):
+            raise FileNotFoundError(f"Atlas not found: {atlas_path}")
+
+        atlas_img = nib.load(atlas_path)
+        atlas_data = np.asarray(atlas_img.dataobj, dtype=np.int32)
+
+        if atlas_data.shape != self.img_size:
+            raise ValueError(
+                f"Atlas shape {atlas_data.shape} != expected {self.img_size}. "
+                f"Resample the atlas to match your data dimensions."
+            )
+
+        self.patch_region_membership, self.region_labels = build_patch_region_mapping(
+            atlas_data, self.img_size, self.patch_size
+        )
+
+    def _compute_static_importance(self):
+        """Compute per-patch static importance: s_i = sum_k(r_{i,k} * w_k)."""
+        if self.patch_region_membership is None:
+            return
+
+        K = len(self.region_labels)
+        region_weights = torch.zeros(K)
+        for i, label in enumerate(self.region_labels):
+            name = AAL116_LABEL_NAMES.get(label, f"Region_{label}")
+            level = _get_region_importance(name)
+            if level == 'high':
+                region_weights[i] = self.w_high
+            elif level == 'mid':
+                region_weights[i] = self.w_mid
+            else:
+                region_weights[i] = self.w_low
+
+        self.static_importance = self.patch_region_membership @ region_weights  # [N]
+
+        # Penalize non-brain patches (< 10% brain coverage)
+        brain_coverage = self.patch_region_membership.sum(dim=1)
+        non_brain = brain_coverage < 0.1
+        self.static_importance[non_brain] = self.w_low * 0.5
+
+    # -----------------------------------------------------------------
+    # Dynamic importance from EMA teacher
+    # -----------------------------------------------------------------
+
+    def _aggregate_to_regions(self, patch_attention):
+        """Aggregate per-patch attention to region level.
+
+        w_k = (1/|P_k|) * sum_{i in P_k} a_i
+        """
+        if self.patch_region_membership is None:
+            return None
+        M = self.patch_region_membership  # [N, K]
+        numerator = M.t() @ patch_attention  # [K]
+        denominator = M.sum(dim=0).clamp(min=1e-8)  # [K]
+        return numerator / denominator
+
+    def update_dynamic_importance(self, patch_attention):
+        """Update dynamic importance from EMA teacher CLS attention.
+
+        Aggregates to region level for smoothing if atlas available,
+        otherwise uses raw per-patch attention.
+        """
+        patch_attention = patch_attention.detach().cpu()
+        momentum = 0.9
+
+        if self.patch_region_membership is not None:
+            region_imp = self._aggregate_to_regions(patch_attention)
+            if self.dynamic_region_importance is None:
+                self.dynamic_region_importance = region_imp
+            else:
+                self.dynamic_region_importance = (
+                    momentum * self.dynamic_region_importance
+                    + (1 - momentum) * region_imp
+                )
+        else:
+            if self.dynamic_patch_importance is None:
+                self.dynamic_patch_importance = patch_attention
+            else:
+                self.dynamic_patch_importance = (
+                    momentum * self.dynamic_patch_importance
+                    + (1 - momentum) * patch_attention
+                )
+
+    def _get_dynamic_scores(self):
+        """Convert dynamic importance to per-patch scores."""
+        if self.dynamic_region_importance is not None and self.patch_region_membership is not None:
+            return self.patch_region_membership @ self.dynamic_region_importance
+        return self.dynamic_patch_importance
+
+    # -----------------------------------------------------------------
+    # Temperature and curriculum
+    # -----------------------------------------------------------------
+
+    def get_temperature(self, epoch, total_epochs):
+        """Three-phase curriculum temperature.
+
+        Phase 1 (0 to phase1_end):      tau = inf  (uniform masking)
+        Phase 2 (phase1_end to phase2_end): cosine anneal start -> target
+        Phase 3 (phase2_end to 1.0):     tau = target (stable)
+        """
+        progress = epoch / max(total_epochs, 1)
+        if progress < self.phase1_end:
+            return float('inf')
+        elif progress < self.phase2_end:
+            phase_progress = (progress - self.phase1_end) / (self.phase2_end - self.phase1_end)
+            return self.temperature_target + 0.5 * (
+                self.temperature_start - self.temperature_target
+            ) * (1.0 + math.cos(math.pi * phase_progress))
+        else:
+            return self.temperature_target
+
+    # -----------------------------------------------------------------
+    # Combined importance scores
+    # -----------------------------------------------------------------
+
+    def get_importance_scores(self):
+        """Get combined per-patch importance scores based on importance_mode."""
+        static = self.static_importance
+        dynamic = self._get_dynamic_scores()
+
+        if self.importance_mode == 'static':
+            return static
+        elif self.importance_mode == 'dynamic':
+            return dynamic if dynamic is not None else static
+        else:  # combined
+            if static is None and dynamic is None:
+                return None
+            if dynamic is None:
+                return static
+            if static is None:
+                return dynamic
+            # Normalize both to [0, 1] before combining
+            s_norm = (static - static.min()) / (static.max() - static.min() + 1e-8)
+            d_norm = (dynamic - dynamic.min()) / (dynamic.max() - dynamic.min() + 1e-8)
+            alpha = self.dynamic_weight
+            return (1 - alpha) * s_norm + alpha * d_norm
+
+    # -----------------------------------------------------------------
+    # Main API
+    # -----------------------------------------------------------------
+
+    def get_mask_probs(self, epoch, total_epochs):
+        """Get per-patch masking probabilities.
+
+        Returns:
+            [N_patches] tensor (sums to 1), or None for uniform masking.
+            Higher value = more likely to be masked.
+        """
+        tau = self.get_temperature(epoch, total_epochs)
+        if tau == float('inf'):
+            return None
+
+        scores = self.get_importance_scores()
+        if scores is None:
+            return None
+
+        return torch.softmax(scores / tau, dim=0)
+
+    def get_curriculum_info(self, epoch, total_epochs):
+        """Get curriculum state for logging."""
+        tau = self.get_temperature(epoch, total_epochs)
+        progress = epoch / max(total_epochs, 1)
+
+        if progress < self.phase1_end:
+            phase = 1
+        elif progress < self.phase2_end:
+            phase = 2
+        else:
+            phase = 3
+
+        info = {'phase': phase, 'temperature': tau if tau != float('inf') else -1.0}
+
+        scores = self.get_importance_scores()
+        if scores is not None:
+            info['importance_min'] = scores.min().item()
+            info['importance_max'] = scores.max().item()
+            info['importance_mean'] = scores.mean().item()
+
+        probs = self.get_mask_probs(epoch, total_epochs)
+        if probs is not None:
+            info['prob_max'] = probs.max().item()
+            info['prob_min'] = probs.min().item()
+            info['prob_ratio'] = (probs.max() / probs.min().clamp(min=1e-10)).item()
+
+        return info
+
+    # -----------------------------------------------------------------
+    # Checkpointing
+    # -----------------------------------------------------------------
+
+    def state_dict(self):
+        return {
+            'static_importance': self.static_importance,
+            'dynamic_region_importance': self.dynamic_region_importance,
+            'dynamic_patch_importance': self.dynamic_patch_importance,
+        }
+
+    def load_state_dict(self, state_dict):
+        if state_dict is None:
+            return
+        self.static_importance = state_dict.get('static_importance')
+        self.dynamic_region_importance = state_dict.get('dynamic_region_importance')
+        self.dynamic_patch_importance = state_dict.get('dynamic_patch_importance')
+
+
+# =============================================================================
+# EMA Teacher Utilities
+# =============================================================================
+
+@torch.no_grad()
+def create_ema_teacher(model):
+    """Create an EMA copy of the model (no gradients)."""
+    teacher = copy.deepcopy(model)
+    for p in teacher.parameters():
+        p.requires_grad = False
+    return teacher
+
+
+@torch.no_grad()
+def update_ema_teacher(teacher, student, momentum=0.998):
+    """Update EMA teacher: theta_t = m * theta_t + (1 - m) * theta_s."""
+    student_model = student.module if hasattr(student, 'module') else student
+    teacher_model = teacher.module if hasattr(teacher, 'module') else teacher
+    for t_param, s_param in zip(teacher_model.parameters(), student_model.parameters()):
+        t_param.data.mul_(momentum).add_(s_param.data, alpha=1 - momentum)
+
+
+@torch.no_grad()
+def extract_teacher_attention(teacher, images, observed, num_global_tokens=1):
+    """Extract CLS-to-patch attention from the EMA teacher's last encoder layer.
+
+    Runs a full (unmasked) forward pass through the teacher encoder and
+    extracts attention weights from the final transformer block.
+
+    Args:
+        teacher: EMA teacher model (MultiMAE3D, not DDP-wrapped)
+        images: [B, 4, D, H, W]
+        observed: [B, 4]
+        num_global_tokens: number of CLS tokens (default 1)
+
+    Returns:
+        patch_attention: [num_patches] averaged CLS attention scores
+    """
+    from models.multimae3d_utils import patchify
+
+    teacher_model = teacher.module if hasattr(teacher, 'module') else teacher
+    teacher_model.eval()
+
+    B = images.shape[0]
+    device = images.device
+    batch = teacher_model._split_modalities(images)
+
+    # Tokenize all patches (no masking)
+    tokens_list = []
+    for i, name in enumerate(teacher_model.MODALITY_NAMES):
+        patches = patchify(batch[name], teacher_model.patch_size)
+        tok = teacher_model.input_adapters[name](patches)
+        pos_emb = teacher_model.pos_embed.expand(B, -1, -1)
+        tok = tok + pos_emb
+        mod_mask = observed[:, i:i + 1].unsqueeze(-1)
+        tok = tok * mod_mask
+        tokens_list.append(tok)
+
+    input_tokens = torch.cat(tokens_list, dim=1)
+
+    if teacher_model.num_global_tokens > 0:
+        cls = teacher_model.global_tokens.unsqueeze(0).expand(B, -1, -1)
+        input_tokens = torch.cat([cls, input_tokens], dim=1)
+
+    # Attention mask for missing modalities
+    total_tokens = input_tokens.shape[1]
+    num_patches = teacher_model.num_patches
+    attn_mask = torch.zeros(B, 1, 1, total_tokens, device=device)
+    mod_offset = num_global_tokens
+    for i in range(len(teacher_model.MODALITY_NAMES)):
+        start, end = mod_offset, mod_offset + num_patches
+        missing = (observed[:, i] < 0.5)
+        if missing.any():
+            attn_mask[missing, :, :, start:end] = float("-inf")
+        mod_offset = end
+    if (attn_mask == 0).all():
+        attn_mask = None
+
+    # Forward through encoder layers 0..L-2
+    x = input_tokens
+    for block in teacher_model.encoder[:-1]:
+        x = block(x, attn_mask=attn_mask)
+
+    # Last layer: manually extract attention weights
+    last_block = teacher_model.encoder[-1]
+    x_norm = last_block.norm1(x)
+    B_, N_, C_ = x_norm.shape
+    num_heads = last_block.attn.num_heads
+    head_dim = C_ // num_heads
+
+    qkv = last_block.attn.qkv(x_norm)
+    qkv = qkv.reshape(B_, N_, 3, num_heads, head_dim).permute(2, 0, 3, 1, 4)
+    q, k, _ = qkv.unbind(0)
+
+    attn_weights = (q @ k.transpose(-2, -1)) * last_block.attn.scale
+    if attn_mask is not None:
+        attn_weights = attn_weights + attn_mask
+    attn_weights = attn_weights.softmax(dim=-1)  # [B, heads, N, N]
+
+    # CLS (token 0) attention to patch tokens (skip global tokens)
+    cls_attn = attn_weights[:, :, 0, num_global_tokens:]  # [B, heads, 4*num_patches]
+    cls_attn = cls_attn.mean(dim=1)  # avg over heads: [B, 4*num_patches]
+
+    # Reshape to per-modality and average
+    num_modalities = len(teacher_model.MODALITY_NAMES)
+    per_mod_attn = cls_attn.reshape(B, num_modalities, num_patches)
+
+    observed_expanded = observed.unsqueeze(-1)  # [B, 4, 1]
+    weighted_attn = (per_mod_attn * observed_expanded).sum(dim=1)  # [B, num_patches]
+    count = observed.sum(dim=1, keepdim=True).clamp(min=1)
+    avg_attn = weighted_attn / count
+    patch_attention = avg_attn.mean(dim=0)  # [num_patches]
+
+    return patch_attention

BrainAnytime/downstream_dataloader.py

@@ -0,0 +1,737 @@
+import os
+import numpy as np
+import pandas as pd
+import nibabel as nib
+import torch
+from torch.utils.data import Dataset, DataLoader
+import torchio as tio
+from typing import List, Dict, Tuple, Optional, Union
+import random
+from itertools import combinations, compress
+
+
+class MultiModalDownstreamDataset(Dataset):
+    """
+    多模态3D医学图像下游任务数据集
+    
+    特点:
+    - 支持多个标签（AGE, MMSE, CN vs. MCI, CN vs. AD）
+    - AGE使用Z-score归一化（regression_norm），MMSE自动进行Min-Max归一化
+    - 自动过滤缺失所需标签的样本
+    - 支持数据增强（Spatial transforms）
+    - 支持指定特定模态列表加载
+    - 支持intersection模式（所有模态都存在）和union模式（至少一种模态存在）
+    - 支持模态组合数据增强：训练阶段随机drop模态，验证阶段扩展所有模态组合
+    """
+    
+    # 统一的模态顺序
+    MODALITY_ORDER = ['T1', 'T2', 'Flair', 'PET']
+    
+    # 模态简写映射（用于模态组合索引）
+    MODALITY_SHORT = {'T1': 'T', 'T2': 'M', 'Flair': 'F', 'PET': 'P'}
+    
+    # MMSE的全局Min-Max值（基于train/val/test全集计算）
+    # 注意：MMSE过滤了<10的离群值
+    # - MMSE < 10表示重度认知障碍，在ADNI数据集中极少（16个样本，0.54%）
+    # - 这些样本可能是数据质量问题或极端异常值
+    # - 过滤后保留2952个样本（99.46%），MMSE范围[10.0, 30.0]
+    GLOBAL_MIN_MAX = {
+        'MMSE': {'min': 10.0, 'max': 30.0},  # 过滤了MMSE<10的离群值（16个样本）
+    }
+    
+    # DX编码映射: 1=CN, 2=MCI, 3=AD
+    DX_MAPPING = {1: 'CN', 2: 'MCI', 3: 'AD'}
+    
+    def __init__(
+        self,
+        excel_path: str,
+        labels: List[str],
+        image_size: Tuple[int, int, int] = (128, 128, 128),
+        augmentation: bool = True,
+        cache_data: bool = False,
+        base_dir: str = "/home/data/Downstream/ADNI/",
+        modalities: Optional[List[str]] = None,
+        intersection: bool = True,
+        exclusive_modalities: bool = False,
+        phase: str = 'train',
+        modality_dropout: bool = True,
+        expand_val_combinations: bool = True,
+        regression_norm: Optional[Dict[str, Tuple[float, float]]] = None,
+    ):
+        """
+        Args:
+            excel_path: Excel文件路径（train/val/test）
+            labels: 需要加载的标签列表，支持:
+                - 'AGE' 或 'Age': 年龄（回归任务，自动归一化）
+                - 'MMSE': MMSE分数（回归任务，自动归一化）
+                - 'CN vs MCI': 二分类（CN=0, MCI=1）
+                - 'CN vs AD': 二分类（CN=0, AD=1）
+            image_size: 图像尺寸 (D, H, W)
+            augmentation: 是否进行数据增强
+            cache_data: 是否缓存加载的数据到内存
+            base_dir: 图像文件的基础目录，用于拼接相对路径
+            modalities: 要加载的模态列表，如 ['T1', 'T2']。如果为None，则加载所有模态
+            intersection: 模态过滤模式
+                - True: 只加载所有指定模态都存在的样本（交集模式）
+                - False: 只要包含其中一种指定模态就可以（并集模式）
+            exclusive_modalities: 是否只加载仅包含指定模态的样本（排除有其他模态的样本）
+                - True: 只加载样本中存在的模态完全等于指定模态的样本
+                - False: 只要包含指定模态即可（默认行为）
+                例如：如果指定modalities=['T1']，exclusive_modalities=True时，只加载只有T1的样本，排除同时有T1和T2的样本
+            phase: 数据集阶段，'train', 'val', 或 'test'
+            modality_dropout: 是否在训练阶段启用模态dropout增强（仅phase='train'时有效）
+            expand_val_combinations: 是否在验证阶段扩展所有模态组合（仅phase='val'时有效）
+        """
+        self.excel_path = excel_path
+        self.labels = [label.upper() for label in labels]  # 统一转为大写
+        self.image_size = image_size
+        self.augmentation = augmentation
+        self.cache_data = cache_data
+        self.base_dir = base_dir
+        self.cache = {}
+        self.phase = phase
+        self.modality_dropout = modality_dropout
+        self.expand_val_combinations = expand_val_combinations
+        self.regression_norm = regression_norm or {}
+        
+        # 处理模态参数
+        if modalities is None:
+            # 默认加载所有模态
+            self.modalities = self.MODALITY_ORDER.copy()
+        else:
+            # 验证模态名称
+            modalities_upper = [m.upper() for m in modalities]
+            valid_modalities = {m.upper() for m in self.MODALITY_ORDER}
+            for mod in modalities_upper:
+                if mod not in valid_modalities:
+                    raise ValueError(f"Invalid modality: {mod}. Valid modalities are: {self.MODALITY_ORDER}")
+            # 保持模态顺序与MODALITY_ORDER一致
+            self.modalities = [m for m in self.MODALITY_ORDER if m.upper() in modalities_upper]
+        
+        self.intersection = intersection
+        self.exclusive_modalities = exclusive_modalities
+        
+        # 生成模态组合索引映射
+        modality_short_list = [self.MODALITY_SHORT[m] for m in self.MODALITY_ORDER]
+        self.combination_to_index = self._get_modality_combinations(modality_short_list)
+        
+        # 验证标签
+        valid_labels = {'AGE', 'MMSE', 'CN VS MCI', 'CN VS AD'}
+        for label in self.labels:
+            if label not in valid_labels:
+                raise ValueError(f"Invalid label: {label}. Valid labels are: {valid_labels}")
+        
+        # 加载并过滤样本
+        self.samples = self._load_and_filter_samples()
+        
+        # 在验证阶段扩展所有模态组合
+        if self.phase == 'val' and self.expand_val_combinations:
+            self.samples = self._expand_val_combinations()
+            print(f"After expanding validation combinations: {len(self.samples)} samples")
+        
+        print(f"Loaded {len(self.samples)} samples from {excel_path}")
+        print(f"Requested labels: {self.labels}")
+        print(f"Requested modalities: {self.modalities}")
+        print(f"Phase: {self.phase}")
+        print(f"Modality filter mode: {'intersection' if self.intersection else 'union'}")
+        if self.exclusive_modalities:
+            print(f"Exclusive mode: Only loading samples that contain EXACTLY the specified modalities")
+        if self.phase == 'train' and self.modality_dropout:
+            print(f"Modality dropout augmentation: Enabled (will randomly drop modalities during training)")
+        if self.phase == 'val' and self.expand_val_combinations:
+            print(f"Validation combination expansion: Enabled (each sample expanded to all possible modality subsets)")
+        
+        # 初始化数据增强
+        if self.augmentation:
+            self.spatial_transform = tio.OneOf({
+                tio.RandomFlip(axes=0, flip_probability=0.5): 0.33,
+                tio.RandomAffine(scales=(0.9, 1.2), degrees=10, p=0.5): 0.33,
+                tio.RandomElasticDeformation(
+                    num_control_points=(10, 10, 10),
+                    max_displacement=8,
+                    locked_borders=2,
+                    p=0.5
+                ): 0.34,
+            })
+    
+    def _get_modality_combinations(self, modalities: List[str]) -> Dict[str, int]:
+        """
+        生成所有可能的模态组合并创建组合字符串到索引的映射
+        
+        Args:
+            modalities: 模态简写列表，如 ['T', 'M', 'F', 'P']
+        
+        Returns:
+            组合字符串到索引的字典，如 {'T': 0, 'M': 1, ..., 'TMFP': 14}
+        """
+        all_combinations = []
+        for i in range(len(modalities), 0, -1):
+            comb = list(combinations(modalities, i))
+            all_combinations.extend(comb)
+        
+        # 创建映射字典
+        combination_to_index = {''.join(sorted(comb)): idx for idx, comb in enumerate(all_combinations)}
+        return combination_to_index
+    
+    def _observed_to_combination(self, observed: List[int]) -> str:
+        """
+        将observed列表转换为模态组合字符串
+        
+        Args:
+            observed: 观察到的模态列表，如 [1, 1, 1, 0] 表示 T1, T2, Flair 存在，PET 不存在
+        
+        Returns:
+            模态组合字符串，如 "TMF"
+        """
+        modality_short_list = [self.MODALITY_SHORT[m] for m in self.MODALITY_ORDER]
+        available_modalities = list(compress(modality_short_list, observed))
+        return ''.join(sorted(available_modalities))
+    
+    def _expand_val_combinations(self) -> List[Dict]:
+        """
+        在验证阶段，将每个样本扩展为所有可能的非空模态子集
+        
+        Returns:
+            扩展后的样本列表
+        """
+        expanded_samples = []
+        modality_short_list = [self.MODALITY_SHORT[m] for m in self.MODALITY_ORDER]
+        
+        for sample in self.samples:
+            # 获取样本中可用的模态（只考虑在指定模态列表中的）
+            available_modalities = [m for m in sample['modalities'].keys() if m in self.modalities]
+            available_modalities_short = [self.MODALITY_SHORT[m] for m in available_modalities]
+            
+            if len(available_modalities_short) == 0:
+                continue
+            
+            # 生成所有可能的非空子集
+            for r in range(1, len(available_modalities_short) + 1):
+                for subset in combinations(available_modalities_short, r):
+                    # 将简写转换回完整模态名
+                    subset_full = [m for m in self.MODALITY_ORDER if self.MODALITY_SHORT[m] in subset]
+                    
+                    # 创建新的样本，只包含子集中的模态
+                    new_sample = {
+                        'subject_id': sample['subject_id'],
+                        'dataset': sample['dataset'],
+                        'modalities': {mod: sample['modalities'][mod] for mod in subset_full},
+                        'labels': sample['labels'].copy(),
+                        'original_observed': [1 if m in subset_full else 0 for m in self.MODALITY_ORDER],
+                        'diag_group': sample.get('diag_group', None),
+                    }
+                    expanded_samples.append(new_sample)
+        
+        return expanded_samples
+    
+    def _load_and_filter_samples(self) -> List[Dict]:
+        """加载Excel文件并过滤出包含所有所需标签的样本"""
+        if not os.path.exists(self.excel_path):
+            raise FileNotFoundError(f"Excel file not found: {self.excel_path}")
+        
+        df = pd.read_excel(self.excel_path)
+        samples = []
+        
+        # 统计信息
+        total_rows = len(df)
+        missing_labels_count = 0
+        missing_modalities_count = 0
+        union_passed_count = 0
+        intersection_passed_count = 0
+        exclusive_filtered_count = 0
+        all_4_modalities_count = 0  # 统计同时包含所有4个指定模态的样本数
+        modality_stats = {mod: 0 for mod in self.modalities}  # 统计每个指定模态的出现次数
+        
+        # 模态列名映射
+        modality_columns = {'T1': 'T1', 'T2': 'T2', 'Flair': 'Flair', 'PET': 'PET'}
+        
+        for idx, row in df.iterrows():
+            sample = {
+                'subject_id': row.get('SubjectID', f'sample_{idx}'),
+                'dataset': row.get('Dataset', 'Unknown'),
+                'modalities': {},
+                'labels': {},
+                'diag_group': None,
+            }
+
+            # Diagnosis group for CN/MCI/AD (used for AGE CN-only train/val and test stratified metrics)
+            if 'DX' in df.columns:
+                dx = row.get('DX', None)
+                if pd.notna(dx):
+                    try:
+                        sample['diag_group'] = self.DX_MAPPING.get(int(dx), None)
+                    except (TypeError, ValueError):
+                        pass
+            
+            # 加载模态路径
+            for unified_name, col_name in modality_columns.items():
+                if col_name in df.columns:
+                    path = row[col_name]
+                    if pd.notna(path) and isinstance(path, str):
+                        # 如果是相对路径，则与base_dir拼接
+                        if not os.path.isabs(path):
+                            full_path = os.path.join(self.base_dir, path)
+                        else:
+                            full_path = path
+                        
+                        if os.path.exists(full_path):
+                            sample['modalities'][unified_name] = full_path
+            
+            # 检查并加载标签
+            has_all_labels = True
+            for label in self.labels:
+                label_value = None
+                
+                if label == 'AGE':
+                    if 'Age' in df.columns:
+                        age = row['Age']
+                        if pd.notna(age):
+                            try:
+                                age_f = float(age)
+                                if np.isfinite(age_f) and 'AGE' in self.regression_norm:
+                                    mean, std = self.regression_norm['AGE']
+                                    label_value = (age_f - float(mean)) / float(std)
+                            except (TypeError, ValueError):
+                                pass
+                
+                elif label == 'MMSE':
+                    if 'MMSE' in df.columns:
+                        mmse = row['MMSE']
+                        if pd.notna(mmse) and mmse >= self.GLOBAL_MIN_MAX['MMSE']['min']:
+                            # 归一化到[0, 1]
+                            # 注意：过滤了MMSE < 10的离群值（重度认知障碍，可能是数据质量问题）
+                            min_val = self.GLOBAL_MIN_MAX['MMSE']['min']
+                            max_val = self.GLOBAL_MIN_MAX['MMSE']['max']
+                            label_value = (mmse - min_val) / (max_val - min_val)
+                
+                elif label == 'CN VS MCI':
+                    if 'DX' in df.columns:
+                        dx = row['DX']
+                        if pd.notna(dx):
+                            dx_int = int(dx)
+                            if dx_int == 1:  # CN
+                                label_value = 0.0
+                            elif dx_int == 2:  # MCI
+                                label_value = 1.0
+                            # DX=3 (AD) 不包含在此任务中，设为None
+                
+                elif label == 'CN VS AD':
+                    if 'DX' in df.columns:
+                        dx = row['DX']
+                        if pd.notna(dx):
+                            dx_int = int(dx)
+                            if dx_int == 1:  # CN
+                                label_value = 0.0
+                            elif dx_int == 3:  # AD
+                                label_value = 1.0
+                            # DX=2 (MCI) 不包含在此任务中，设为None
+                
+                if label_value is None:
+                    has_all_labels = False
+                    break
+                else:
+                    sample['labels'][label] = label_value
+            
+            # 根据intersection参数过滤模态
+            if has_all_labels:
+                # 检查样本中存在的指定模态
+                available_modalities = [mod for mod in self.modalities if mod in sample['modalities']]
+                
+                # 更新模态统计
+                for mod in available_modalities:
+                    modality_stats[mod] += 1
+                
+                # 统计同时包含所有4个指定模态的样本数
+                if len(available_modalities) == len(self.modalities):
+                    all_4_modalities_count += 1
+                
+                # 检查是否通过intersection/union过滤
+                passed_modality_filter = False
+                if self.intersection:
+                    # 交集模式：所有指定模态都必须存在
+                    if len(available_modalities) == len(self.modalities):
+                        passed_modality_filter = True
+                        intersection_passed_count += 1
+                    else:
+                        missing_modalities_count += 1
+                else:
+                    # 并集模式：至少包含一种指定模态
+                    if len(available_modalities) >= 1:
+                        passed_modality_filter = True
+                        union_passed_count += 1
+                    else:
+                        missing_modalities_count += 1
+                
+                # 如果通过了intersection/union过滤，再检查exclusive_modalities
+                if passed_modality_filter:
+                    if self.exclusive_modalities:
+                        # 检查样本中存在的模态集合是否完全等于指定的模态集合
+                        sample_modalities_set = set(sample['modalities'].keys())
+                        specified_modalities_set = set(self.modalities)
+                        if sample_modalities_set == specified_modalities_set:
+                            samples.append(sample)
+                        else:
+                            exclusive_filtered_count += 1
+                    else:
+                        # 非exclusive模式，直接添加
+                        samples.append(sample)
+            else:
+                missing_labels_count += 1
+        
+        # 输出详细的统计信息
+        print(f"\n数据加载统计信息:")
+        print(f"  Excel总行数: {total_rows}")
+        print(f"  缺失标签的样本数: {missing_labels_count}")
+        print(f"  缺失模态的样本数: {missing_modalities_count}")
+        print(f"  各指定模态在样本中的出现次数:")
+        for mod, count in modality_stats.items():
+            print(f"    {mod}: {count} 次")
+        print(f"  同时包含所有{len(self.modalities)}个指定模态的样本数: {all_4_modalities_count}")
+        if self.intersection:
+            print(f"  交集模式: 需要所有指定模态({self.modalities})都存在 (通过: {intersection_passed_count})")
+        else:
+            print(f"  并集模式: 至少一种指定模态存在即可 (通过: {union_passed_count})")
+        if self.exclusive_modalities:
+            print(f"  独占模式: 只加载样本中存在的模态完全等于指定模态的样本")
+            print(f"    因独占模式被过滤的样本数: {exclusive_filtered_count}")
+            print(f"    说明: 只有同时包含所有{len(self.modalities)}个指定模态的样本才会被加载")
+        print(f"  最终加载的样本数: {len(samples)}")
+        
+        if len(samples) == 0:
+            print(f"\n警告: 没有加载到任何样本!")
+            print(f"  可能的原因:")
+            print(f"  1. Excel文件中没有同时包含所有请求标签的样本")
+            print(f"  2. 请求的模态在样本中不存在或文件路径不正确")
+            if self.intersection:
+                print(f"  3. 交集模式要求所有指定模态({self.modalities})都必须存在")
+            else:
+                print(f"  3. 并集模式要求至少一种指定模态({self.modalities})存在")
+            print(f"  建议: 检查Excel文件内容或尝试使用 --intersection False")
+        
+        return samples
+    
+    def _load_nifti(self, path: str) -> np.ndarray:
+        """加载NIfTI文件"""
+        try:
+            nii = nib.load(path)
+            data = nii.get_fdata().astype(np.float32)
+            return data
+        except Exception as e:
+            print(f"Error loading {path}: {e}")
+            return None
+    
+    def __len__(self) -> int:
+        return len(self.samples)
+    
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        sample = self.samples[idx]
+        
+        # 获取原始observed状态（样本中实际存在的模态）
+        if 'original_observed' in sample:
+            # 验证阶段扩展的样本已经有original_observed
+            original_observed = sample['original_observed']
+        else:
+            # 训练/测试阶段，从样本的modalities构建original_observed
+            original_observed = [1 if modality in sample['modalities'] else 0 for modality in self.MODALITY_ORDER]
+        
+        # 应用模态dropout（仅在训练阶段且启用时）
+        observed = original_observed.copy()
+        if self.phase == 'train' and self.modality_dropout:
+            # 获取可用的模态（在指定模态列表中的）
+            available_modalities = []
+            for mod_idx, modality in enumerate(self.MODALITY_ORDER):
+                if observed[mod_idx] == 1 and modality in self.modalities:
+                    available_modalities.append((mod_idx, modality))
+            
+            # 如果有多个可用模态，随机drop一些
+            if len(available_modalities) > 1:
+                m = len(available_modalities)
+                k = random.randint(1, m - 1)  # 随机选择要drop的模态数量 (1 <= k < m)
+                modalities_to_drop = random.sample(available_modalities, k)
+                
+                # 更新observed列表
+                for mod_idx, _ in modalities_to_drop:
+                    observed[mod_idx] = 0
+        
+        # 检查缓存（使用原始observed作为缓存键的一部分，因为dropout是随机的）
+        cache_key = (idx, tuple(original_observed))
+        if self.cache_data and cache_key in self.cache:
+            cached_data = self.cache[cache_key]
+            images = cached_data['images'].clone()
+            cached_observed = cached_data['observed'].clone()
+        else:
+            # 初始化输出张量（始终为4个模态，对应完整的MODALITY_ORDER）
+            num_full_modalities = len(self.MODALITY_ORDER)  # Always 4
+            images = torch.zeros(num_full_modalities, *self.image_size, dtype=torch.float32)
+            cached_observed = torch.zeros(num_full_modalities, dtype=torch.float32)
+            
+            # 加载每个模态（按照完整的MODALITY_ORDER顺序）
+            # 只加载在original_observed中存在的模态（不考虑dropout）
+            for mod_idx, modality in enumerate(self.MODALITY_ORDER):
+                # 只加载在指定模态列表中的模态
+                if modality in self.modalities:
+                    if modality in sample['modalities']:
+                        path = sample['modalities'][modality]
+                        data = self._load_nifti(path)
+                        
+                        if data is not None:
+                            # 确保数据尺寸正确
+                            if data.shape == self.image_size:
+                                images[mod_idx] = torch.from_numpy(data)
+                                cached_observed[mod_idx] = 1.0
+                            else:
+                                print(f"Warning: Size mismatch for {path}, expected {self.image_size}, got {data.shape}")
+                # 如果模态不在指定的模态列表中，cached_observed[mod_idx]保持为0，images[mod_idx]保持为全零
+            
+            # 缓存数据（使用原始observed）
+            if self.cache_data:
+                self.cache[cache_key] = {
+                    'images': images.clone(),
+                    'observed': cached_observed.clone()
+                }
+        
+        # 应用dropout后的observed（将dropout的模态设为0）
+        observed_tensor = torch.tensor(observed, dtype=torch.float32)
+        # 对于被dropout的模态，将图像也设为0
+        images = images.clone()
+        for mod_idx in range(len(observed)):
+            if observed[mod_idx] == 0:
+                images[mod_idx] = torch.zeros_like(images[mod_idx])
+        
+        # 应用空间数据增强（只对observed的模态应用）
+        if self.augmentation:
+            # 只对observed的模态应用增强（按照MODALITY_ORDER顺序）
+            subject_dict = {}
+            for mod_idx, modality in enumerate(self.MODALITY_ORDER):
+                if observed[mod_idx] == 1.0:
+                    # TorchIO需要4D张量 (C, D, H, W)
+                    subject_dict[modality] = tio.ScalarImage(tensor=images[mod_idx:mod_idx+1])
+            
+            if subject_dict:
+                subject = tio.Subject(**subject_dict)
+                transformed = self.spatial_transform(subject)
+                
+                # 将增强后的数据放回images张量
+                for mod_idx, modality in enumerate(self.MODALITY_ORDER):
+                    if modality in subject_dict:
+                        images[mod_idx] = transformed[modality].data[0]
+        
+        # 计算模态组合索引
+        combination_str = self._observed_to_combination(observed)
+        mc = self.combination_to_index.get(combination_str, 0)
+        
+        # 构建标签张量
+        labels_tensor = torch.tensor([sample['labels'][label] for label in self.labels], dtype=torch.float32)
+        
+        return {
+            'images': images,  # (4, D, H, W) - Always 4 modalities in MODALITY_ORDER
+            'observed': observed_tensor,  # (4,) - After modality dropout (if applied)
+            'original_observed': torch.tensor(original_observed, dtype=torch.float32),  # (4,) - Original observed before dropout
+            'labels': labels_tensor,  # (num_labels,)
+            'mc': torch.tensor(mc, dtype=torch.long),  # Modality combination index
+            'subject_id': sample['subject_id'],
+            'diag_group': sample.get('diag_group', None),
+        }
+
+
+def create_downstream_dataloader(
+    excel_path: str,
+    labels: List[str],
+    batch_size: int = 4,
+    num_workers: int = 8,
+    augmentation: bool = True,
+    shuffle: bool = True,
+    pin_memory: bool = True,
+    cache_data: bool = False,
+    image_size: Tuple[int, int, int] = (128, 128, 128),
+    base_dir: str = "/home/data/Downstream/ADNI/",
+    modalities: Optional[List[str]] = None,
+    intersection: bool = True,
+    exclusive_modalities: bool = False,
+    phase: str = 'train',
+    modality_dropout: bool = True,
+    expand_val_combinations: bool = True,
+    regression_norm: Optional[Dict[str, Tuple[float, float]]] = None,
+) -> DataLoader:
+    """
+    创建下游任务数据加载器
+    
+    Args:
+        excel_path: Excel文件路径（train/val/test）
+        labels: 需要加载的标签列表，支持:
+            - 'AGE' 或 'Age': 年龄（回归任务，自动归一化）
+            - 'MMSE': MMSE分数（回归任务，自动归一化）
+            - 'CN vs MCI': 二分类（CN=0, MCI=1）
+            - 'CN vs AD': 二分类（CN=0, AD=1）
+        batch_size: 批量大小
+        num_workers: 数据加载进程数
+        augmentation: 是否数据增强
+        shuffle: 是否打乱数据
+        pin_memory: 是否使用pinned memory
+        cache_data: 是否缓存数据到内存
+        image_size: 图像尺寸 (D, H, W)
+        base_dir: 图像文件的基础目录，用于拼接相对路径
+        modalities: 要加载的模态列表，如 ['T1', 'T2']。如果为None，则加载所有模态
+        intersection: 模态过滤模式
+            - True: 只加载所有指定模态都存在的样本（交集模式）
+            - False: 只要包含其中一种指定模态就可以（并集模式）
+        exclusive_modalities: 是否只加载仅包含指定模态的样本（排除有其他模态的样本）
+            - True: 只加载样本中存在的模态完全等于指定模态的样本
+            - False: 只要包含指定模态即可（默认行为）
+            例如：如果指定modalities=['T1']，exclusive_modalities=True时，只加载只有T1的样本，排除同时有T1和T2的样本
+        phase: 数据集阶段，'train', 'val', 或 'test'
+        modality_dropout: 是否在训练阶段启用模态dropout增强（仅phase='train'时有效）
+        expand_val_combinations: 是否在验证阶段扩展所有模态组合（仅phase='val'时有效）
+    
+    Returns:
+        DataLoader实例
+    """
+    # Label-efficiency compatibility: if excel_path is relative and not found from cwd,
+    # try resolving it under base_dir. This keeps old behavior for absolute paths.
+    resolved_excel_path = excel_path
+    if not os.path.isabs(resolved_excel_path) and not os.path.exists(resolved_excel_path):
+        candidate = os.path.join(base_dir, resolved_excel_path)
+        if os.path.exists(candidate):
+            resolved_excel_path = candidate
+
+    dataset = MultiModalDownstreamDataset(
+        excel_path=resolved_excel_path,
+        labels=labels,
+        image_size=image_size,
+        augmentation=augmentation,
+        cache_data=cache_data,
+        base_dir=base_dir,
+        modalities=modalities,
+        intersection=intersection,
+        exclusive_modalities=exclusive_modalities,
+        phase=phase,
+        modality_dropout=modality_dropout,
+        expand_val_combinations=expand_val_combinations,
+        regression_norm=regression_norm,
+    )
+    
+    dataloader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        num_workers=num_workers,
+        pin_memory=pin_memory,
+        drop_last=False,  # 下游任务通常不drop last
+    )
+    
+    return dataloader
+
+
+# ============== 使用示例 ==============
+if __name__ == '__main__':
+    print("=" * 60)
+    print("多模态3D医学图像下游任务数据加载器")
+    print("=" * 60)
+    
+    # 示例1: 训练阶段，启用模态dropout增强
+    print("\n示例1: 训练阶段，启用模态dropout增强")
+    dataloader = create_downstream_dataloader(
+        excel_path="/home/data/Downstream/ADNI_Division/modality_data_train.xlsx",
+        labels= ["CN vs AD"],
+        batch_size=2,
+        num_workers=4,
+        augmentation=True,
+        modalities=["T1","T2","Flair","PET"],
+        intersection=False,
+        shuffle=True,
+        phase='train',
+        modality_dropout=True,  # 启用模态dropout
+        expand_val_combinations=False,
+    )
+    
+    print(f"\n数据集大小: {len(dataloader.dataset)}")
+    print(f"批量数: {len(dataloader)}")
+    print(f"请求的标签: {dataloader.dataset.labels}")
+    
+    # 测试加载一个批量
+    print("\n测试加载一个批量...")
+    for batch in dataloader:
+        images = batch['images']
+        observed = batch['observed']
+        original_observed = batch['original_observed']
+        labels = batch['labels']
+        mc = batch['mc']
+        
+        print(f"\n批量数据形状:")
+        print(f"  images: {images.shape}")  # (B, 4, 128, 128, 128)
+        print(f"  observed: {observed.shape}")  # (B, 4) - After dropout
+        print(f"  original_observed: {original_observed.shape}")  # (B, 4) - Before dropout
+        print(f"  labels: {labels.shape}")  # (B, num_labels)
+        print(f"  mc (modality combination): {mc.shape}")  # (B,)
+        print(f"  subject_ids: {batch['subject_id']}")
+        print(f"  示例: original_observed[0]={original_observed[0].numpy()}, observed[0]={observed[0].numpy()}, mc[0]={mc[0].item()}")
+        break
+    
+    # 示例2: 验证阶段，扩展所有模态组合
+    print("\n" + "=" * 60)
+    print("示例2: 验证阶段，扩展所有模态组合")
+    print("=" * 60)
+
+    dataloader2 = create_downstream_dataloader(
+        excel_path="/home/data/Downstream/ADNI_Division/modality_data_val.xlsx",
+        labels=["CN vs AD"],
+        batch_size=2,
+        modalities=["T1","T2","Flair","PET"],
+        intersection=False,
+        augmentation=False,
+        shuffle=False,
+        phase='val',
+        modality_dropout=False,  # 验证阶段不启用dropout
+        expand_val_combinations=True,  # 启用模态组合扩展
+    )
+    
+    print(f"\n数据集大小: {len(dataloader2.dataset)}")
+    print(f"请求的标签: {dataloader2.dataset.labels}")
+    print(f"注意: 验证集已扩展为所有可能的模态组合，样本数会增加")
+    
+    # 测试加载一个批量
+    print("\n测试加载一个批量...")
+    for batch in dataloader2:
+        images = batch['images']
+        observed = batch['observed']
+        original_observed = batch['original_observed']
+        labels = batch['labels']
+        mc = batch['mc']
+        
+        print(f"\n批量数据形状:")
+        print(f"  images: {images.shape}")
+        print(f"  observed: {observed.shape}")
+        print(f"  original_observed: {original_observed.shape}")
+        print(f"  labels: {labels.shape}")
+        print(f"  mc (modality combination): {mc.shape}")
+        print(f"  示例: original_observed[0]={original_observed[0].numpy()}, observed[0]={observed[0].numpy()}, mc[0]={mc[0].item()}")
+        break
+    
+    # 示例3: 测试阶段，不使用任何增强
+    print("\n" + "=" * 60)
+    print("示例3: 测试阶段，不使用任何增强")
+    print("=" * 60)
+
+    dataloader3 = create_downstream_dataloader(
+        excel_path="/home/data/Downstream/ADNI_Division/modality_data_test.xlsx",
+        labels=["CN vs AD"],
+        batch_size=2,
+        modalities=["T1","T2","Flair","PET"],
+        intersection=False,
+        augmentation=False,
+        shuffle=False,
+        phase='test',
+        modality_dropout=False,
+        expand_val_combinations=False,
+        exclusive_modalities=False,
+    )
+    
+    print(f"\n数据集大小: {len(dataloader3.dataset)}")
+    print(f"请求的标签: {dataloader3.dataset.labels}")
+    
+    print("\n" + "=" * 60)
+    print("数据加载测试完成!")
+    print("=" * 60)
+
+
+

BrainAnytime/finetune_main.py

@@ -0,0 +1,854 @@
+#!/usr/bin/env python
+"""
+MultiMAE3D Finetuning for Downstream Tasks
+
+Full finetuning: train the entire model end-to-end.
+
+Tasks: CN vs AD, CN vs MCI, MMSE, AGE
+Each task runs with multiple seeds and reports mean +/- std metrics.
+
+Usage:
+    # Finetune on all 4 tasks (3 seeds each)
+    python finetune_main.py --pretrained ./pretrain_checkpoints/multimae/best_model.pth
+
+    # Specific task only
+    python finetune_main.py --pretrained ./pretrain_checkpoints/multimae/best_model.pth --tasks "CN vs AD"
+"""
+
+import os
+import sys
+import gc
+import random
+import warnings
+from copy import deepcopy
+from collections import defaultdict
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+from tqdm import tqdm, trange
+from scipy.stats import pearsonr
+
+warnings.filterwarnings("ignore")
+
+_BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, _BASE_DIR)
+
+from models.multimae3d import create_multimae3d, MultiMAE3D
+from downstream_dataloader import create_downstream_dataloader
+
+
+# =========================================================================
+# Utilities
+# =========================================================================
+
+def seed_everything(seed: int):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+def str2bool(v):
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ('true', '1', 'yes'):
+        return True
+    if v.lower() in ('false', '0', 'no'):
+        return False
+    raise ValueError(f'Boolean value expected, got: {v}')
+
+
+def setup_logger(log_dir, name, filename):
+    """Simple logger setup."""
+    import logging
+    os.makedirs(log_dir, exist_ok=True)
+    logger = logging.getLogger(name)
+    logger.setLevel(logging.INFO)
+    fh = logging.FileHandler(os.path.join(log_dir, filename))
+    fh.setLevel(logging.INFO)
+    formatter = logging.Formatter('%(asctime)s - %(message)s')
+    fh.setFormatter(formatter)
+    logger.addHandler(fh)
+    return logger
+
+
+# =========================================================================
+# Model: MultiMAE3D encoder + downstream task head
+# =========================================================================
+
+class MultiMAE3DForDownstream(nn.Module):
+    """
+    MultiMAE3D encoder + task head for downstream classification/regression.
+
+    Uses encoder.encode() to get all-patch features, pools to a single vector
+    (CLS token or mean pooling), then applies a linear head.
+    """
+
+    def __init__(
+        self,
+        encoder: MultiMAE3D,
+        embed_dim: int = 768,
+        num_outputs: int = 1,
+        pool: str = 'cls',
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+        self.encoder = encoder
+        self.pool = pool
+        self.num_patches_per_modality = encoder.num_patches
+        self.num_global_tokens = encoder.num_global_tokens
+
+        self.norm = nn.LayerNorm(embed_dim)
+        self.head = nn.Sequential(
+            nn.Dropout(dropout),
+            nn.Linear(embed_dim, num_outputs),
+        )
+
+    def forward(self, images: torch.Tensor, observed: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            images:   [B, 4, D, H, W]
+            observed: [B, 4] float mask (1.0=present, 0.0=missing)
+        Returns:
+            logits: [B, num_outputs]
+        """
+        # encode() returns [B, 1 + 4*num_patches, embed_dim]
+        encoder_out = self.encoder.encode(images, observed)
+
+        if self.pool == 'cls':
+            features = encoder_out[:, 0]  # CLS token -> [B, D]
+        elif self.pool == 'mean':
+            # Mean pool over modality tokens with masking for missing modalities
+            tokens = encoder_out[:, self.num_global_tokens:]  # [B, 4*N_p, D]
+            B, _, D = tokens.shape
+            N = self.num_patches_per_modality
+            # Build per-token mask: repeat each modality's observed flag N times
+            mask = observed.unsqueeze(-1).expand(-1, -1, N)  # [B, 4, N]
+            mask = mask.reshape(B, 4 * N).unsqueeze(-1)  # [B, 4*N, 1]
+            features = (tokens * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1.0)
+        else:
+            raise ValueError(f"Unknown pool type: {self.pool}")
+
+        features = self.norm(features)
+        logits = self.head(features)
+        return logits
+
+
+# =========================================================================
+# Pretrained weight loading
+# =========================================================================
+
+def load_pretrained_weights(model: MultiMAE3D, checkpoint_path: str, device='cpu'):
+    """
+    Load pretrained encoder weights into a MultiMAE3D model.
+
+    Supports checkpoint formats:
+      - 'encoder_state_dict': encoder-only (from periodic/best saves)
+      - 'full_model_state_dict': full model (from best_model.pth)
+      - 'model_state_dict': full model (from latest.pth)
+      - raw state_dict (no wrapper key)
+    """
+    print(f"Loading pretrained weights from: {checkpoint_path}")
+    ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
+
+    # Pick the best available state dict
+    if 'encoder_state_dict' in ckpt:
+        state_dict = ckpt['encoder_state_dict']
+        source = 'encoder_state_dict'
+    elif 'full_model_state_dict' in ckpt:
+        state_dict = ckpt['full_model_state_dict']
+        source = 'full_model_state_dict'
+    elif 'model_state_dict' in ckpt:
+        state_dict = ckpt['model_state_dict']
+        source = 'model_state_dict'
+    else:
+        state_dict = ckpt
+        source = 'raw'
+
+    # If loading from full model, filter to encoder keys only
+    encoder_prefixes = ('encoder.', 'input_adapters.', 'pos_embed', 'global_tokens')
+    if source in ('full_model_state_dict', 'model_state_dict', 'raw'):
+        state_dict = {
+            k: v for k, v in state_dict.items()
+            if any(k.startswith(p) for p in encoder_prefixes)
+        }
+
+    missing, unexpected = model.load_state_dict(state_dict, strict=False)
+
+    # Output adapters (decoders) are expected to be missing — we don't use them
+    truly_missing = [k for k in missing if not k.startswith('output_adapters.')]
+
+    epoch_info = ckpt.get('epoch', '?')
+    loss_info = ckpt.get('loss', '?')
+    if isinstance(loss_info, float):
+        loss_info = f"{loss_info:.4f}"
+    print(f"  Source: {source}, Epoch: {epoch_info}, Pretrain Loss: {loss_info}")
+    print(f"  Loaded {len(state_dict)} parameter tensors")
+    if truly_missing:
+        print(f"  WARNING: {len(truly_missing)} encoder keys missing: {truly_missing[:5]}...")
+    if unexpected:
+        print(f"  WARNING: {len(unexpected)} unexpected keys (ignored)")
+
+    del ckpt
+    return model
+
+
+# =========================================================================
+# Metrics
+# =========================================================================
+
+def calc_regression_metrics(preds, labels):
+    """Calculate MAE, RMSE, Pearson correlation."""
+    preds, labels = np.array(preds), np.array(labels)
+    mae = np.mean(np.abs(preds - labels))
+    rmse = np.sqrt(np.mean((preds - labels) ** 2))
+    if len(preds) > 1 and np.std(preds) > 0 and np.std(labels) > 0:
+        r, _ = pearsonr(preds, labels)
+    else:
+        r = 0.0
+    return {'mae': mae, 'rmse': rmse, 'pearson': r}
+
+
+def calc_classification_metrics(preds, labels, probs):
+    """Calculate ACC, AUC, Sensitivity, Specificity, F1."""
+    from sklearn.metrics import accuracy_score, f1_score, roc_auc_score, confusion_matrix
+
+    preds, labels = np.array(preds), np.array(labels)
+    probs = np.array(probs)
+
+    acc = accuracy_score(labels, preds)
+
+    try:
+        auc = roc_auc_score(labels, probs[:, 1])
+    except ValueError:
+        auc = 0.0
+
+    f1 = f1_score(labels, preds, average='binary')
+
+    cm = confusion_matrix(labels, preds)
+    if cm.shape == (2, 2):
+        tn, fp, fn, tp = cm.ravel()
+        sensitivity = tp / (tp + fn) if (tp + fn) > 0 else 0.0
+        specificity = tn / (tn + fp) if (tn + fp) > 0 else 0.0
+    else:
+        sensitivity, specificity = 0.0, 0.0
+
+    return {
+        'acc': acc, 'auc': auc,
+        'sensitivity': sensitivity, 'specificity': specificity,
+        'f1': f1,
+    }
+
+
+def calc_metrics_by_combo(preds, labels, probs, combos, task_type):
+    """Calculate metrics grouped by modality combination string."""
+    is_cls = task_type in ('CN vs AD', 'CN vs MCI')
+    grouped = defaultdict(lambda: {'preds': [], 'labels': [], 'probs': []})
+    for i, combo in enumerate(combos):
+        grouped[combo]['preds'].append(preds[i])
+        grouped[combo]['labels'].append(labels[i])
+        grouped[combo]['probs'].append(probs[i])
+
+    results = {}
+    for combo, data in grouped.items():
+        p = np.array(data['preds'])
+        l = np.array(data['labels'])
+        pr = np.array(data['probs'])
+        if is_cls:
+            try:
+                m = calc_classification_metrics(p, l, pr)
+            except Exception:
+                m = {'acc': 0, 'auc': 0, 'sensitivity': 0, 'specificity': 0, 'f1': 0}
+        else:
+            m = calc_regression_metrics(p, l)
+        m['n_samples'] = len(p)
+        results[combo] = m
+    return results
+
+
+# =========================================================================
+# Training & Evaluation
+# =========================================================================
+
+MODALITY_NAMES = ['T1', 'T2', 'Flair', 'PET']
+
+
+def run_epoch(loader, model, criterion, device, task_type,
+              is_training=False, optimizer=None):
+    """Run one epoch of training or evaluation."""
+    all_preds, all_labels, all_probs = [], [], []
+    modality_combos = []
+    total_loss, n_batches = 0.0, 0
+    is_cls = task_type in ('CN vs AD', 'CN vs MCI')
+
+    model.train() if is_training else model.eval()
+
+    with torch.set_grad_enabled(is_training):
+        for batch in tqdm(loader, leave=False,
+                          desc='Train' if is_training else 'Eval'):
+            images = batch['images'].to(device, non_blocking=True)
+            observed = batch['observed'].to(device, non_blocking=True)
+            labels = batch['labels'][:, 0].to(device, non_blocking=True)
+
+            logits = model(images, observed)       # [B, 1]
+            logits_flat = logits.squeeze(-1)       # [B]
+            loss = criterion(logits_flat, labels.float())
+
+            total_loss += loss.item()
+            n_batches += 1
+
+            if is_training:
+                optimizer.zero_grad()
+                loss.backward()
+                nn.utils.clip_grad_norm_(
+                    [p for p in model.parameters() if p.requires_grad],
+                    max_norm=1.0,
+                )
+                optimizer.step()
+
+            # Collect predictions
+            if is_cls:
+                prob_pos = torch.sigmoid(logits_flat).detach().cpu().numpy()
+                pred = (prob_pos > 0.5).astype(int)
+                probs_2d = np.stack([1 - prob_pos, prob_pos], axis=1)
+                all_preds.extend(pred)
+                all_labels.extend(labels.cpu().numpy())
+                all_probs.extend(probs_2d)
+            else:
+                pred_vals = logits_flat.detach().cpu().numpy()
+                all_preds.extend(pred_vals)
+                all_labels.extend(labels.cpu().numpy())
+                all_probs.extend(pred_vals.reshape(-1, 1))
+
+            # Track modality combos (eval only)
+            if not is_training:
+                B = images.shape[0]
+                for i in range(B):
+                    present = [
+                        MODALITY_NAMES[j]
+                        for j in range(4)
+                        if observed[i, j] > 0.5
+                    ]
+                    combo_str = ('+'.join(sorted(present))
+                                 if present else 'None')
+                    modality_combos.append(combo_str)
+
+    avg_loss = total_loss / max(n_batches, 1)
+    return avg_loss, all_preds, all_labels, all_probs, modality_combos
+
+
+# =========================================================================
+# Single task + seed pipeline
+# =========================================================================
+
+def train_and_evaluate(args, task_type, seed, device):
+    """
+    Full train/val/test pipeline for one (task, seed) combination.
+    Returns a dict of test metrics.
+    """
+    seed_everything(seed)
+    torch.cuda.empty_cache()
+
+    is_cls = task_type in ('CN vs AD', 'CN vs MCI')
+
+    # ---- Data loaders ----
+    loader_kwargs = dict(
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        pin_memory=True,
+        cache_data=False,
+        image_size=tuple(args.image_size),
+        base_dir=args.base_dir,
+        modalities=args.modalities,
+        intersection=args.intersection,
+    )
+
+    print(f"\nLoading data for task={task_type}, seed={seed}, mode=finetune")
+    train_loader = create_downstream_dataloader(
+        excel_path=args.train_excel, labels=[task_type],
+        augmentation=True, shuffle=True,
+        phase='train', modality_dropout=True, expand_val_combinations=False,
+        **loader_kwargs,
+    )
+    val_loader = create_downstream_dataloader(
+        excel_path=args.val_excel, labels=[task_type],
+        augmentation=False, shuffle=False,
+        phase='val', modality_dropout=False, expand_val_combinations=True,
+        **loader_kwargs,
+    )
+    test_loader = create_downstream_dataloader(
+        excel_path=args.test_excel, labels=[task_type],
+        augmentation=False, shuffle=False,
+        phase='test', modality_dropout=False, expand_val_combinations=False,
+        exclusive_modalities=False,
+        **loader_kwargs,
+    )
+    print(f"  Train: {len(train_loader.dataset)}, "
+          f"Val: {len(val_loader.dataset)}, "
+          f"Test: {len(test_loader.dataset)}")
+
+    # ---- Model ----
+    encoder = create_multimae3d(
+        img_size=args.img_size,
+        patch_size=args.patch_size,
+        embed_dim=args.embed_dim,
+        depth=args.depth,
+        num_heads=args.num_heads,
+        decoder_embed_dim=args.decoder_embed_dim,
+        decoder_depth=args.decoder_depth,
+        decoder_num_heads=args.decoder_num_heads,
+    )
+
+    # Load pretrained weights
+    if args.pretrained and os.path.isfile(args.pretrained):
+        load_pretrained_weights(encoder, args.pretrained, device='cpu')
+    else:
+        print("  No pretrained weights loaded (training from scratch)")
+
+    model = MultiMAE3DForDownstream(
+        encoder=encoder,
+        embed_dim=args.embed_dim,
+        num_outputs=1,
+        pool=args.pool,
+        dropout=args.dropout,
+    ).to(device)
+
+    total_params = sum(p.numel() for p in model.parameters())
+
+    # ---- Freeze encoder if requested ----
+    freeze_epochs = getattr(args, 'freeze_epochs', 0)
+    if freeze_epochs > 0:
+        # Freeze all pretrained encoder parameters
+        for param in model.encoder.parameters():
+            param.requires_grad = False
+        trainable_params = sum(p.numel() for p in model.parameters()
+                               if p.requires_grad)
+        encoder_params = sum(p.numel() for p in model.encoder.parameters())
+        print(f"  Model: {total_params:,} total, {trainable_params:,} trainable "
+              f"(encoder frozen: {encoder_params:,} params for first {freeze_epochs} epochs)")
+    else:
+        trainable_params = sum(p.numel() for p in model.parameters()
+                               if p.requires_grad)
+        print(f"  Model: {total_params:,} total, {trainable_params:,} trainable")
+
+    # ---- Helper: build optimizer + scheduler ----
+    warmup_ep = args.warmup_epochs
+    total_ep = args.epochs
+
+    def build_optimizer_and_scheduler(model, lr, remaining_epochs, warmup):
+        trainable = [p for p in model.parameters() if p.requires_grad]
+        optimizer = torch.optim.AdamW(trainable, lr=lr,
+                                      weight_decay=args.weight_decay)
+
+        def lr_lambda(epoch):
+            if epoch < warmup:
+                return (epoch + 1) / max(warmup, 1)
+            progress = (epoch - warmup) / max(remaining_epochs - warmup, 1)
+            return 0.5 * (1.0 + np.cos(np.pi * progress))
+
+        scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+        return optimizer, scheduler
+
+    optimizer, scheduler = build_optimizer_and_scheduler(
+        model, args.lr, total_ep, warmup_ep)
+
+    # Criterion
+    criterion = (nn.BCEWithLogitsLoss() if is_cls
+                 else nn.MSELoss()).to(device)
+
+    # ---- Training loop ----
+    best_metric = 0.0 if is_cls else float('inf')
+    best_model_state = None
+    patience_counter = 0
+
+    for epoch in range(total_ep):
+        # Unfreeze encoder after freeze_epochs
+        if freeze_epochs > 0 and epoch == freeze_epochs:
+            print(f"\n  >>> Epoch {epoch}: Unfreezing pretrained encoder <<<")
+            for param in model.encoder.parameters():
+                param.requires_grad = True
+            unfrozen_trainable = sum(p.numel() for p in model.parameters()
+                                     if p.requires_grad)
+            print(f"  Trainable params: {unfrozen_trainable:,} (all parameters)")
+            # Rebuild optimizer & scheduler for joint training phase
+            remaining = total_ep - epoch
+            optimizer, scheduler = build_optimizer_and_scheduler(
+                model, args.lr, remaining, warmup_ep)
+            print(f"  New optimizer created: lr={args.lr}, "
+                  f"remaining_epochs={remaining}, warmup={warmup_ep}")
+
+        # Train
+        train_loss, tr_preds, tr_labels, tr_probs, _ = run_epoch(
+            train_loader, model, criterion, device, task_type,
+            is_training=True, optimizer=optimizer,
+        )
+        # Validate
+        val_loss, val_preds, val_labels, val_probs, _ = run_epoch(
+            val_loader, model, criterion, device, task_type,
+            is_training=False,
+        )
+
+        scheduler.step()
+        torch.cuda.empty_cache()
+
+        # Compute metrics
+        if is_cls:
+            tr_m = calc_classification_metrics(tr_preds, tr_labels, tr_probs)
+            val_m = calc_classification_metrics(val_preds, val_labels, val_probs)
+            current = val_m['acc']
+            improved = current > best_metric
+        else:
+            tr_m = calc_regression_metrics(tr_preds, tr_labels)
+            val_m = calc_regression_metrics(val_preds, val_labels)
+            current = val_m['mae']
+            improved = current < best_metric
+
+        if improved:
+            best_metric = current
+            best_model_state = deepcopy(model.state_dict())
+            patience_counter = 0
+
+            # Save best checkpoint
+            mode_suffix = 'freeze_then_finetune' if freeze_epochs > 0 else 'finetune'
+            save_dir = os.path.join(_BASE_DIR, 'saves', f'multimae_{mode_suffix}')
+            os.makedirs(save_dir, exist_ok=True)
+            task_str = task_type.replace(' ', '_')
+            torch.save({
+                'epoch': epoch + 1,
+                'model_state_dict': best_model_state,
+                'task': task_type,
+                'seed': seed,
+                'best_metric': best_metric,
+                'freeze_epochs': freeze_epochs,
+            }, os.path.join(save_dir, f'{task_str}_seed_{seed}_best.pth'))
+        else:
+            patience_counter += 1
+
+        # Print progress periodically or on improvement
+        if (epoch + 1) % 5 == 0 or improved:
+            if is_cls:
+                print(
+                    f"  Epoch {epoch+1:3d}/{total_ep} | "
+                    f"TrLoss: {train_loss:.4f}, TrAcc: {tr_m['acc']*100:.1f}% | "
+                    f"ValAcc: {val_m['acc']*100:.1f}%, "
+                    f"ValAUC: {val_m['auc']*100:.1f}%"
+                    f"{'  ***' if improved else ''}"
+                )
+            else:
+                print(
+                    f"  Epoch {epoch+1:3d}/{total_ep} | "
+                    f"TrLoss: {train_loss:.4f}, TrMAE: {tr_m['mae']:.4f} | "
+                    f"ValMAE: {val_m['mae']:.4f}, "
+                    f"ValPearson: {val_m['pearson']:.4f}"
+                    f"{'  ***' if improved else ''}"
+                )
+
+        # Early stopping
+        if patience_counter >= args.patience:
+            print(f"  Early stopping at epoch {epoch+1}")
+            break
+
+    # ---- Test evaluation ----
+    if best_model_state is not None:
+        model.load_state_dict(best_model_state)
+
+    print("\n  Evaluating on test set...")
+    test_loss, test_preds, test_labels, test_probs, test_combos = run_epoch(
+        test_loader, model, criterion, device, task_type,
+        is_training=False,
+    )
+
+    # Overall test metrics
+    if is_cls:
+        test_m = calc_classification_metrics(
+            test_preds, test_labels, test_probs)
+        print(
+            f"  Test: Acc={test_m['acc']*100:.2f}%, "
+            f"AUC={test_m['auc']*100:.2f}%, "
+            f"Sen={test_m['sensitivity']*100:.2f}%, "
+            f"Spe={test_m['specificity']*100:.2f}%, "
+            f"F1={test_m['f1']*100:.2f}%"
+        )
+    else:
+        test_m = calc_regression_metrics(test_preds, test_labels)
+        print(
+            f"  Test: MAE={test_m['mae']:.4f}, "
+            f"RMSE={test_m['rmse']:.4f}, "
+            f"Pearson={test_m['pearson']:.4f}"
+        )
+
+    # Per-modality-combination breakdown
+    combo_results = calc_metrics_by_combo(
+        test_preds, test_labels, test_probs, test_combos, task_type)
+    if combo_results:
+        print(f"\n  Per-modality-combination results:")
+        for combo in sorted(combo_results.keys()):
+            r = combo_results[combo]
+            n = r['n_samples']
+            if is_cls:
+                print(f"    {combo:25s} (n={n:3d}) | "
+                      f"Acc={r['acc']*100:.1f}%, AUC={r['auc']*100:.1f}%")
+            else:
+                print(f"    {combo:25s} (n={n:3d}) | "
+                      f"MAE={r['mae']:.4f}, Pearson={r['pearson']:.4f}")
+
+        # Save per-combo results to Excel
+        freeze_epochs = getattr(args, 'freeze_epochs', 0)
+        mode_tag = (f"freeze{freeze_epochs}_finetune"
+                    if freeze_epochs > 0 else "finetune")
+        _save_combo_results(combo_results, task_type, seed, mode_tag, is_cls)
+
+    # Cleanup
+    del model, encoder, optimizer, train_loader, val_loader, test_loader
+    del best_model_state
+    torch.cuda.empty_cache()
+    gc.collect()
+
+    return test_m
+
+
+def _save_combo_results(combo_results, task_type, seed, mode, is_cls):
+    """Save per-modality-combination results to Excel."""
+    results_dir = os.path.join(_BASE_DIR, 'results')
+    os.makedirs(results_dir, exist_ok=True)
+
+    rows = []
+    for combo in sorted(combo_results.keys()):
+        r = combo_results[combo]
+        row = {'Modality': combo, 'N': r['n_samples']}
+        if is_cls:
+            row.update({
+                'Acc': r['acc'] * 100,
+                'AUC': r['auc'] * 100,
+                'Sensitivity': r['sensitivity'] * 100,
+                'Specificity': r['specificity'] * 100,
+                'F1': r['f1'] * 100,
+            })
+        else:
+            row.update({
+                'MAE': r['mae'],
+                'RMSE': r['rmse'],
+                'Pearson': r['pearson'],
+            })
+        rows.append(row)
+
+    task_str = task_type.replace(' ', '_')
+    path = os.path.join(
+        results_dir,
+        f'multimae_{mode}_{task_str}_seed_{seed}_by_combo.xlsx',
+    )
+    pd.DataFrame(rows).to_excel(path, index=False)
+    print(f"  Saved: {path}")
+
+
+# =========================================================================
+# Argument parsing
+# =========================================================================
+
+def parse_args():
+    import argparse
+    p = argparse.ArgumentParser(
+        description='MultiMAE3D Finetuning for Downstream Tasks')
+
+    # Mode
+    p.add_argument('--mode', type=str, default='finetune',
+                   choices=['finetune', 'freeze_then_finetune'],
+                   help='finetune: train all parameters end-to-end; '
+                        'freeze_then_finetune: freeze encoder for N epochs then unfreeze')
+    p.add_argument(
+        '--pretrained', type=str,
+        default=os.path.join(
+            _BASE_DIR, 'pretrain_checkpoints', 'multimae', 'best_model.pth'),
+        help='Path to pretrained MultiMAE checkpoint')
+
+    # Tasks & seeds
+    p.add_argument('--tasks', type=str, nargs='+',
+                   default=['CN vs AD', 'CN vs MCI', 'MMSE', 'AGE'],
+                   help='Tasks to evaluate')
+    p.add_argument('--n_seeds', type=int, default=3,
+                   help='Number of random seeds per task')
+
+    # Data
+    p.add_argument('--train_excel', type=str,
+                   default='./data/Downstream/'
+                           'ADNI_Division/modality_data_train.xlsx')
+    p.add_argument('--val_excel', type=str,
+                   default='./data/Downstream/'
+                           'ADNI_Division/modality_data_val.xlsx')
+    p.add_argument('--test_excel', type=str,
+                   default='./data/Downstream/'
+                           'ADNI_Division/modality_data_test.xlsx')
+    p.add_argument('--base_dir', type=str,
+                   default='./data/Downstream/ADNI/')
+    p.add_argument('--modalities', type=str, nargs='+',
+                   default=['T1', 'T2', 'Flair', 'PET'])
+    p.add_argument('--intersection', type=str2bool, default=False)
+    p.add_argument('--image_size', type=int, nargs=3,
+                   default=[128, 128, 128])
+    p.add_argument('--batch_size', type=int, default=4)
+    p.add_argument('--num_workers', type=int, default=8)
+
+    # MultiMAE encoder architecture (must match pretrained checkpoint)
+    p.add_argument('--img_size', type=int, default=128)
+    p.add_argument('--patch_size', type=int, default=16)
+    p.add_argument('--embed_dim', type=int, default=768)
+    p.add_argument('--depth', type=int, default=12)
+    p.add_argument('--num_heads', type=int, default=12)
+    p.add_argument('--decoder_embed_dim', type=int, default=384)
+    p.add_argument('--decoder_depth', type=int, default=2)
+    p.add_argument('--decoder_num_heads', type=int, default=12)
+
+    # Downstream head
+    p.add_argument('--pool', type=str, default='cls',
+                   choices=['cls', 'mean'],
+                   help='Feature pooling: cls token or mean pool')
+    p.add_argument('--dropout', type=float, default=0.1)
+
+    # Training
+    p.add_argument('--epochs', type=int, default=50)
+    p.add_argument('--lr', type=float, default=5e-5,
+                   help='Learning rate')
+    p.add_argument('--weight_decay', type=float, default=0.05)
+    p.add_argument('--warmup_epochs', type=int, default=5)
+    p.add_argument('--patience', type=int, default=15,
+                   help='Early stopping patience')
+    p.add_argument('--freeze_epochs', type=int, default=0,
+                   help='Number of epochs to freeze pretrained encoder '
+                        '(0 = no freeze, full finetune from start)')
+
+    # Device
+    p.add_argument('--device', type=int, default=0)
+
+    return p.parse_args()
+
+
+# =========================================================================
+# Main: loop over tasks x seeds
+# =========================================================================
+
+def main():
+    args = parse_args()
+    device = torch.device(
+        f'cuda:{args.device}' if torch.cuda.is_available() else 'cpu')
+
+    print("=" * 80)
+    print(f"MultiMAE3D Downstream Evaluation")
+    print(f"  Mode         : {args.mode}")
+    print(f"  Tasks        : {args.tasks}")
+    print(f"  Seeds        : {args.n_seeds}")
+    print(f"  Pretrained   : {args.pretrained}")
+    print(f"  Pool         : {args.pool}")
+    print(f"  Device       : {device}")
+    print(f"  LR           : {args.lr}")
+    print(f"  Epochs       : {args.epochs}")
+    print(f"  Batch size   : {args.batch_size}")
+    if args.freeze_epochs > 0:
+        print(f"  Freeze epochs: {args.freeze_epochs} (encoder frozen, then joint training)")
+    print("=" * 80)
+
+    # Logger
+    log_dir = os.path.join(_BASE_DIR, 'logs')
+    logger = setup_logger(log_dir, 'multimae_ft',
+                          f'multimae_{args.mode}.txt')
+
+    all_results = {}
+
+    for task_type in args.tasks:
+        print(f"\n{'='*80}")
+        print(f"TASK: {task_type}")
+        print(f"{'='*80}")
+
+        is_cls = task_type in ('CN vs AD', 'CN vs MCI')
+        seed_results = []
+
+        for seed in range(args.n_seeds):
+            print(f"\n--- Seed {seed} ---")
+            metrics = train_and_evaluate(args, task_type, seed, device)
+            seed_results.append(metrics)
+
+        all_results[task_type] = seed_results
+
+        # Per-task summary
+        print(f"\n{task_type} Summary ({args.n_seeds} seeds):")
+        summary_str = f"[{args.mode}] {task_type}: "
+        if is_cls:
+            for key in ['acc', 'auc', 'sensitivity', 'specificity', 'f1']:
+                vals = [r[key] * 100 for r in seed_results]
+                msg = f"{np.mean(vals):.2f} +/- {np.std(vals):.2f}%"
+                print(f"  {key:>12s}: {msg}")
+                summary_str += f"{key}={msg}, "
+        else:
+            for key in ['mae', 'rmse', 'pearson']:
+                vals = [r[key] for r in seed_results]
+                msg = f"{np.mean(vals):.4f} +/- {np.std(vals):.4f}"
+                print(f"  {key:>12s}: {msg}")
+                summary_str += f"{key}={msg}, "
+        logger.info(summary_str)
+
+    # ---- Final summary table ----
+    print("\n" + "=" * 80)
+    print("FINAL SUMMARY")
+    print("=" * 80)
+
+    summary_rows = []
+
+    for task_type in args.tasks:
+        results = all_results[task_type]
+        is_cls = task_type in ('CN vs AD', 'CN vs MCI')
+        row = {'Task': task_type, 'Mode': args.mode}
+
+        if is_cls:
+            for key in ['acc', 'auc', 'sensitivity', 'specificity', 'f1']:
+                vals = [r[key] * 100 for r in results]
+                row[f'{key}_mean'] = np.mean(vals)
+                row[f'{key}_std'] = np.std(vals)
+                row[key] = f"{np.mean(vals):.2f}+/-{np.std(vals):.2f}"
+            # Per-seed values
+            for i, r in enumerate(results):
+                row[f'seed_{i}_acc'] = r['acc'] * 100
+                row[f'seed_{i}_auc'] = r['auc'] * 100
+
+            vals_acc = [r['acc'] * 100 for r in results]
+            vals_auc = [r['auc'] * 100 for r in results]
+            print(f"  {task_type:12s} | "
+                  f"Acc: {np.mean(vals_acc):.2f}+/-{np.std(vals_acc):.2f}% | "
+                  f"AUC: {np.mean(vals_auc):.2f}+/-{np.std(vals_auc):.2f}%")
+        else:
+            for key in ['mae', 'rmse', 'pearson']:
+                vals = [r[key] for r in results]
+                row[f'{key}_mean'] = np.mean(vals)
+                row[f'{key}_std'] = np.std(vals)
+                row[key] = f"{np.mean(vals):.4f}+/-{np.std(vals):.4f}"
+            for i, r in enumerate(results):
+                row[f'seed_{i}_mae'] = r['mae']
+                row[f'seed_{i}_pearson'] = r['pearson']
+
+            vals_mae = [r['mae'] for r in results]
+            vals_r = [r['pearson'] for r in results]
+            print(f"  {task_type:12s} | "
+                  f"MAE: {np.mean(vals_mae):.4f}+/-{np.std(vals_mae):.4f} | "
+                  f"Pearson: {np.mean(vals_r):.4f}+/-{np.std(vals_r):.4f}")
+
+        summary_rows.append(row)
+
+    # Save summary Excel
+    results_dir = os.path.join(_BASE_DIR, 'results')
+    os.makedirs(results_dir, exist_ok=True)
+    freeze_epochs = getattr(args, 'freeze_epochs', 0)
+    summary_tag = (f"freeze{freeze_epochs}_finetune"
+                   if freeze_epochs > 0 else "finetune")
+    summary_path = os.path.join(
+        results_dir, f'multimae_{summary_tag}_summary.xlsx')
+    pd.DataFrame(summary_rows).to_excel(summary_path, index=False)
+    print(f"\nSummary saved to: {summary_path}")
+    print("=" * 80)
+
+
+if __name__ == '__main__':
+    main()

BrainAnytime/models/multimae3d.py

@@ -0,0 +1,997 @@
+"""
+MultiMAE3D: Multi-modal Masked Autoencoder for 3D Medical Images
+
+Architecture:
+- Per-modality input adapters (Conv3D patch embedding)
+- Shared ViT encoder
+- Per-modality output adapters (cross-attn decoder)
+- Handles arbitrary missing modalities via observed mask
+
+Based on MultiMAE_reference, simplified for our use case:
+- Fixed input size 128^3, 4 modalities (T1, T2, Flair, PET)
+- Pure reconstruction pretraining (MSE loss)
+- No Hydra/Lightning dependencies
+"""
+
+import copy
+import math
+from typing import Union, Tuple, Dict, List, Optional
+from collections import OrderedDict
+from functools import partial
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from timm.layers import DropPath
+from einops import rearrange
+
+from models.multimae3d_utils import (
+    to_3tuple,
+    calc_patchified_dim,
+    patchify,
+    unpatchify,
+    shuffle_patches,
+    unshuffle_patches,
+    build_3d_sincos_position_embedding,
+    mask_data,
+)
+
+
+# =============================================================================
+# Input Adapter: Conv3D patch embedding (per modality)
+# =============================================================================
+
+class PatchedInputAdapter(nn.Module):
+    """
+    Converts a single-channel 3D volume into patch tokens.
+    Input:  [B, N_selected, 1, pd, ph, pw] (selected shuffled patches)
+    Output: [B, N_selected, embed_dim]
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 1,
+        patch_size: Union[int, Tuple[int, int, int]] = 16,
+        embed_dim: int = 768,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.patch_size = to_3tuple(patch_size)
+        self.embed_dim = embed_dim
+
+        # Conv3D projection: each patch -> embed_dim
+        self.proj = nn.Conv3d(
+            in_channels,
+            embed_dim,
+            kernel_size=self.patch_size,
+            stride=self.patch_size,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        x: [B, N, C, pd, ph, pw] selected patches (already patchified & shuffled)
+        returns: [B, N, embed_dim]
+        """
+        B, N = x.shape[0], x.shape[1]
+        # Merge batch and patch dims for Conv3D
+        x = rearrange(x, "b n c d h w -> (b n) c d h w")
+        x = self.proj(x)  # [(B*N), embed_dim, 1, 1, 1]
+        x = x.flatten(2)  # [(B*N), embed_dim, 1]
+        x = x.squeeze(-1)  # [(B*N), embed_dim]
+        x = rearrange(x, "(b n) d -> b n d", b=B)
+        return x
+
+
+# =============================================================================
+# Cross Attention (for decoder)
+# =============================================================================
+
+class CrossAttention(nn.Module):
+    """Cross attention: query attends to context (encoder output)."""
+
+    def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = True,
+                 attn_drop: float = 0.0, proj_drop: float = 0.0):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim ** -0.5
+
+        self.q = nn.Linear(dim, dim, bias=qkv_bias)
+        self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x: torch.Tensor, context: torch.Tensor) -> torch.Tensor:
+        B, N, C = x.shape
+        _, M, _ = context.shape
+
+        q = self.q(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
+        kv = self.kv(context).reshape(B, M, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        k, v = kv[0], kv[1]
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, -1)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+# =============================================================================
+# Transformer blocks with attention mask support
+# =============================================================================
+
+class Mlp(nn.Module):
+    """Simple MLP with GELU activation."""
+
+    def __init__(self, in_features, hidden_features=None, out_features=None,
+                 act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class MaskedAttention(nn.Module):
+    """Multi-head self-attention with optional additive attention mask."""
+
+    def __init__(self, dim, num_heads=12, qkv_bias=True,
+                 attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, attn_mask=None):
+        """
+        x: [B, N, C]
+        attn_mask: [B, 1, 1, N] additive mask, -inf for tokens to ignore (column masking)
+        """
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)  # each [B, num_heads, N, head_dim]
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale  # [B, num_heads, N, N]
+        if attn_mask is not None:
+            attn = attn + attn_mask
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class MaskedBlock(nn.Module):
+    """Pre-LN Transformer block with optional attention mask support.
+    Used for both encoder (with mask) and decoder (without mask).
+    """
+
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=True,
+                 drop_path=0., act_layer=nn.GELU,
+                 norm_layer=partial(nn.LayerNorm, eps=1e-6)):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = MaskedAttention(dim, num_heads=num_heads, qkv_bias=qkv_bias)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden, act_layer=act_layer)
+
+    def forward(self, x, attn_mask=None):
+        x = x + self.drop_path(self.attn(self.norm1(x), attn_mask=attn_mask))
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
+# =============================================================================
+# Cross-Modal Predictor (for cross-level mutual prediction)
+# =============================================================================
+
+class CrossModalPredictor(nn.Module):
+    """3-layer MLP predictor for cross-modal feature prediction.
+
+    Maps features from one modality space to another.
+    Structure: Linear(D, 2D) → GELU → Linear(2D, 2D) → GELU → Linear(2D, D)
+    """
+
+    def __init__(self, dim: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(dim, dim * 2),
+            nn.GELU(),
+            nn.Linear(dim * 2, dim * 2),
+            nn.GELU(),
+            nn.Linear(dim * 2, dim),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.net(x)
+
+
+# =============================================================================
+# Output Adapter: Decoder (per modality)
+# =============================================================================
+
+class SpatialOutputAdapter(nn.Module):
+    """
+    Per-modality decoder.
+    Takes encoder tokens, adds mask tokens, applies cross-attention + self-attention,
+    then projects back to patch pixel space.
+
+    Architecture:
+    1. Project encoder tokens from encoder_dim -> decoder_dim
+    2. Create mask tokens for masked positions
+    3. Add positional embedding to query (mask + selected tokens)
+    4. Cross-attention: query attends to encoder context
+    5. Self-attention transformer blocks
+    6. Linear projection to patch pixel dimension
+    """
+
+    def __init__(
+        self,
+        out_channels: int = 1,
+        img_size: Union[int, Tuple[int, int, int]] = 128,
+        patch_size: Union[int, Tuple[int, int, int]] = 16,
+        encoder_embed_dim: int = 768,
+        embed_dim: int = 384,
+        num_heads: int = 12,
+        depth: int = 2,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = True,
+    ):
+        super().__init__()
+        self.out_channels = out_channels
+        self.img_size = to_3tuple(img_size)
+        self.patch_size = to_3tuple(patch_size)
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.depth = depth
+
+        self.patchified_dim = calc_patchified_dim(self.img_size, self.patch_size)
+        self.num_patches = self.patchified_dim[0] * self.patchified_dim[1] * self.patchified_dim[2]
+
+        # Project encoder tokens to decoder dimension
+        self.proj_context = nn.Linear(encoder_embed_dim, embed_dim)
+
+        # Learnable mask token
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        nn.init.normal_(self.mask_token, std=0.02)
+
+        # Decoder positional embedding (sincos, frozen)
+        self.pos_embed = build_3d_sincos_position_embedding(
+            grid_size=self.patchified_dim,
+            embed_dim=embed_dim,
+        )
+
+        # Cross-attention + MLP (MultiMAE style)
+        self.xattn = CrossAttention(
+            dim=embed_dim, num_heads=num_heads, qkv_bias=qkv_bias,
+        )
+        norm_layer = partial(nn.LayerNorm, eps=1e-6)
+        self.context_norm = norm_layer(embed_dim)
+        self.query_norm = norm_layer(embed_dim)
+        self.out_norm = norm_layer(embed_dim)
+        mlp_hidden = int(embed_dim * mlp_ratio)
+        self.mlp = nn.Sequential(
+            nn.Linear(embed_dim, mlp_hidden),
+            nn.GELU(),
+            nn.Linear(mlp_hidden, embed_dim),
+        )
+
+        # Self-attention transformer blocks (decoder: no attention mask needed)
+        self.blocks = nn.Sequential(*[
+            MaskedBlock(
+                dim=embed_dim,
+                num_heads=num_heads,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                act_layer=nn.GELU,
+                norm_layer=norm_layer,
+            )
+            for _ in range(depth)
+        ]) if depth > 0 else nn.Identity()
+
+        # Output projection: decoder_dim -> patch_pixels
+        dim_patch = self.patch_size[0] * self.patch_size[1] * self.patch_size[2] * out_channels
+        self.out_proj = nn.Linear(embed_dim, dim_patch)
+
+    def forward(
+        self,
+        encoder_tokens: torch.Tensor,
+        task_range: Tuple[int, int],
+        perm_idx: torch.Tensor,
+        num_patches: int,
+    ) -> torch.Tensor:
+        """
+        Args:
+            encoder_tokens: [B, total_visible_tokens, encoder_dim] (last layer output)
+            task_range: (start, end) indices of this modality's tokens in the concat
+            perm_idx: [B, num_patches] permutation indices for this modality
+            num_patches: total number of patches for this modality
+
+        Returns:
+            output: [B, num_patches, out_channels, pd, ph, pw] (all patches, unshuffled order)
+        """
+        B = encoder_tokens.shape[0]
+
+        # 1. Project encoder tokens to decoder dim
+        context = self.proj_context(encoder_tokens)
+
+        # 2. Extract this modality's selected tokens from the context
+        num_selected = task_range[1] - task_range[0]
+        selected_tokens = context[:, task_range[0]:task_range[1]]
+
+        # 3. Create mask tokens for masked positions
+        num_masked = num_patches - num_selected
+        mask_tokens = self.mask_token.repeat(B, num_masked, 1)
+
+        # 4. Concatenate: [selected, masked] in shuffled order
+        query = torch.cat([selected_tokens, mask_tokens], dim=1)  # [B, num_patches, dim]
+
+        # 5. Add positional embedding (following the permutation order)
+        pos_emb = self.pos_embed.expand(B, -1, -1)  # [B, num_patches, dim]
+        pos_emb_shuffled = pos_emb[torch.arange(B, device=pos_emb.device)[:, None], perm_idx]
+        query = query + pos_emb_shuffled
+
+        # 6. Cross-attention + MLP
+        x = self.xattn(self.query_norm(query), self.context_norm(context))
+        x = x + self.mlp(self.out_norm(x))
+
+        # 7. Self-attention blocks
+        if self.depth > 0:
+            x = self.blocks(x)
+
+        # 8. Project to patch pixel space
+        x = self.out_proj(x)  # [B, num_patches, patch_pixels]
+
+        # 9. Reshape to patch format
+        x = rearrange(
+            x,
+            "b n (c pd ph pw) -> b n c pd ph pw",
+            c=self.out_channels,
+            pd=self.patch_size[0],
+            ph=self.patch_size[1],
+            pw=self.patch_size[2],
+        )
+
+        # 10. Unshuffle back to spatial order
+        x = unshuffle_patches(x, perm_idx)
+
+        return x
+
+
+# =============================================================================
+# MultiMAE3D: Main Model
+# =============================================================================
+
+class MultiMAE3D(nn.Module):
+    """
+    Multi-modal Masked Autoencoder for 3D Medical Images.
+
+    Handles 4 modalities (T1, T2, Flair, PET) with arbitrary missing modalities.
+
+    Forward pass:
+    1. Split stacked input into per-modality volumes
+    2. Patchify and mask each modality (missing → 100% masked)
+    3. Tokenize visible patches via per-modality input adapters
+    4. Add positional embeddings + CLS token
+    5. Concatenate all visible tokens → shared ViT encoder
+    6. Per-modality decoder → reconstruct masked patches
+    7. Compute MSE loss only on present modalities' masked patches
+    """
+
+    MODALITY_NAMES = ["T1", "T2", "Flair", "PET"]
+
+    def __init__(
+        self,
+        img_size: Union[int, Tuple[int, int, int]] = 128,
+        patch_size: Union[int, Tuple[int, int, int]] = 16,
+        embed_dim: int = 768,
+        depth: int = 12,
+        num_heads: int = 12,
+        mlp_ratio: float = 4.0,
+        decoder_embed_dim: int = 384,
+        decoder_depth: int = 2,
+        decoder_num_heads: int = 12,
+        mask_ratio: float = 0.75,
+        use_dirichlet: bool = True,
+        dirichlet_alpha: float = 1.0,
+        num_global_tokens: int = 1,
+        qkv_bias: bool = True,
+        drop_path_rate: float = 0.0,
+        enable_cross_modal: bool = False,
+    ):
+        super().__init__()
+
+        self.img_size = to_3tuple(img_size)
+        self.patch_size = to_3tuple(patch_size)
+        self.embed_dim = embed_dim
+        self.depth = depth
+        self.mask_ratio = mask_ratio
+        self.use_dirichlet = use_dirichlet
+        self.dirichlet_alpha = dirichlet_alpha
+        self.num_global_tokens = num_global_tokens
+        self.enable_cross_modal = enable_cross_modal
+
+        self.patchified_dim = calc_patchified_dim(self.img_size, self.patch_size)
+        self.num_patches = self.patchified_dim[0] * self.patchified_dim[1] * self.patchified_dim[2]
+
+        # ----- Input adapters (per modality) -----
+        self.input_adapters = nn.ModuleDict({
+            name: PatchedInputAdapter(
+                in_channels=1,
+                patch_size=patch_size,
+                embed_dim=embed_dim,
+            )
+            for name in self.MODALITY_NAMES
+        })
+
+        # ----- Encoder positional embedding (sincos, frozen) -----
+        self.pos_embed = build_3d_sincos_position_embedding(
+            grid_size=self.patchified_dim,
+            embed_dim=embed_dim,
+        )
+
+        # ----- CLS token -----
+        if num_global_tokens > 0:
+            self.global_tokens = nn.Parameter(torch.zeros(num_global_tokens, embed_dim))
+            nn.init.normal_(self.global_tokens, std=0.02)
+
+        # ----- Shared Transformer encoder (ModuleList for attn_mask support) -----
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]
+        norm_layer = partial(nn.LayerNorm, eps=1e-6)
+        self.encoder = nn.ModuleList([
+            MaskedBlock(
+                dim=embed_dim,
+                num_heads=num_heads,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                drop_path=dpr[i],
+                act_layer=nn.GELU,
+                norm_layer=norm_layer,
+            )
+            for i in range(depth)
+        ])
+
+        # ----- Output adapters / decoders (per modality) -----
+        self.output_adapters = nn.ModuleDict({
+            name: SpatialOutputAdapter(
+                out_channels=1,
+                img_size=img_size,
+                patch_size=patch_size,
+                encoder_embed_dim=embed_dim,
+                embed_dim=decoder_embed_dim,
+                num_heads=decoder_num_heads,
+                depth=decoder_depth,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+            )
+            for name in self.MODALITY_NAMES
+        })
+
+        # Initialize weights
+        self._initialize_weights()
+
+        # ----- Cross-modal mutual prediction components -----
+        if self.enable_cross_modal:
+            # Teacher encoder (EMA copy of student) — no gradients
+            self.teacher_input_adapters = copy.deepcopy(self.input_adapters)
+            for p in self.teacher_input_adapters.parameters():
+                p.requires_grad = False
+
+            self.teacher_encoder = copy.deepcopy(self.encoder)
+            for p in self.teacher_encoder.parameters():
+                p.requires_grad = False
+
+            # Teacher global tokens stored as buffer (auto-moves with .to(device))
+            if self.num_global_tokens > 0:
+                self.register_buffer(
+                    "teacher_global_tokens",
+                    self.global_tokens.data.clone(),
+                )
+
+            # Cross-modal predictors (student-only, learnable)
+            self.predictor_mri_to_pet = CrossModalPredictor(embed_dim)
+            self.predictor_pet_to_mri = CrossModalPredictor(embed_dim)
+            # Initialize predictor weights
+            self.predictor_mri_to_pet.apply(self._init_weights)
+            self.predictor_pet_to_mri.apply(self._init_weights)
+
+    def _initialize_weights(self):
+        self.apply(self._init_weights)
+        # Special init for Conv3D projection (following MAE)
+        for name, m in self.named_modules():
+            if isinstance(m, nn.Linear):
+                if "qkv" in name:
+                    val = math.sqrt(6.0 / float(m.weight.shape[0] // 3 + m.weight.shape[1]))
+                    nn.init.uniform_(m.weight, -val, val)
+                elif "kv" in name:
+                    val = math.sqrt(6.0 / float(m.weight.shape[0] // 2 + m.weight.shape[1]))
+                    nn.init.uniform_(m.weight, -val, val)
+            if isinstance(m, nn.Conv3d):
+                if ".proj" in name:
+                    w = m.weight.data
+                    nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+
+    @staticmethod
+    def _init_weights(m):
+        if isinstance(m, nn.Linear):
+            nn.init.xavier_uniform_(m.weight)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def _split_modalities(self, images: torch.Tensor) -> Dict[str, torch.Tensor]:
+        """Split stacked [B, 4, D, H, W] into per-modality dict {name: [B, 1, D, H, W]}."""
+        return {
+            name: images[:, i:i+1]
+            for i, name in enumerate(self.MODALITY_NAMES)
+        }
+
+    # -----------------------------------------------------------------
+    # Cross-modal mutual prediction helpers
+    # -----------------------------------------------------------------
+
+    def _encode_with(
+        self,
+        selected_patches: Dict[str, torch.Tensor],
+        perm_indices: Dict[str, torch.Tensor],
+        observed: torch.Tensor,
+        input_adapters: nn.ModuleDict,
+        global_tokens,
+        encoder_blocks: nn.ModuleList,
+    ):
+        """
+        Shared encoding logic used by both student and teacher.
+
+        Returns:
+            encoder_output: [B, total_tokens, D] or None
+            task_ranges: OrderedDict {modality_name: (start, end)}
+        """
+        B = observed.shape[0]
+        device = observed.device
+
+        tokens = {}
+        for name in self.MODALITY_NAMES:
+            sel = selected_patches[name]
+            if sel.shape[1] == 0:
+                continue
+            tok = input_adapters[name](sel)
+            perm = perm_indices[name]
+            pos_emb = self.pos_embed.expand(B, -1, -1)
+            pos_emb_selected = pos_emb[
+                torch.arange(B, device=device)[:, None], perm[:, :sel.shape[1]]
+            ]
+            tok = tok + pos_emb_selected
+            tokens[name] = tok
+
+        token_list = []
+        task_ranges = OrderedDict()
+        offset = self.num_global_tokens
+
+        for name in self.MODALITY_NAMES:
+            if name in tokens:
+                n_tok = tokens[name].shape[1]
+                task_ranges[name] = (offset, offset + n_tok)
+                token_list.append(tokens[name])
+                offset += n_tok
+            else:
+                task_ranges[name] = (offset, offset)
+
+        if len(token_list) == 0:
+            return None, task_ranges
+
+        input_tokens = torch.cat(token_list, dim=1)
+
+        if self.num_global_tokens > 0 and global_tokens is not None:
+            if global_tokens.dim() == 2:
+                cls = global_tokens.unsqueeze(0).expand(B, -1, -1)
+            else:
+                cls = global_tokens.expand(B, -1, -1)
+            input_tokens = torch.cat([cls, input_tokens], dim=1)
+
+        # Column masking for missing modalities
+        total_tokens = input_tokens.shape[1]
+        attn_mask = torch.zeros(B, 1, 1, total_tokens, device=device)
+        for i, name in enumerate(self.MODALITY_NAMES):
+            start, end = task_ranges[name]
+            if start == end:
+                continue
+            missing = (observed[:, i] < 0.5)
+            if missing.any():
+                attn_mask[missing, :, :, start:end] = float("-inf")
+        if (attn_mask == 0).all():
+            attn_mask = None
+
+        encoder_output = input_tokens
+        for block in encoder_blocks:
+            encoder_output = block(encoder_output, attn_mask=attn_mask)
+
+        return encoder_output, task_ranges
+
+    def _compute_cross_modal_loss(
+        self,
+        selected_patches: Dict[str, torch.Tensor],
+        perm_indices: Dict[str, torch.Tensor],
+        observed: torch.Tensor,
+        student_encoder_output: torch.Tensor,
+        task_ranges: OrderedDict,
+    ) -> torch.Tensor:
+        """
+        Cross-level mutual prediction loss (simplified global-average-pooling version).
+
+        Two groups:
+          - MRI group: all T1 + T2 + Flair tokens  →  z_MRI  (D-dim vector)
+          - PET group: all PET tokens               →  z_PET  (D-dim vector)
+
+        Predictions (student → teacher):
+          - predictor_mri_to_pet(z_MRI_student) → predict z_PET_teacher
+          - predictor_pet_to_mri(z_PET_student) → predict z_MRI_teacher
+
+        Loss: negative cosine similarity, averaged over paired samples only.
+        """
+        B = observed.shape[0]
+        device = observed.device
+
+        # Paired = has at least one MRI modality AND PET
+        has_mri = (observed[:, :3].sum(dim=1) > 0.5)  # [B]
+        has_pet = (observed[:, 3] > 0.5)               # [B]
+        is_paired = has_mri & has_pet                   # [B]
+
+        if not is_paired.any():
+            return torch.tensor(0.0, device=device, requires_grad=True)
+
+        # --- Teacher forward (no gradients) ---
+        with torch.no_grad():
+            teacher_gt = (
+                self.teacher_global_tokens
+                if self.num_global_tokens > 0 else None
+            )
+            teacher_output, _ = self._encode_with(
+                selected_patches, perm_indices, observed,
+                self.teacher_input_adapters, teacher_gt,
+                self.teacher_encoder,
+            )
+        if teacher_output is None:
+            return torch.tensor(0.0, device=device, requires_grad=True)
+
+        # --- Build group masks [B, L] ---
+        total_tokens = student_encoder_output.shape[1]
+        mri_mask = torch.zeros(B, total_tokens, device=device)
+        pet_mask = torch.zeros(B, total_tokens, device=device)
+
+        # MRI group: T1 (idx 0), T2 (idx 1), Flair (idx 2)
+        for idx, name in enumerate(["T1", "T2", "Flair"]):
+            start, end = task_ranges[name]
+            if start < end:
+                mri_mask[:, start:end] = observed[:, idx:idx+1].expand(-1, end - start)
+
+        # PET group: idx 3
+        start, end = task_ranges["PET"]
+        if start < end:
+            pet_mask[:, start:end] = observed[:, 3:4].expand(-1, end - start)
+
+        # --- Global average pooling per group ---
+        mri_count = mri_mask.sum(dim=1, keepdim=True).clamp(min=1)
+        pet_count = pet_mask.sum(dim=1, keepdim=True).clamp(min=1)
+
+        z_mri_s = (student_encoder_output * mri_mask.unsqueeze(-1)).sum(dim=1) / mri_count  # [B, D]
+        z_pet_s = (student_encoder_output * pet_mask.unsqueeze(-1)).sum(dim=1) / pet_count  # [B, D]
+
+        z_mri_t = (teacher_output * mri_mask.unsqueeze(-1)).sum(dim=1) / mri_count  # [B, D]
+        z_pet_t = (teacher_output * pet_mask.unsqueeze(-1)).sum(dim=1) / pet_count  # [B, D]
+
+        # --- L2 normalize onto unit hypersphere ---
+        z_mri_s = F.normalize(z_mri_s, dim=-1)
+        z_pet_s = F.normalize(z_pet_s, dim=-1)
+        z_mri_t = F.normalize(z_mri_t, dim=-1)
+        z_pet_t = F.normalize(z_pet_t, dim=-1)
+
+        # --- Cross-modal predictions + normalize ---
+        pred_pet = F.normalize(self.predictor_mri_to_pet(z_mri_s), dim=-1)  # [B, D]
+        pred_mri = F.normalize(self.predictor_pet_to_mri(z_pet_s), dim=-1)  # [B, D]
+
+        # --- Negative cosine similarity: L = 2 - 2·cos(pred, target) ---
+        loss_m2p = 2 - 2 * (pred_pet * z_pet_t.detach()).sum(dim=-1)  # [B]
+        loss_p2m = 2 - 2 * (pred_mri * z_mri_t.detach()).sum(dim=-1)  # [B]
+
+        # Average only over paired samples
+        paired_f = is_paired.float()
+        n_paired = paired_f.sum().clamp(min=1)
+
+        loss_m2p = (loss_m2p * paired_f).sum() / n_paired
+        loss_p2m = (loss_p2m * paired_f).sum() / n_paired
+
+        return 0.5 * (loss_m2p + loss_p2m)
+
+    @torch.no_grad()
+    def update_teacher(self, momentum: float):
+        """EMA update: θ_teacher ← m·θ_teacher + (1-m)·θ_student."""
+        if not self.enable_cross_modal:
+            return
+
+        for p_s, p_t in zip(
+            self.input_adapters.parameters(),
+            self.teacher_input_adapters.parameters(),
+        ):
+            p_t.data.mul_(momentum).add_(p_s.data, alpha=1 - momentum)
+
+        if self.num_global_tokens > 0:
+            self.teacher_global_tokens.mul_(momentum).add_(
+                self.global_tokens.data, alpha=1 - momentum
+            )
+
+        for p_s, p_t in zip(
+            self.encoder.parameters(),
+            self.teacher_encoder.parameters(),
+        ):
+            p_t.data.mul_(momentum).add_(p_s.data, alpha=1 - momentum)
+
+    @torch.no_grad()
+    def init_teacher_from_student(self):
+        """Copy current student weights to teacher (call after loading checkpoint)."""
+        if not self.enable_cross_modal:
+            return
+
+        for p_s, p_t in zip(
+            self.input_adapters.parameters(),
+            self.teacher_input_adapters.parameters(),
+        ):
+            p_t.data.copy_(p_s.data)
+
+        if self.num_global_tokens > 0:
+            self.teacher_global_tokens.copy_(self.global_tokens.data)
+
+        for p_s, p_t in zip(
+            self.encoder.parameters(),
+            self.teacher_encoder.parameters(),
+        ):
+            p_t.data.copy_(p_s.data)
+
+    def forward(
+        self,
+        images: torch.Tensor,
+        observed: torch.Tensor,
+        return_loss: bool = True,
+        patch_mask_probs: torch.Tensor = None,
+    ) -> Dict[str, torch.Tensor]:
+        """
+        Args:
+            images: [B, 4, D, H, W] stacked multi-modal 3D volumes
+            observed: [B, 4] float tensor, 1.0=present, 0.0=missing
+            return_loss: if True, compute and return reconstruction loss
+            patch_mask_probs: optional [N_patches] per-patch masking probability
+                from anatomy-aware masking (higher = more likely to be masked)
+
+        Returns:
+            dict with:
+                'loss': scalar MSE loss (if return_loss=True)
+                'per_modality_loss': {name: loss} for each present modality
+                'mask_ratios': {name: float} actual mask ratios used
+        """
+        B = images.shape[0]
+        device = images.device
+
+        # 1. Split into per-modality dict
+        batch = self._split_modalities(images)
+
+        # 2. Mask data (patchify + shuffle + split)
+        #    When patch_mask_probs is provided, uses anatomy-aware weighted sampling
+        selected_patches, masked_patches, perm_indices, mask_ratios = mask_data(
+            batch=batch,
+            modality_names=self.MODALITY_NAMES,
+            observed=observed,
+            mask_ratio=self.mask_ratio,
+            patch_size=self.patch_size,
+            use_dirichlet=self.use_dirichlet if self.training else False,
+            dirichlet_alpha=self.dirichlet_alpha,
+            patch_mask_probs=patch_mask_probs if self.training else None,
+        )
+
+        # 3-6. Student encoding (tokenize → concat → attn mask → encoder)
+        encoder_output, task_ranges = self._encode_with(
+            selected_patches, perm_indices, observed,
+            self.input_adapters, self.global_tokens, self.encoder,
+        )
+
+        if encoder_output is None:
+            return {
+                "loss": torch.tensor(0.0, device=device),
+                "cross_modal_loss": torch.tensor(0.0, device=device),
+                "per_modality_loss": {},
+                "mask_ratios": mask_ratios,
+            }
+
+        # 7. Per-modality decoder
+        reconstructed = {}
+        for name in self.MODALITY_NAMES:
+            reconstructed[name] = self.output_adapters[name](
+                encoder_tokens=encoder_output,
+                task_range=task_ranges[name],
+                perm_idx=perm_indices[name],
+                num_patches=self.num_patches,
+            )
+            # reconstructed[name]: [B, num_patches, 1, pd, ph, pw] in spatial order
+
+        # 8. Compute reconstruction loss (MSE, only on present modalities' masked patches)
+        if return_loss:
+            total_loss = torch.tensor(0.0, device=device)
+            per_mod_loss = {}
+            num_present = 0
+
+            for i, name in enumerate(self.MODALITY_NAMES):
+                # Only compute loss on present modalities
+                mod_observed = observed[:, i]  # [B]
+                if mod_observed.sum() < 0.5:
+                    continue
+
+                # Ground truth: all patches in spatial order
+                gt_patches = patchify(batch[name], self.patch_size)  # [B, num_patches, 1, pd, ph, pw]
+                pred_patches = reconstructed[name]  # [B, num_patches, 1, pd, ph, pw]
+
+                # Create per-patch mask: 1 = masked (should reconstruct), 0 = visible
+                perm = perm_indices[name]
+                num_selected = selected_patches[name].shape[1]
+                # In shuffled order: first num_selected are visible, rest masked
+                # Convert to spatial order mask (vectorized, no Python loop)
+                mask = torch.ones(B, self.num_patches, device=device)
+                if num_selected > 0:
+                    selected_perm = perm[:, :num_selected]  # [B, num_selected]
+                    mask.scatter_(1, selected_perm, 0.0)
+
+                # Per-sample observed mask: zero out loss for missing samples
+                sample_mask = mod_observed.float()  # [B]
+
+                # Patch normalization (per-patch zero-mean unit-variance, like original MAE)
+                gt_mean = gt_patches.mean(dim=(2, 3, 4, 5), keepdim=True)
+                gt_var = gt_patches.var(dim=(2, 3, 4, 5), keepdim=True)
+                gt_patches_norm = (gt_patches - gt_mean) / (gt_var + 1e-6).sqrt()
+
+                # Compute MSE on masked patches only (against normalized targets)
+                per_patch_mse = ((pred_patches - gt_patches_norm) ** 2).mean(dim=(2, 3, 4, 5))  # [B, num_patches]
+                masked_mse = (per_patch_mse * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)  # [B]
+                mod_loss = (masked_mse * sample_mask).sum() / sample_mask.sum().clamp(min=1)
+
+                per_mod_loss[name] = mod_loss
+                total_loss = total_loss + mod_loss
+                num_present += 1
+
+            if num_present > 0:
+                total_loss = total_loss / num_present
+
+            # 9. Cross-modal mutual prediction loss
+            cross_modal_loss = torch.tensor(0.0, device=device)
+            if self.enable_cross_modal:
+                cross_modal_loss = self._compute_cross_modal_loss(
+                    selected_patches, perm_indices, observed,
+                    encoder_output, task_ranges,
+                )
+
+            return {
+                "loss": total_loss,
+                "cross_modal_loss": cross_modal_loss,
+                "per_modality_loss": per_mod_loss,
+                "mask_ratios": mask_ratios,
+            }
+
+        return {
+            "reconstructed": reconstructed,
+            "cross_modal_loss": torch.tensor(0.0, device=device),
+            "mask_ratios": mask_ratios,
+        }
+
+    def encode(
+        self,
+        images: torch.Tensor,
+        observed: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Encode without masking (for downstream use).
+        Returns encoder output tokens [B, num_global + 4*num_patches, embed_dim].
+        """
+        B = images.shape[0]
+        device = images.device
+        batch = self._split_modalities(images)
+
+        tokens_list = []
+        offset = self.num_global_tokens
+
+        for i, name in enumerate(self.MODALITY_NAMES):
+            img = batch[name]  # [B, 1, D, H, W]
+            patches = patchify(img, self.patch_size)  # [B, num_patches, 1, pd, ph, pw]
+
+            # Tokenize all patches (no masking)
+            tok = self.input_adapters[name](patches)  # [B, num_patches, embed_dim]
+
+            # Add positional embedding
+            pos_emb = self.pos_embed.expand(B, -1, -1)
+            tok = tok + pos_emb
+
+            # Zero out tokens for missing modalities
+            mod_mask = observed[:, i:i+1].unsqueeze(-1)  # [B, 1, 1]
+            tok = tok * mod_mask
+
+            tokens_list.append(tok)
+            offset += self.num_patches
+
+        input_tokens = torch.cat(tokens_list, dim=1)
+
+        # Add CLS token
+        if self.num_global_tokens > 0:
+            cls = self.global_tokens.unsqueeze(0).expand(B, -1, -1)
+            input_tokens = torch.cat([cls, input_tokens], dim=1)
+
+        # Build attention mask: prevent attending to tokens from missing modalities
+        total_tokens = input_tokens.shape[1]
+        attn_mask = torch.zeros(B, 1, 1, total_tokens, device=device)
+        mod_offset = self.num_global_tokens
+        for i, name in enumerate(self.MODALITY_NAMES):
+            start = mod_offset
+            end = mod_offset + self.num_patches
+            missing = (observed[:, i] < 0.5)  # [B]
+            if missing.any():
+                attn_mask[missing, :, :, start:end] = float("-inf")
+            mod_offset = end
+        if (attn_mask == 0).all():
+            attn_mask = None
+
+        # Encode with attention mask
+        encoder_output = input_tokens
+        for block in self.encoder:
+            encoder_output = block(encoder_output, attn_mask=attn_mask)
+
+        return encoder_output
+
+
+def create_multimae3d(
+    img_size: int = 128,
+    patch_size: int = 16,
+    embed_dim: int = 768,
+    depth: int = 12,
+    num_heads: int = 12,
+    decoder_embed_dim: int = 384,
+    decoder_depth: int = 2,
+    decoder_num_heads: int = 12,
+    mask_ratio: float = 0.75,
+    use_dirichlet: bool = True,
+    enable_cross_modal: bool = False,
+    **kwargs,
+) -> MultiMAE3D:
+    """Factory function to create MultiMAE3D with default ViT-B config."""
+    return MultiMAE3D(
+        img_size=img_size,
+        patch_size=patch_size,
+        embed_dim=embed_dim,
+        depth=depth,
+        num_heads=num_heads,
+        decoder_embed_dim=decoder_embed_dim,
+        decoder_depth=decoder_depth,
+        decoder_num_heads=decoder_num_heads,
+        mask_ratio=mask_ratio,
+        use_dirichlet=use_dirichlet,
+        enable_cross_modal=enable_cross_modal,
+        **kwargs,
+    )

BrainAnytime/models/multimae3d_utils.py

@@ -0,0 +1,344 @@
+"""
+MultiMAE 3D Utility Functions
+- Patchify / Unpatchify
+- Patch shuffling for masking
+- 3D sinusoidal positional embeddings
+- Dirichlet masking with missing modality support
+"""
+
+from typing import Union, Tuple, Dict, List
+
+import torch
+import torch.nn as nn
+from torch.distributions import Dirichlet
+from einops import rearrange
+
+
+def to_3tuple(x):
+    if isinstance(x, (list, tuple)):
+        assert len(x) == 3
+        return tuple(x)
+    return (x, x, x)
+
+
+def calc_patchified_dim(
+    img_size: Union[int, Tuple[int, int, int]],
+    patch_size: Union[int, Tuple[int, int, int]],
+) -> Tuple[int, int, int]:
+    img_size = to_3tuple(img_size)
+    patch_size = to_3tuple(patch_size)
+    return tuple(img_size[i] // patch_size[i] for i in range(3))
+
+
+def patchify(
+    image: torch.Tensor,
+    patch_size: Union[int, Tuple[int, int, int]],
+) -> torch.Tensor:
+    """
+    Convert image to patches.
+    image: [B, C, D, H, W]
+    returns: [B, num_patches, C, pd, ph, pw]
+    """
+    patch_size = to_3tuple(patch_size)
+    img_size = image.shape[-3:]
+    patchified_dim = calc_patchified_dim(img_size, patch_size)
+    patches = rearrange(
+        image,
+        "b c (nd pd) (nh ph) (nw pw) -> b (nd nh nw) c pd ph pw",
+        pd=patch_size[0],
+        ph=patch_size[1],
+        pw=patch_size[2],
+        nd=patchified_dim[0],
+        nh=patchified_dim[1],
+        nw=patchified_dim[2],
+    )
+    return patches
+
+
+def unpatchify(
+    patches: torch.Tensor,
+    img_size: Union[int, Tuple[int, int, int]],
+    patch_size: Union[int, Tuple[int, int, int]],
+) -> torch.Tensor:
+    """
+    Convert patches back to image.
+    patches: [B, num_patches, C, pd, ph, pw]
+    returns: [B, C, D, H, W]
+    """
+    patch_size = to_3tuple(patch_size)
+    img_size = to_3tuple(img_size)
+    patchified_dim = calc_patchified_dim(img_size, patch_size)
+    image = rearrange(
+        patches,
+        "b (nd nh nw) c pd ph pw -> b c (nd pd) (nh ph) (nw pw)",
+        pd=patch_size[0],
+        ph=patch_size[1],
+        pw=patch_size[2],
+        nd=patchified_dim[0],
+        nh=patchified_dim[1],
+        nw=patchified_dim[2],
+    )
+    return image
+
+
+def shuffle_patches(
+    patches: torch.Tensor,
+    permutations: torch.Tensor = None,
+    mask_probs: torch.Tensor = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Shuffle patches along the patch dimension.
+
+    When mask_probs is None: uniform random shuffle.
+    When mask_probs is provided: Gumbel-top-k weighted shuffle.
+      Patches with higher mask_probs end up at higher indices (masked),
+      patches with lower mask_probs end up at lower indices (visible).
+
+    Args:
+        patches: [B, N, ...]
+        permutations: optional pre-computed permutation indices [B, N]
+        mask_probs: optional [N] per-patch masking probability (sums to 1)
+
+    Returns:
+        (shuffled_patches, perm_indices)
+    """
+    batch_size, num_patches = patches.shape[0], patches.shape[1]
+    if permutations is not None:
+        perm_idx = permutations
+    else:
+        rand = torch.rand(batch_size, num_patches, device=patches.device)
+
+        if mask_probs is not None:
+            # Gumbel-top-k trick for weighted sampling without replacement.
+            # key_i = log(p_i) + Gumbel(0,1)_i
+            # Top-k of keys = sample from Multinomial(p, k)
+            # After ascending argsort: low keys → visible, high keys → masked.
+            mask_probs = mask_probs.to(patches.device)
+            gumbel = -torch.log(-torch.log(rand.clamp(1e-20, 1.0 - 1e-20)))
+            log_probs = torch.log(mask_probs.clamp(min=1e-20))  # [N]
+            keys = gumbel + log_probs.unsqueeze(0)  # [B, N]
+            perm_idx = torch.argsort(keys, dim=1)
+        else:
+            perm_idx = torch.argsort(rand, dim=1)
+
+    shuffled = patches[torch.arange(batch_size, device=patches.device)[:, None], perm_idx]
+    return shuffled, perm_idx
+
+
+def unshuffle_patches(
+    patches: torch.Tensor,
+    perm_idx: torch.Tensor,
+) -> torch.Tensor:
+    """
+    Inverse of shuffle_patches.
+    """
+    batch_size = patches.shape[0]
+    inv_idx = torch.argsort(perm_idx, dim=1)
+    return patches[torch.arange(batch_size, device=patches.device)[:, None], inv_idx]
+
+
+def build_3d_sincos_position_embedding(
+    grid_size: Tuple[int, int, int],
+    embed_dim: int,
+    temperature: float = 10000.0,
+) -> nn.Parameter:
+    """
+    Build 3D sinusoidal positional embedding.
+    returns: [1, num_patches, embed_dim] (frozen parameter)
+    """
+    grid_size = to_3tuple(grid_size)
+    h, w, d = grid_size
+
+    assert embed_dim % 6 == 0, \
+        f"embed_dim ({embed_dim}) must be divisible by 6 for 3D sincos pos embed"
+
+    pos_dim = embed_dim // 6
+    omega = torch.arange(pos_dim, dtype=torch.float32) / pos_dim
+    omega = 1.0 / (temperature ** omega)
+
+    grid_h = torch.arange(h, dtype=torch.float32)
+    grid_w = torch.arange(w, dtype=torch.float32)
+    grid_d = torch.arange(d, dtype=torch.float32)
+
+    out_h = torch.einsum("m,d->md", grid_h.flatten(), omega)
+    out_w = torch.einsum("m,d->md", grid_w.flatten(), omega)
+    out_d = torch.einsum("m,d->md", grid_d.flatten(), omega)
+
+    # Expand to full grid: [H*W*D, pos_dim] for each axis
+    # Use meshgrid ordering to get correct spatial layout
+    grid_h_idx, grid_w_idx, grid_d_idx = torch.meshgrid(
+        torch.arange(h), torch.arange(w), torch.arange(d), indexing="ij"
+    )
+    grid_h_flat = grid_h_idx.flatten()  # [H*W*D]
+    grid_w_flat = grid_w_idx.flatten()
+    grid_d_flat = grid_d_idx.flatten()
+
+    pos_emb = torch.cat([
+        torch.sin(out_h[grid_h_flat]),
+        torch.cos(out_h[grid_h_flat]),
+        torch.sin(out_w[grid_w_flat]),
+        torch.cos(out_w[grid_w_flat]),
+        torch.sin(out_d[grid_d_flat]),
+        torch.cos(out_d[grid_d_flat]),
+    ], dim=1)[None, :, :]  # [1, num_patches, embed_dim]
+
+    pos_emb = nn.Parameter(pos_emb)
+    pos_emb.requires_grad = False
+    return pos_emb
+
+
+def generate_dirichlet_mask_ratios(
+    num_modalities: int,
+    alpha: float,
+    overall_mask_ratio: float,
+) -> torch.Tensor:
+    """
+    Sample per-modality mask ratios from a Dirichlet distribution.
+    The total visible budget is distributed among modalities.
+
+    Returns: [num_modalities] tensor of per-modality mask ratios
+    """
+    dirichlet = Dirichlet(torch.tensor([float(alpha)] * num_modalities))
+    visible_ratio = 1.0 - overall_mask_ratio
+    total_visible = visible_ratio * num_modalities
+    visible_per_mod = total_visible * dirichlet.sample()
+    # Clamp to [0, 1]
+    mask_ratios = (1.0 - visible_per_mod).clamp(0.0, 1.0)
+    return mask_ratios
+
+
+def compute_mask_ratios(
+    modality_names: List[str],
+    observed: torch.Tensor,
+    mask_ratio: float = 0.75,
+    use_dirichlet: bool = True,
+    dirichlet_alpha: float = 1.0,
+) -> Dict[str, float]:
+    """
+    Compute per-modality mask ratios, respecting observed mask.
+    Missing modalities (observed=0) get mask_ratio=1.0.
+    Present modalities get Dirichlet or uniform masking.
+
+    Args:
+        modality_names: list of modality names, e.g. ['T1', 'T2', 'Flair', 'PET']
+        observed: [M] bool/float tensor, 1.0=present, 0.0=missing
+            NOTE: This is per-sample, called once per sample in the batch.
+            For simplicity, we use the same mask ratio for the whole batch
+            (based on which modalities are present in the majority of the batch).
+        mask_ratio: overall target mask ratio for present modalities
+        use_dirichlet: whether to use Dirichlet distribution
+        dirichlet_alpha: Dirichlet concentration parameter
+
+    Returns:
+        dict mapping modality_name -> mask_ratio (float)
+    """
+    ratios = {}
+    present_mods = [name for i, name in enumerate(modality_names) if observed[i] > 0.5]
+    missing_mods = [name for i, name in enumerate(modality_names) if observed[i] <= 0.5]
+
+    # Missing modalities: fully masked
+    for name in missing_mods:
+        ratios[name] = 1.0
+
+    # Present modalities: Dirichlet or uniform
+    if len(present_mods) > 0:
+        if use_dirichlet and len(present_mods) > 1:
+            # Dirichlet masking among present modalities
+            dir_ratios = generate_dirichlet_mask_ratios(
+                num_modalities=len(present_mods),
+                alpha=dirichlet_alpha,
+                overall_mask_ratio=mask_ratio,
+            )
+            for i, name in enumerate(present_mods):
+                ratios[name] = dir_ratios[i].item()
+        else:
+            # Uniform masking
+            for name in present_mods:
+                ratios[name] = mask_ratio
+
+    return ratios
+
+
+def mask_data(
+    batch: Dict[str, torch.Tensor],
+    modality_names: List[str],
+    observed: torch.Tensor,
+    mask_ratio: float = 0.75,
+    patch_size: Union[int, Tuple[int, int, int]] = 16,
+    use_dirichlet: bool = True,
+    dirichlet_alpha: float = 1.0,
+    patch_mask_probs: torch.Tensor = None,
+) -> Tuple[
+    Dict[str, torch.Tensor],
+    Dict[str, torch.Tensor],
+    Dict[str, torch.Tensor],
+    Dict[str, float],
+]:
+    """
+    Core masking function for MultiMAE pretraining.
+
+    For each modality:
+    - Patchify the image
+    - Shuffle patches (optionally weighted by anatomy importance)
+    - Split into selected (visible) and masked based on mask_ratio
+    - Missing modalities (observed=0) get 100% masking
+
+    Args:
+        batch: dict mapping modality name -> [B, 1, D, H, W] tensor
+        modality_names: ordered list of modality names
+        observed: [B, M] tensor indicating which modalities are present
+        mask_ratio: target mask ratio for present modalities
+        patch_size: patch size for patchification
+        use_dirichlet: whether to use Dirichlet distribution
+        dirichlet_alpha: Dirichlet concentration parameter
+        patch_mask_probs: optional [N_patches] per-patch masking probability
+            from anatomy-aware masking. When provided, uses Gumbel-top-k
+            weighted sampling instead of uniform random shuffling.
+            Higher probability = more likely to be masked.
+
+    Returns:
+        selected_patches: {modality: [B, num_selected, C, pd, ph, pw]}
+        masked_patches: {modality: [B, num_masked, C, pd, ph, pw]}
+        perm_indices: {modality: [B, num_patches]}
+        mask_ratios: {modality: float}
+    """
+    patch_size = to_3tuple(patch_size)
+    batch_size = observed.shape[0]
+
+    # Union strategy: if ANY sample in the batch has a modality, it gets
+    # partial masking. Samples where this modality is missing contribute
+    # zero-valued patches (harmless in encoder, excluded from loss).
+    # This ensures no information is wasted when modalities are present
+    # in a minority of samples.
+    batch_observed = (observed.max(dim=0).values > 0.5).float()  # [M]
+    mask_ratios = compute_mask_ratios(
+        modality_names=modality_names,
+        observed=batch_observed,
+        mask_ratio=mask_ratio,
+        use_dirichlet=use_dirichlet,
+        dirichlet_alpha=dirichlet_alpha,
+    )
+
+    selected_patches = {}
+    masked_patches = {}
+    perm_indices = {}
+
+    for mod_name in modality_names:
+        # Patchify: [B, 1, D, H, W] -> [B, num_patches, 1, pd, ph, pw]
+        patches = patchify(batch[mod_name], patch_size)
+        num_patches = patches.shape[1]
+
+        # Shuffle patches (weighted by anatomy importance if provided)
+        shuffled, perm_idx = shuffle_patches(patches, mask_probs=patch_mask_probs)
+        perm_indices[mod_name] = perm_idx
+
+        # Split into selected and masked
+        mod_mask_ratio = mask_ratios[mod_name]
+        num_selected = int((1.0 - mod_mask_ratio) * num_patches)
+        # Ensure at least 0 selected (for fully masked modalities)
+        num_selected = max(0, num_selected)
+
+        selected_patches[mod_name] = shuffled[:, :num_selected]
+        masked_patches[mod_name] = shuffled[:, num_selected:]
+
+    return selected_patches, masked_patches, perm_indices, mask_ratios

BrainAnytime/pretrain_dataloader_v2.py

@@ -0,0 +1,313 @@
+import os
+import numpy as np
+import pandas as pd
+import nibabel as nib
+import torch
+from torch.utils.data import Dataset, DataLoader
+import torchio as tio
+from typing import List, Dict, Tuple, Optional
+import random
+
+
+class MultiModalPretrainDataset(Dataset):
+    """
+    多模态3D医学图像预训练数据集
+    
+    特点:
+    - 支持多个数据集（A4, ADNIDOD, AIBL, BraTS, NACC）
+    - 缺失模态填充为0，并提供observed_indicator
+    - 支持数据增强（Spatial transforms）
+    - 支持Modality Dropout增加组合多样性
+    """
+    
+    # 统一的模态顺序
+    MODALITY_ORDER = ['T1', 'T2', 'Flair', 'PET']  # 统一为4个模态
+    
+    # 每个数据集的模态列名映射到统一名称
+    MODALITY_MAPPING = {
+        'modality_data_A4.xlsx': {'T1': 'T1', 'T2': 'T2', 'Flair': 'Flair', 'Amy_PET': 'PET'},
+        'modality_data_ADNIDOD.xlsx': {'T1': 'T1', 'T2': 'T2', 'Flair': 'Flair', 'PET': 'PET'},
+        'modality_data_AIBL.xlsx': {'T1': 'T1', 'T2': 'T2', 'Flair': 'Flair', 'PET': 'PET'},
+        'modality_data_BraTS.xlsx': {'T1w': 'T1', 'T2w': 'T2', 'Flair': 'Flair', 'PET': 'PET'},
+        'modality_data_NACC.xlsx': {'T1': 'T1', 'T2': 'T2', 'Flair': 'Flair', 'Amyloid': 'PET'},
+    }
+    
+    # Path prefix replacement: Excel paths use the old server prefix,
+    # remap to the local data directory.
+    OLD_PATH_PREFIX = "/home/data/Pretrain"
+    NEW_PATH_PREFIX = "./data/Pretrain"
+
+    def __init__(
+        self,
+        excel_dir: str = "./data/Match_data_path/pretraining_processed",
+        image_size: Tuple[int, int, int] = (128, 128, 128),
+        augmentation: bool = True,
+        modality_dropout_prob: float = 0.3,
+        min_modalities: int = 1,
+        cache_data: bool = False,
+    ):
+        """
+        Args:
+            excel_dir: Excel文件目录路径
+            image_size: 图像尺寸 (D, H, W)
+            augmentation: 是否进行数据增强
+            modality_dropout_prob: 每个模态被dropout的概率
+            min_modalities: 至少保留的模态数量
+            cache_data: 是否缓存加载的数据到内存
+        """
+        self.excel_dir = excel_dir
+        self.image_size = image_size
+        self.augmentation = augmentation
+        self.modality_dropout_prob = modality_dropout_prob
+        self.min_modalities = min_modalities
+        self.cache_data = cache_data
+        self.cache = {}
+        
+        # 加载所有样本
+        self.samples = self._load_all_samples()
+        print(f"Loaded {len(self.samples)} samples from {len(self.MODALITY_MAPPING)} datasets")
+        
+        # 初始化数据增强
+        if self.augmentation:
+            self.spatial_transform = tio.OneOf({
+                tio.RandomFlip(axes=0, flip_probability=0.5): 0.33,
+                tio.RandomAffine(scales=(0.9, 1.2), degrees=10, p=0.5): 0.33,
+                tio.RandomElasticDeformation(
+                    num_control_points=(10, 10, 10),
+                    max_displacement=8,
+                    locked_borders=2,
+                    p=0.5
+                ): 0.34,
+            })
+        
+    def _load_all_samples(self) -> List[Dict]:
+        """加载所有Excel文件中的样本"""
+        samples = []
+        
+        for excel_file, modality_map in self.MODALITY_MAPPING.items():
+            excel_path = os.path.join(self.excel_dir, excel_file)
+            if not os.path.exists(excel_path):
+                print(f"Warning: Excel file not found: {excel_path}")
+                continue
+                
+            df = pd.read_excel(excel_path)
+            dataset_name = excel_file.replace('modality_data_', '').replace('.xlsx', '')
+            
+            for idx, row in df.iterrows():
+                sample = {
+                    'dataset': dataset_name,
+                    'subject_id': row.get('SubjectID', f'{dataset_name}_{idx}'),
+                    'modalities': {}
+                }
+                
+                # 映射模态路径
+                for orig_col, unified_name in modality_map.items():
+                    if orig_col in df.columns:
+                        path = row[orig_col]
+                        if pd.notna(path) and isinstance(path, str):
+                            # Remap old server path prefix to local path
+                            if path.startswith(self.OLD_PATH_PREFIX):
+                                path = self.NEW_PATH_PREFIX + path[len(self.OLD_PATH_PREFIX):]
+                            if os.path.exists(path):
+                                sample['modalities'][unified_name] = path
+                
+                # 只添加至少有一个模态的样本
+                if len(sample['modalities']) >= 1:
+                    samples.append(sample)
+        
+        return samples
+    
+    def _load_nifti(self, path: str) -> np.ndarray:
+        """加载NIfTI文件"""
+        try:
+            nii = nib.load(path)
+            data = nii.get_fdata().astype(np.float32)
+            return data
+        except Exception as e:
+            print(f"Error loading {path}: {e}")
+            return None
+    
+    def _apply_modality_dropout(self, available_modalities: List[str]) -> List[str]:
+        """
+        应用Modality Dropout
+        随机丢弃一些模态以增加组合多样性
+        """
+        if len(available_modalities) <= self.min_modalities:
+            return available_modalities
+        
+        kept_modalities = []
+        for mod in available_modalities:
+            if random.random() > self.modality_dropout_prob:
+                kept_modalities.append(mod)
+        
+        # 确保至少保留min_modalities个模态
+        if len(kept_modalities) < self.min_modalities:
+            # 随机选择需要保留的模态
+            kept_modalities = random.sample(available_modalities, self.min_modalities)
+        
+        return kept_modalities
+    
+    def __len__(self) -> int:
+        return len(self.samples)
+    
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        sample = self.samples[idx]
+        
+        # 检查缓存
+        if self.cache_data and idx in self.cache:
+            cached_data = self.cache[idx]
+            images = cached_data['images'].clone()
+            original_observed = cached_data['observed'].clone()
+        else:
+            # 初始化输出张量
+            num_modalities = len(self.MODALITY_ORDER)
+            images = torch.zeros(num_modalities, *self.image_size, dtype=torch.float32)
+            original_observed = torch.zeros(num_modalities, dtype=torch.float32)
+            
+            # 加载每个模态
+            for i, modality in enumerate(self.MODALITY_ORDER):
+                if modality in sample['modalities']:
+                    path = sample['modalities'][modality]
+                    data = self._load_nifti(path)
+                    
+                    if data is not None:
+                        # 确保数据尺寸正确
+                        if data.shape == self.image_size:
+                            images[i] = torch.from_numpy(data)
+                            original_observed[i] = 1.0
+                        else:
+                            print(f"Warning: Size mismatch for {path}, expected {self.image_size}, got {data.shape}")
+            
+            # 缓存数据
+            if self.cache_data:
+                self.cache[idx] = {
+                    'images': images.clone(),
+                    'observed': original_observed.clone()
+                }
+        
+        # 不再应用Modality Dropout，直接使用原始observed
+        observed = original_observed.clone()
+        
+        # 应用空间数据增强
+        if self.augmentation:
+            # 只对observed的模态应用增强
+            # 创建TorchIO Subject
+            subject_dict = {}
+            for i, modality in enumerate(self.MODALITY_ORDER):
+                if observed[i] == 1.0:
+                    # TorchIO需要4D张量 (C, D, H, W)
+                    subject_dict[modality] = tio.ScalarImage(tensor=images[i:i+1])
+            
+            if subject_dict:
+                subject = tio.Subject(**subject_dict)
+                transformed = self.spatial_transform(subject)
+                
+                # 将增强后的数据放回images张量
+                for i, modality in enumerate(self.MODALITY_ORDER):
+                    if modality in subject_dict:
+                        images[i] = transformed[modality].data[0]
+        
+        return {
+            'images': images,  # (num_modalities, D, H, W)
+            'observed': observed,  # (num_modalities,)
+
+        }
+
+
+def create_pretrain_dataloader(
+    excel_dir: str = "/home/data/Match_data_path/pretraining_processed",
+    batch_size: int = 4,
+    num_workers: int = 8,
+    augmentation: bool = True,
+    modality_dropout_prob: float = 0.3,
+    min_modalities: int = 1,
+    shuffle: bool = True,
+    pin_memory: bool = True,
+    cache_data: bool = False,
+) -> DataLoader:
+    """
+    创建预训练数据加载器
+    
+    Args:
+        excel_dir: Excel文件目录
+        batch_size: 批量大小
+        num_workers: 数据加载进程数
+        augmentation: 是否数据增强
+        modality_dropout_prob: 模态dropout概率
+        min_modalities: 至少保留的模态数
+        shuffle: 是否打乱数据
+        pin_memory: 是否使用pinned memory
+        cache_data: 是否缓存数据到内存
+    
+    Returns:
+        DataLoader实例
+    """
+    dataset = MultiModalPretrainDataset(
+        excel_dir=excel_dir,
+        augmentation=augmentation,
+        modality_dropout_prob=modality_dropout_prob,
+        min_modalities=min_modalities,
+        cache_data=cache_data,
+    )
+    
+    dataloader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        num_workers=num_workers,
+        pin_memory=pin_memory,
+        drop_last=True,
+    )
+    
+    return dataloader
+
+
+def collate_fn_with_info(batch: List[Dict]) -> Dict[str, torch.Tensor]:
+    """
+    自定义collate函数，处理批量数据
+    """
+    images = torch.stack([item['images'] for item in batch])
+    observed = torch.stack([item['observed'] for item in batch])
+
+    return {
+        'images': images,  # (B, num_modalities, D, H, W)
+        'observed': observed,  # (B, num_modalities)
+    }
+
+
+# ============== 使用示例 ==============
+if __name__ == '__main__':
+    print("=" * 60)
+    print("多模态3D医学图像预训练数据加载器")
+    print("=" * 60)
+    
+    # 创建数据加载器
+    dataloader = create_pretrain_dataloader(
+        excel_dir="/home/data/Match_data_path/pretraining_processed",
+        batch_size=2,
+        num_workers=4,
+        augmentation=True,
+        modality_dropout_prob=0.3,
+        min_modalities=1,
+        shuffle=True,
+    )
+    
+    print(f"\n数据集大小: {len(dataloader.dataset)}")
+    print(f"批量数: {len(dataloader)}")
+    print(f"模态顺序: {MultiModalPretrainDataset.MODALITY_ORDER}")
+    
+    # 测试加载一个批量
+    print("\n测试加载一个批量...")
+    for batch in dataloader:
+        images = batch['images']
+        observed = batch['observed']
+        
+        print(f"\n批量数据形状:")
+        print(f"  images: {images.shape}")  # (B, 4, 128, 128, 128)
+        print(f"  observed: {observed.shape}")  # (B, 4)
+        
+    
+    print("\n" + "=" * 60)
+    print("数据加载测试完成!")
+    print("=" * 60)
+

BrainAnytime/test_main.py

@@ -0,0 +1,355 @@
+#!/usr/bin/env python
+"""
+MultiMAE3D Test-Only Evaluation
+
+Load saved finetuned checkpoints and evaluate on the test set.
+Reuses model/data/metric utilities from finetune_main.py.
+
+Usage:
+    # Test all tasks for finetune mode
+    python test_main.py --mode finetune
+
+    # Test a specific task
+    python test_main.py --mode finetune --tasks "CN vs AD"
+
+    # Test with custom checkpoint directory
+    python test_main.py --mode finetune --checkpoint_dir ./saves/multimae_finetune/
+"""
+
+import os
+import sys
+import gc
+import random
+import warnings
+from collections import defaultdict
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+from tqdm import tqdm
+from scipy.stats import pearsonr
+
+warnings.filterwarnings("ignore")
+
+_BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, _BASE_DIR)
+
+from models.multimae3d import create_multimae3d, MultiMAE3D
+from downstream_dataloader import create_downstream_dataloader
+from finetune_main import (
+    seed_everything,
+    str2bool,
+    MultiMAE3DForDownstream,
+    run_epoch,
+    calc_classification_metrics,
+    calc_regression_metrics,
+    calc_metrics_by_combo,
+    _save_combo_results,
+)
+
+
+# =========================================================================
+# Test-only evaluation for a single (task, seed)
+# =========================================================================
+
+def test_evaluate(args, task_type, seed, device, checkpoint_path):
+    """Load a saved checkpoint and evaluate on test set."""
+    seed_everything(seed)
+    torch.cuda.empty_cache()
+
+    is_cls = task_type in ('CN vs AD', 'CN vs MCI')
+
+    # ---- Test data loader ----
+    loader_kwargs = dict(
+        batch_size=args.batch_size,
+        num_workers=args.num_workers,
+        pin_memory=True,
+        cache_data=False,
+        image_size=tuple(args.image_size),
+        base_dir=args.base_dir,
+        modalities=args.modalities,
+        intersection=args.intersection,
+    )
+
+    print(f"\nLoading test data for task={task_type}, seed={seed}")
+    test_loader = create_downstream_dataloader(
+        excel_path=args.test_excel, labels=[task_type],
+        augmentation=False, shuffle=False,
+        phase='test', modality_dropout=False, expand_val_combinations=False,
+        exclusive_modalities=False,
+        **loader_kwargs,
+    )
+    print(f"  Test: {len(test_loader.dataset)} samples")
+
+    # ---- Model ----
+    encoder = create_multimae3d(
+        img_size=args.img_size,
+        patch_size=args.patch_size,
+        embed_dim=args.embed_dim,
+        depth=args.depth,
+        num_heads=args.num_heads,
+        decoder_embed_dim=args.decoder_embed_dim,
+        decoder_depth=args.decoder_depth,
+        decoder_num_heads=args.decoder_num_heads,
+    )
+
+    model = MultiMAE3DForDownstream(
+        encoder=encoder,
+        embed_dim=args.embed_dim,
+        num_outputs=1,
+        pool=args.pool,
+        dropout=args.dropout,
+    ).to(device)
+
+    # Load checkpoint
+    print(f"  Loading checkpoint: {checkpoint_path}")
+    ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    model.load_state_dict(ckpt['model_state_dict'])
+    print(f"  Loaded (epoch={ckpt.get('epoch', '?')}, "
+          f"best_metric={ckpt.get('best_metric', '?')})")
+
+    # Criterion
+    criterion = (nn.BCEWithLogitsLoss() if is_cls
+                 else nn.MSELoss()).to(device)
+
+    # ---- Test evaluation ----
+    print("  Evaluating on test set...")
+    test_loss, test_preds, test_labels, test_probs, test_combos = run_epoch(
+        test_loader, model, criterion, device, task_type,
+        is_training=False,
+    )
+
+    # Overall test metrics
+    if is_cls:
+        test_m = calc_classification_metrics(
+            test_preds, test_labels, test_probs)
+        print(
+            f"  Test: Acc={test_m['acc']*100:.2f}%, "
+            f"AUC={test_m['auc']*100:.2f}%, "
+            f"Sen={test_m['sensitivity']*100:.2f}%, "
+            f"Spe={test_m['specificity']*100:.2f}%, "
+            f"F1={test_m['f1']*100:.2f}%"
+        )
+    else:
+        test_m = calc_regression_metrics(test_preds, test_labels)
+        print(
+            f"  Test: MAE={test_m['mae']:.4f}, "
+            f"RMSE={test_m['rmse']:.4f}, "
+            f"Pearson={test_m['pearson']:.4f}"
+        )
+
+    # Per-modality-combination breakdown
+    combo_results = calc_metrics_by_combo(
+        test_preds, test_labels, test_probs, test_combos, task_type)
+    if combo_results:
+        print(f"\n  Per-modality-combination results:")
+        for combo in sorted(combo_results.keys()):
+            r = combo_results[combo]
+            n = r['n_samples']
+            if is_cls:
+                print(f"    {combo:25s} (n={n:3d}) | "
+                      f"Acc={r['acc']*100:.1f}%, AUC={r['auc']*100:.1f}%")
+            else:
+                print(f"    {combo:25s} (n={n:3d}) | "
+                      f"MAE={r['mae']:.4f}, Pearson={r['pearson']:.4f}")
+
+        # Save per-combo results
+        mode_tag = args.mode
+        _save_combo_results(combo_results, task_type, seed,
+                            f"test_{mode_tag}", is_cls)
+
+    # Cleanup
+    del model, encoder, test_loader
+    torch.cuda.empty_cache()
+    gc.collect()
+
+    return test_m
+
+
+# =========================================================================
+# Argument parsing
+# =========================================================================
+
+def parse_args():
+    import argparse
+    p = argparse.ArgumentParser(
+        description='MultiMAE3D Test-Only Evaluation')
+
+    # Mode & checkpoints
+    p.add_argument('--mode', type=str, default='finetune',
+                   choices=['finetune', 'freeze_then_finetune'],
+                   help='Which training mode checkpoints to load')
+    p.add_argument('--checkpoint_dir', type=str, default=None,
+                   help='Directory containing saved checkpoints. '
+                        'Defaults to saves/multimae_{mode}/')
+
+    # Tasks & seeds
+    p.add_argument('--tasks', type=str, nargs='+',
+                   default=['CN vs AD', 'CN vs MCI', 'MMSE', 'AGE'],
+                   help='Tasks to evaluate')
+    p.add_argument('--n_seeds', type=int, default=3,
+                   help='Number of random seeds per task')
+
+    # Data
+    p.add_argument('--test_excel', type=str,
+                   default='./data/Downstream/'
+                           'ADNI_Division/modality_data_test.xlsx')
+    p.add_argument('--base_dir', type=str,
+                   default='./data/Downstream/ADNI/')
+    p.add_argument('--modalities', type=str, nargs='+',
+                   default=['T1', 'T2', 'Flair', 'PET'])
+    p.add_argument('--intersection', type=str2bool, default=False)
+    p.add_argument('--image_size', type=int, nargs=3,
+                   default=[128, 128, 128])
+    p.add_argument('--batch_size', type=int, default=4)
+    p.add_argument('--num_workers', type=int, default=8)
+
+    # MultiMAE encoder architecture (must match checkpoint)
+    p.add_argument('--img_size', type=int, default=128)
+    p.add_argument('--patch_size', type=int, default=16)
+    p.add_argument('--embed_dim', type=int, default=768)
+    p.add_argument('--depth', type=int, default=12)
+    p.add_argument('--num_heads', type=int, default=12)
+    p.add_argument('--decoder_embed_dim', type=int, default=384)
+    p.add_argument('--decoder_depth', type=int, default=2)
+    p.add_argument('--decoder_num_heads', type=int, default=12)
+
+    # Downstream head
+    p.add_argument('--pool', type=str, default='cls',
+                   choices=['cls', 'mean'])
+    p.add_argument('--dropout', type=float, default=0.1)
+
+    # Device
+    p.add_argument('--device', type=int, default=0)
+
+    return p.parse_args()
+
+
+# =========================================================================
+# Main
+# =========================================================================
+
+def main():
+    args = parse_args()
+    device = torch.device(
+        f'cuda:{args.device}' if torch.cuda.is_available() else 'cpu')
+
+    # Resolve checkpoint directory
+    if args.checkpoint_dir is None:
+        args.checkpoint_dir = os.path.join(
+            _BASE_DIR, 'saves', f'multimae_{args.mode}')
+
+    print("=" * 80)
+    print(f"MultiMAE3D Test-Only Evaluation")
+    print(f"  Mode           : {args.mode}")
+    print(f"  Tasks          : {args.tasks}")
+    print(f"  Seeds          : {args.n_seeds}")
+    print(f"  Checkpoint dir : {args.checkpoint_dir}")
+    print(f"  Test data      : {args.test_excel}")
+    print(f"  Pool           : {args.pool}")
+    print(f"  Device         : {device}")
+    print("=" * 80)
+
+    all_results = {}
+
+    for task_type in args.tasks:
+        print(f"\n{'='*80}")
+        print(f"TASK: {task_type}")
+        print(f"{'='*80}")
+
+        is_cls = task_type in ('CN vs AD', 'CN vs MCI')
+        seed_results = []
+        task_str = task_type.replace(' ', '_')
+
+        for seed in range(args.n_seeds):
+            ckpt_name = f'{task_str}_seed_{seed}_best.pth'
+            ckpt_path = os.path.join(args.checkpoint_dir, ckpt_name)
+
+            if not os.path.isfile(ckpt_path):
+                print(f"\n--- Seed {seed} --- SKIPPED (checkpoint not found: {ckpt_name})")
+                continue
+
+            print(f"\n--- Seed {seed} ---")
+            metrics = test_evaluate(args, task_type, seed, device, ckpt_path)
+            seed_results.append(metrics)
+
+        if not seed_results:
+            print(f"  No checkpoints found for {task_type}, skipping.")
+            continue
+
+        all_results[task_type] = seed_results
+
+        # Per-task summary
+        n = len(seed_results)
+        print(f"\n{task_type} Summary ({n} seeds):")
+        if is_cls:
+            for key in ['acc', 'auc', 'sensitivity', 'specificity', 'f1']:
+                vals = [r[key] * 100 for r in seed_results]
+                print(f"  {key:>12s}: {np.mean(vals):.2f} +/- {np.std(vals):.2f}%")
+        else:
+            for key in ['mae', 'rmse', 'pearson']:
+                vals = [r[key] for r in seed_results]
+                print(f"  {key:>12s}: {np.mean(vals):.4f} +/- {np.std(vals):.4f}")
+
+    # ---- Final summary table ----
+    print("\n" + "=" * 80)
+    print("FINAL SUMMARY")
+    print("=" * 80)
+
+    summary_rows = []
+
+    for task_type in args.tasks:
+        if task_type not in all_results:
+            continue
+        results = all_results[task_type]
+        is_cls = task_type in ('CN vs AD', 'CN vs MCI')
+        row = {'Task': task_type, 'Mode': args.mode, 'N_seeds': len(results)}
+
+        if is_cls:
+            for key in ['acc', 'auc', 'sensitivity', 'specificity', 'f1']:
+                vals = [r[key] * 100 for r in results]
+                row[f'{key}_mean'] = np.mean(vals)
+                row[f'{key}_std'] = np.std(vals)
+                row[key] = f"{np.mean(vals):.2f}+/-{np.std(vals):.2f}"
+            for i, r in enumerate(results):
+                row[f'seed_{i}_acc'] = r['acc'] * 100
+                row[f'seed_{i}_auc'] = r['auc'] * 100
+
+            vals_acc = [r['acc'] * 100 for r in results]
+            vals_auc = [r['auc'] * 100 for r in results]
+            print(f"  {task_type:12s} | "
+                  f"Acc: {np.mean(vals_acc):.2f}+/-{np.std(vals_acc):.2f}% | "
+                  f"AUC: {np.mean(vals_auc):.2f}+/-{np.std(vals_auc):.2f}%")
+        else:
+            for key in ['mae', 'rmse', 'pearson']:
+                vals = [r[key] for r in results]
+                row[f'{key}_mean'] = np.mean(vals)
+                row[f'{key}_std'] = np.std(vals)
+                row[key] = f"{np.mean(vals):.4f}+/-{np.std(vals):.4f}"
+            for i, r in enumerate(results):
+                row[f'seed_{i}_mae'] = r['mae']
+                row[f'seed_{i}_pearson'] = r['pearson']
+
+            vals_mae = [r['mae'] for r in results]
+            vals_r = [r['pearson'] for r in results]
+            print(f"  {task_type:12s} | "
+                  f"MAE: {np.mean(vals_mae):.4f}+/-{np.std(vals_mae):.4f} | "
+                  f"Pearson: {np.mean(vals_r):.4f}+/-{np.std(vals_r):.4f}")
+
+        summary_rows.append(row)
+
+    # Save summary Excel
+    if summary_rows:
+        results_dir = os.path.join(_BASE_DIR, 'results')
+        os.makedirs(results_dir, exist_ok=True)
+        summary_path = os.path.join(
+            results_dir, f'multimae_test_{args.mode}_summary.xlsx')
+        pd.DataFrame(summary_rows).to_excel(summary_path, index=False)
+        print(f"\nSummary saved to: {summary_path}")
+
+    print("=" * 80)
+
+
+if __name__ == '__main__':
+    main()

README.md

@@ -4,7 +4,7 @@ emoji: 🧠
 colorFrom: purple
 colorTo: yellow
 sdk: gradio
-sdk_version: 6.14.0
+sdk_version: 4.44.0
 python_version: '3.10'
 app_file: app.py
 pinned: false
@@ -14,160 +14,47 @@ models:
   - Simmonstt/BrainAnytime
 ---
 
-# BrainAnytime
-
-Official implementation of **BrainAnytime: Anatomy-Aware Cross-Modal Pretraining for Brain Image Analysis with Arbitrary Modality Availability**.
-
-## Paper
+# BrainAnytime Demo
 
 **BrainAnytime: Anatomy-Aware Cross-Modal Pretraining for Brain Image Analysis with Arbitrary Modality Availability**
 
-- arXiv: [2605.13059](https://arxiv.org/abs/2605.13059)
-- PDF: [https://arxiv.org/pdf/2605.13059](https://arxiv.org/pdf/2605.13059)
-
-## Congrats: This paper has been early accepted (top 9%) by MICCAI 2026.
-
-## Pretrained Weights
-
-**Finetuned checkpoints are available on Hugging Face: [Simmonstt/BrainAnytime](https://huggingface.co/Simmonstt/BrainAnytime).**
-
-Pretrained weights are also available at [Google Drive](https://drive.google.com/file/d/1L49zJ_Apj2jJe88_iy6jLcmd6KUlnc5h/view?usp=sharing).
-
-## Overview
-
-BrainAnytime is a self-supervised pretraining framework for multi-modal 3D brain imaging (T1, T2, Flair, PET) that handles **arbitrary missing modality combinations** at both training and inference time.
-
-### Key Features
-
-- **Multi-modal Masked Autoencoder (MultiMAE3D)**: Shared ViT encoder with per-modality input/output adapters, supporting 4 modalities (T1, T2, Flair, PET)
-- **Cross-Modal Mutual Prediction**: EMA teacher-student framework for MRI-PET cross-level feature alignment
-- **Anatomy-Aware Adaptive Masking**: Three-phase curriculum masking guided by AAL116 brain atlas and AD-relevant region priors
-- **Missing Modality Robustness**: Handles arbitrary missing modality combinations via attention masking and observed indicators
+This Hugging Face Space provides an interactive demo for the BrainAnytime model, which supports:
 
-## Project Structure
+- **Multi-modal brain imaging**: T1, T2, Flair, PET
+- **Missing modality robustness**: Handles arbitrary missing modality combinations
+- **Multiple downstream tasks**: 
+  - CN vs AD classification
+  - CN vs MCI classification
+  - MMSE score regression
+  - Age prediction
 
-```
-BrainAnytime/
-├── models/
-│   ├── multimae3d.py              # MultiMAE3D model architecture
-│   └── multimae3d_utils.py        # Patchify, masking, positional embeddings
-├── anatomy_masking.py             # Anatomy-aware adaptive masking module
-├── pretrain_dataloader_v2.py      # Multi-modal pretraining data loader
-├── train_multimae.py              # Pretraining script (single/multi-GPU DDP)
-├── finetune_main.py               # Downstream finetuning (CN vs AD, CN vs MCI, MMSE, AGE)
-├── test_main.py                   # Test-only evaluation
-└── altas/
-    └── AAL116_standard.nii.gz     # AAL116 brain atlas (128x128x128)
-```
-
-## Requirements
-
-- Python >= 3.8
-- PyTorch >= 1.12
-- torchio
-- nibabel
-- timm
-- einops
-- tensorboardX
-- scikit-learn
-- pandas
-- scipy
-- tqdm
-
-## Data Preparation
-
-Organize your data as follows:
-
-```
-./data/
-├── Match_data_path/
-│   └── pretraining_processed/     # Pretraining Excel files
-│       ├── modality_data_A4.xlsx
-│       ├── modality_data_ADNIDOD.xlsx
-│       ├── modality_data_AIBL.xlsx
-│       ├── modality_data_BraTS.xlsx
-│       └── modality_data_NACC.xlsx
-├── Pretrain/                      # Preprocessed NIfTI files for pretraining
-└── Downstream/
-    └── ADNI/                      # Downstream task data
-        └── ADNI_Division/
-            ├── modality_data_train.xlsx
-            ├── modality_data_val.xlsx
-            └── modality_data_test.xlsx
-```
+## Features
 
-Each Excel file should contain columns for subject IDs and file paths to the corresponding NIfTI images for each modality.
+### Online Inference Demo
+Select a task and modality combination, then run inference on pre-selected ADNI samples.
 
-## Attention
+### Supported Modality Combinations
+- **T**: T1 only
+- **TF**: T1 + Flair
+- **TMF**: T1 + T2 + Flair
+- **TFP**: T1 + Flair + PET
+- **TMFP**: All modalities (Full)
 
-To save training time, we preprocess the 3D multimodal image by following skull stripping, MN152 template co-registration, min–max normalization, and resampling to 128 × 128 ×128 in advance. The data loader only contains data augmentation during training.
+## Links
 
-## Usage
-
-### Pretraining
-
-```bash
-# Single GPU
-python train_multimae.py --batch_size 4
-
-# Multi-GPU DDP (8 GPUs)
-torchrun --nproc_per_node=8 train_multimae.py \
-    --batch_size 16 \
-    --enable_cross_modal \
-    --use_anatomy_masking \
-    --atlas_path altas/AAL116_standard.nii.gz
-```
-
-### Downstream Finetuning
-
-```bash
-# Finetune on all tasks (3 seeds each)
-python finetune_main.py \
-    --pretrained ./pretrain_checkpoints/multimae/best_model.pth
-
-# Specific task only
-python finetune_main.py \
-    --pretrained ./pretrain_checkpoints/multimae/best_model.pth \
-    --tasks "CN vs AD"
-```
-
-### Testing
-
-```bash
-# Test all tasks for finetune mode
-python test_main.py --mode finetune
-
-# Test a specific task
-python test_main.py --mode finetune --tasks "CN vs AD"
-```
-
-## Downstream Tasks
-
-| Task | Type | Metric |
-|------|------|--------|
-| CN vs AD | Classification | ACC, AUC, Sensitivity, Specificity, F1 |
-| CN vs MCI | Classification | ACC, AUC, Sensitivity, Specificity, F1 |
-| MMSE | Regression | MAE, RMSE, Pearson |
-| AGE | Regression | MAE, RMSE, Pearson |
-
-## License
-
-This project is released for academic research purposes only.
+- **GitHub**: https://github.com/guangqianyang/BrainAnytime
+- **Model Weights**: https://huggingface.co/Simmonstt/BrainAnytime
+- **Paper**: [arXiv:2605.13059](https://arxiv.org/abs/2605.13059)
 
 ## Citation
 
-If you use BrainAnytime in your research, please cite:
-
 ```bibtex
-@misc{yang2026brainanytimeanatomyawarecrossmodalpretraining,
-      title={BrainAnytime: Anatomy-Aware Cross-Modal Pretraining for Brain Image Analysis with Arbitrary Modality Availability},
-      author={Guangqian Yang and Tong Ding and Wenlong Hou and Yue Xun and Ye Du and Qian Niu and Shujun Wang},
-      year={2026},
-      eprint={2605.13059},
-      archivePrefix={arXiv},
-      primaryClass={cs.CV},
-      url={https://arxiv.org/abs/2605.13059},
+@misc{yang2026brainanytime,
+  title={BrainAnytime: Anatomy-Aware Cross-Modal Pretraining for Brain Image Analysis with Arbitrary Modality Availability},
+  author={Yang, Guangqian and Ding, Tong and Hou, Wenlong and Xun, Yue and Du, Ye and Niu, Qian and Wang, Shujun},
+  year={2026},
+  eprint={2605.13059},
+  archivePrefix={arXiv},
+  primaryClass={cs.CV}
 }
 ```
-
-Paper page: https://arxiv.org/abs/2605.13059

app.py

@@ -1,55 +1,532 @@
+#!/usr/bin/env python3
+"""
+BrainAnytime Hugging Face Space Demo
+
+Interactive demo for brain image analysis with multi-modal support.
+Supports 4 tasks and 5 modality combinations.
+"""
+
+import os
+import sys
+import json
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
 import gradio as gr
+import numpy as np
+from PIL import Image
 
-MODEL_REPO = "Simmonstt/BrainAnytime"
-GITHUB_REPO = "https://github.com/guangqianyang/BrainAnytime"
-CHECKPOINTS = [
-    "CN_vs_AD_seed_0_best.pth",
-    "CN_vs_MCI_seed_0_best.pth",
-    "MMSE_seed_0_best.pth",
-    "AGE_seed_0_best.pth",
-]
+# 添加当前目录到路径（用于导入 inference_engine）
+BASE_DIR = Path(__file__).parent
+sys.path.insert(0, str(BASE_DIR))
 
-INTRO = """
-# BrainAnytime Demo
+# 导入推理引擎
+try:
+    from inference_engine import (
+        BrainAnytimeInference,
+        TASKS,
+        MODALITY_ORDER,
+        SHORT_TO_FULL,
+        format_result,
+    )
+    INFERENCE_AVAILABLE = True
+except ImportException as e:
+    print(f"Warning: Inference engine not available: {e}")
+    INFERENCE_AVAILABLE = False
 
-**BrainAnytime: Anatomy-Aware Cross-Modal Pretraining for Brain Image Analysis with Arbitrary Modality Availability**
 
-This Hugging Face Space hosts the official code from GitHub. Full 3D multi-modal inference
-requires preprocessed NIfTI volumes and GPU resources. Use the linked model repository for
-finetuned checkpoints and run `finetune_main.py` / `test_main.py` locally for evaluation.
-"""
+# =============================================================================
+# 配置
+# =============================================================================
 
+# 模态组合（5种实验配置）
+MODALITY_COMBOS = {
+    "T": {"name": "T (仅T1)", "modalities": ["T1"]},
+    "TF": {"name": "TF (T1+Flair)", "modalities": ["T1", "Flair"]},
+    "TMF": {"name": "TMF (T1+T2+Flair)", "modalities": ["T1", "T2", "Flair"]},
+    "TFP": {"name": "TFP (T1+Flair+PET)", "modalities": ["T1", "Flair", "PET"]},
+    "TMFP": {"name": "TMFP (全模态)", "modalities": ["T1", "T2", "Flair", "PET"]},
+}
 
-def show_project_info():
-    checkpoint_lines = "\n".join(f"- `{name}`" for name in CHECKPOINTS)
-    return f"""{INTRO}
+# 任务配置
+TASK_CONFIG = {
+    "CN_vs_AD": {
+        "name": "CN vs AD",
+        "type": "classification",
+        "description": "区分认知正常 (CN) 与阿尔茨海默病 (AD)",
+        "classes": ["CN", "AD"],
+    },
+    "CN_vs_MCI": {
+        "name": "CN vs MCI",
+        "type": "classification",
+        "description": "区分认知正常 (CN) 与轻度认知障碍 (MCI)",
+        "classes": ["CN", "MCI"],
+    },
+    "MMSE": {
+        "name": "MMSE Score",
+        "type": "regression",
+        "description": "预测 MMSE 认知评分 (10-30分)",
+        "unit": "points",
+        "range": [10, 30],
+    },
+    "AGE": {
+        "name": "Age Prediction",
+        "type": "regression",
+        "description": "预测年龄",
+        "unit": "years",
+        "range": [50, 100],
+    },
+}
 
-## Links
-- GitHub: {GITHUB_REPO}
-- Model weights: https://huggingface.co/{MODEL_REPO}
+# 样本数据路径
+SAMPLES_DIR = BASE_DIR / "demo_samples"
+SAMPLES_INDEX = SAMPLES_DIR / "samples.json"
 
-## Available finetuned checkpoints
-{checkpoint_lines}
 
-## Supported downstream tasks
-- CN vs AD (classification)
-- CN vs MCI (classification)
-- MMSE (regression)
-- AGE (regression)
+# =============================================================================
+# 全局状态
+# =============================================================================
 
-## Quick start (local)
-```bash
-git clone {GITHUB_REPO}.git
-cd BrainAnytime
-pip install -r requirements.txt
-python finetune_main.py --pretrained <path/to/pretrained.pth>
-```
-"""
+_inference_engine: Optional[BrainAnytimeInference] = None
+_samples_cache: Optional[Dict] = None
+
+
+def get_inference_engine() -> Optional[BrainAnytimeInference]:
+    """获取或创建推理引擎（单例）"""
+    global _inference_engine
+    
+    if _inference_engine is None and INFERENCE_AVAILABLE:
+        # 尝试从本地加载，否则从 HF Hub 下载
+        checkpoints_dir = os.environ.get(
+            "CHECKPOINTS_DIR",
+            "/home/23037125r/code/random/multimae_freeze_then_finetune/"
+        )
+        
+        if not os.path.exists(checkpoints_dir):
+            checkpoints_dir = None  # 将从 HF Hub 下载
+        
+        try:
+            _inference_engine = BrainAnytimeInference(
+                checkpoints_dir=checkpoints_dir
+            )
+        except Exception as e:
+            print(f"Failed to initialize inference engine: {e}")
+            return None
+    
+    return _inference_engine
+
+
+def load_samples_index() -> Optional[Dict]:
+    """加载样本索引"""
+    global _samples_cache
+    
+    if _samples_cache is not None:
+        return _samples_cache
+    
+    if not SAMPLES_INDEX.exists():
+        print(f"Warning: Samples index not found at {SAMPLES_INDEX}")
+        return None
+    
+    try:
+        with open(SAMPLES_INDEX, 'r') as f:
+            _samples_cache = json.load(f)
+        return _samples_cache
+    except Exception as e:
+        print(f"Error loading samples index: {e}")
+        return None
+
+
+def get_sample_info(task: str, combo: str) -> Optional[Dict]:
+    """获取指定任务和组合的样本信息"""
+    index = load_samples_index()
+    if not index:
+        return None
+    
+    for sample in index.get("samples", []):
+        if sample["task"] == task and sample["modality_combination"] == combo:
+            return sample
+    
+    return None
+
+
+def get_preview_images(sample_info: Dict) -> List[Tuple[str, str]]:
+    """获取样本的预览图路径和标题"""
+    if not sample_info:
+        return []
+    
+    images = []
+    sample_dir = SAMPLES_DIR / sample_info["task"] / sample_info["modality_combination"] / sample_info["sample_name"]
+    
+    for mod in MODALITY_ORDER:
+        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"))
+    
+    return images
+
+
+# =============================================================================
+# 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')}"
+
+
+def run_inference(task: str, combo: str):
+    """执行推理"""
+    engine = get_inference_engine()
+    
+    if not engine:
+        return {
+            "error": "Inference engine not available. Please check model checkpoints."
+        }
+    
+    # 获取样本信息
+    sample_info = get_sample_info(task, combo)
+    if not sample_info:
+        return {"error": f"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 result
+        
+    except Exception as e:
+        return {"error": f"Inference error: {str(e)}"}
+
+
+def format_prediction(result: Dict) -> str:
+    """格式化预测结果为 Markdown"""
+    if "error" in result:
+        return f"❌ **Error**: {result['error']}"
+    
+    task = result.get("task", "Unknown")
+    task_type = result.get("task_type", "unknown")
+    
+    lines = [
+        f"## Prediction Result: {TASK_CONFIG.get(task, {}).get('name', task)}",
+        "",
+    ]
+    
+    if task_type == "classification":
+        pred = result.get("prediction", "Unknown")
+        prob = result.get("probability", 0)
+        conf = result.get("confidence", 0)
+        classes = result.get("classes", [])
+        
+        lines.extend([
+            f"**Predicted Class**: {pred}",
+            "",
+            f"**Probability**:",
+            f"- {classes[0]}: {1-prob:.3f}",
+            f"- {classes[1]}: {prob:.3f}",
+            "",
+            f"**Confidence**: {conf:.1%}",
+        ])
+    else:
+        pred = result.get("prediction", 0)
+        unit = result.get("unit", "")
+        ref_range = result.get("reference_range", [0, 100])
+        
+        lines.extend([
+            f"**Predicted Value**: {pred:.2f} {unit}",
+            "",
+            f"**Reference Range**: [{ref_range[0]}, {ref_range[1]}]",
+        ])
+    
+    lines.extend([
+        "",
+        f"**Input Modalities**: {', '.join(result.get('input', {}).get('modalities', []))}",
+    ])
+    
+    return "\n".join(lines)
+
+
+# =============================================================================
+# 创建 Gradio 界面
+# =============================================================================
+
+def create_demo() -> gr.Blocks:
+    """创建 Gradio Demo 界面"""
+    
+    with gr.Blocks(
+        title="BrainAnytime Demo",
+        css="""
+        .preview-image { max-height: 300px; }
+        .result-box { font-size: 16px; }
+        """
+    ) as demo:
+        
+        gr.Markdown("""
+        # 🧠 BrainAnytime 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.
+        
+        ---
+        """)
+        
+        # ==================== Tab 1: 在线推理 ====================
+        with gr.Tab("🎯 在线推理 (Online Inference)"):
+            gr.Markdown("""
+            Select a task and modality combination to run inference on pre-selected samples.
+            Each sample is from the ADNI training set and has been verified for accuracy.
+            """)
+            
+            with gr.Row():
+                # 左侧：选择面板
+                with gr.Column(scale=1):
+                    gr.Markdown("### Configuration")
+                    
+                    task_selector = gr.Radio(
+                        choices=list(TASK_CONFIG.keys()),
+                        value="CN_vs_AD",
+                        label="Task",
+                        info="Select the prediction task"
+                    )
+                    
+                    # 显示任务描述
+                    task_desc = gr.Markdown(
+                        TASK_CONFIG["CN_vs_AD"]["description"]
+                    )
+                    
+                    modality_selector = gr.Radio(
+                        choices=list(MODALITY_COMBOS.keys()),
+                        value="T",
+                        label="Modality Combination",
+                        info="Select available modalities (only experiment-supported combinations)"
+                    )
+                    
+                    # 显示模态详情
+                    modality_desc = gr.Markdown(
+                        "**T**: Only T1-weighted MRI"
+                    )
+                    
+                    run_btn = gr.Button(
+                        "▶️ Run Inference",
+                        variant="primary",
+                        size="lg"
+                    )
+                
+                # 右侧：结果展示
+                with gr.Column(scale=2):
+                    gr.Markdown("### Sample Preview")
+                    
+                    sample_preview = gr.Image(
+                        label="Brain MRI Preview (Axial Slice)",
+                        type="filepath",
+                        height=300,
+                    )
+                    
+                    sample_info = gr.Textbox(
+                        label="Sample Info",
+                        interactive=False,
+                    )
+                    
+                    gr.Markdown("### Prediction Result")
+                    
+                    result_display = gr.Markdown(
+                        "Click 'Run Inference' to see results",
+                        elem_classes=["result-box"]
+                    )
+            
+            # 事件绑定
+            def update_task_desc(task):
+                return TASK_CONFIG.get(task, {}).get("description", "")
+            
+            def update_modality_desc(combo):
+                combo_info = MODALITY_COMBOS.get(combo, {})
+                mods = combo_info.get("modalities", [])
+                return f"**{combo}**: {', '.join(mods)}"
+            
+            task_selector.change(
+                update_task_desc,
+                inputs=task_selector,
+                outputs=task_desc
+            )
+            
+            modality_selector.change(
+                update_modality_desc,
+                inputs=modality_selector,
+                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]
+            )
+            
+            # 运行推理
+            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]
+            )
+        
+        # ==================== Tab 2: 项目信息 ====================
+        with gr.Tab("📖 项目信息 (Project Info)"):
+            gr.Markdown("""
+            ## About BrainAnytime
+            
+            **Paper**: BrainAnytime: Anatomy-Aware Cross-Modal Pretraining for Brain Image Analysis 
+            with Arbitrary Modality Availability
+            
+            **Conference**: MICCAI 2026 (Early Accept, Top 9%)
+            
+            ### Key Features
+            
+            - **Multi-modal Support**: T1, T2, Flair, PET
+            - **Missing Modality Robustness**: Handles arbitrary missing combinations
+            - **Anatomy-Aware**: Uses AAL116 brain atlas for adaptive masking
+            - **Pretrained Model**: Self-supervised pretraining on large-scale datasets
+            
+            ### Supported Tasks
+            
+            | Task | Type | Description |
+            |------|------|-------------|
+            | CN vs AD | Classification | Distinguish Normal vs Alzheimer's |
+            | CN vs MCI | Classification | Distinguish Normal vs Mild Cognitive Impairment |
+            | MMSE | Regression | Predict cognitive score (10-30) |
+            | AGE | Regression | Predict age from brain MRI |
+            
+            ### Links
+            
+            - **GitHub**: https://github.com/guangqianyang/BrainAnytime
+            - **Model Weights**: https://huggingface.co/Simmonstt/BrainAnytime
+            - **Paper**: [arXiv:2605.13059](https://arxiv.org/abs/2605.13059)
+            
+            ### Citation
+            
+            ```bibtex
+            @misc{yang2026brainanytime,
+                title={BrainAnytime: Anatomy-Aware Cross-Modal Pretraining for Brain Image Analysis},
+                author={Yang, Guangqian and Ding, Tong and others},
+                year={2026},
+                eprint={2605.13059},
+                archivePrefix={arXiv},
+            }
+            ```
+            """)
+        
+        # ==================== 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("""
+        ---
+        
+        <center>
+        <small>
+        BrainAnytime Demo | Built with Gradio | 
+        <a href="https://huggingface.co/Simmonstt/BrainAnytime">Model</a> |
+        <a href="https://github.com/guangqianyang/BrainAnytime">GitHub</a>
+        </small>
+        </center>
+        """)
+    
+    return demo
 
 
-with gr.Blocks(title="BrainAnytime Demo") as demo:
-    gr.Markdown(INTRO)
-    gr.Button("Show project details").click(show_project_info, outputs=gr.Markdown())
+# =============================================================================
+# 主入口
+# =============================================================================
 
 if __name__ == "__main__":
-    demo.launch()
+    import argparse
+    
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--host", type=str, default="0.0.0.0")
+    parser.add_argument("--port", type=int, default=7860)
+    parser.add_argument("--share", action="store_true")
+    
+    args = parser.parse_args()
+    
+    # 预加载模型（可选）
+    if INFERENCE_AVAILABLE:
+        print("Initializing inference engine...")
+        engine = get_inference_engine()
+        if engine:
+            print("✅ Inference engine ready")
+        else:
+            print("⚠️ Inference engine failed to initialize")
+    
+    # 检查样本数据
+    if not SAMPLES_INDEX.exists():
+        print(f"⚠️ Warning: Samples not found at {SAMPLES_INDEX}")
+        print("Please run prepare_samples.py first")
+    else:
+        samples = load_samples_index()
+        print(f"✅ Loaded {samples.get('total_samples', 0)} samples")
+    
+    # 启动 Demo
+    demo = create_demo()
+    demo.launch(
+        server_name=args.host,
+        server_port=args.port,
+        share=args.share,
+    )

demo_samples/AGE/T/sample_016/sample_016_T1.nii.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a490ea2db75a7fcd46cd8caf2d5d59be722036062265ccb705f79f404e68d22e
+size 2099660

demo_samples/AGE/T/sample_016/sample_016_meta.json

@@ -0,0 +1,33 @@
+{
+  "task": "AGE",
+  "modality_combination": "T",
+  "modalities": [
+    "T1"
+  ],
+  "modality_indices": [
+    1,
+    0,
+    0,
+    0
+  ],
+  "subject_id": "007_S_4516",
+  "sample_name": "sample_016",
+  "original_labels": {
+    "AGE": 73.4839151266256,
+    "MMSE": 28.0,
+    "DX": 1,
+    "Category": "CN"
+  },
+  "task_label": {
+    "value": 73.4839151266256,
+    "type": "regression"
+  },
+  "files": {
+    "nifti": {
+      "T1": "sample_016_T1.nii.gz"
+    },
+    "previews": {
+      "T1": "sample_016_T1_preview.png"
+    }
+  }
+}

demo_samples/AGE/TF/sample_017/sample_017_Flair.nii.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:85c84cf50563a38c9ef91a9d5ab2943a4ce14193cce042157cd0f9acf90fb880
+size 1110239

demo_samples/AGE/TF/sample_017/sample_017_T1.nii.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a490ea2db75a7fcd46cd8caf2d5d59be722036062265ccb705f79f404e68d22e
+size 2099660

demo_samples/AGE/TF/sample_017/sample_017_meta.json

@@ -0,0 +1,36 @@
+{
+  "task": "AGE",
+  "modality_combination": "TF",
+  "modalities": [
+    "T1",
+    "Flair"
+  ],
+  "modality_indices": [
+    1,
+    0,
+    1,
+    0
+  ],
+  "subject_id": "007_S_4516",
+  "sample_name": "sample_017",
+  "original_labels": {
+    "AGE": 73.4839151266256,
+    "MMSE": 28.0,
+    "DX": 1,
+    "Category": "CN"
+  },
+  "task_label": {
+    "value": 73.4839151266256,
+    "type": "regression"
+  },
+  "files": {
+    "nifti": {
+      "T1": "sample_017_T1.nii.gz",
+      "Flair": "sample_017_Flair.nii.gz"
+    },
+    "previews": {
+      "T1": "sample_017_T1_preview.png",
+      "Flair": "sample_017_Flair_preview.png"
+    }
+  }
+}

demo_samples/AGE/TFP/sample_019/sample_019_Flair.nii.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:85c84cf50563a38c9ef91a9d5ab2943a4ce14193cce042157cd0f9acf90fb880
+size 1110239

demo_samples/AGE/TFP/sample_019/sample_019_PET.nii.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:546ec204487be18e75f950e214ffc8917aba2f7678c0db38a5a88f1547550d3c
+size 2198379

demo_samples/AGE/TFP/sample_019/sample_019_T1.nii.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a490ea2db75a7fcd46cd8caf2d5d59be722036062265ccb705f79f404e68d22e
+size 2099660

demo_samples/AGE/TFP/sample_019/sample_019_meta.json

@@ -0,0 +1,39 @@
+{
+  "task": "AGE",
+  "modality_combination": "TFP",
+  "modalities": [
+    "T1",
+    "Flair",
+    "PET"
+  ],
+  "modality_indices": [
+    1,
+    0,
+    1,
+    1
+  ],
+  "subject_id": "007_S_4516",
+  "sample_name": "sample_019",
+  "original_labels": {
+    "AGE": 73.4839151266256,
+    "MMSE": 28.0,
+    "DX": 1,
+    "Category": "CN"
+  },
+  "task_label": {
+    "value": 73.4839151266256,
+    "type": "regression"
+  },
+  "files": {
+    "nifti": {
+      "T1": "sample_019_T1.nii.gz",
+      "Flair": "sample_019_Flair.nii.gz",
+      "PET": "sample_019_PET.nii.gz"
+    },
+    "previews": {
+      "T1": "sample_019_T1_preview.png",
+      "Flair": "sample_019_Flair_preview.png",
+      "PET": "sample_019_PET_preview.png"
+    }
+  }
+}

demo_samples/AGE/TMF/sample_018/sample_018_Flair.nii.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:85c84cf50563a38c9ef91a9d5ab2943a4ce14193cce042157cd0f9acf90fb880
+size 1110239

demo_samples/AGE/TMF/sample_018/sample_018_T1.nii.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a490ea2db75a7fcd46cd8caf2d5d59be722036062265ccb705f79f404e68d22e
+size 2099660

demo_samples/AGE/TMF/sample_018/sample_018_T2.nii.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9dd616c68eeb474a260b5cb98cdfe5d6fce9d6c2d28356fd8b4e35e0207f94d9
+size 1522248

demo_samples/AGE/TMF/sample_018/sample_018_meta.json

@@ -0,0 +1,39 @@
+{
+  "task": "AGE",
+  "modality_combination": "TMF",
+  "modalities": [
+    "T1",
+    "T2",
+    "Flair"
+  ],
+  "modality_indices": [
+    1,
+    1,
+    1,
+    0
+  ],
+  "subject_id": "007_S_4516",
+  "sample_name": "sample_018",
+  "original_labels": {
+    "AGE": 73.4839151266256,
+    "MMSE": 28.0,
+    "DX": 1,
+    "Category": "CN"
+  },
+  "task_label": {
+    "value": 73.4839151266256,
+    "type": "regression"
+  },
+  "files": {
+    "nifti": {
+      "T1": "sample_018_T1.nii.gz",
+      "T2": "sample_018_T2.nii.gz",
+      "Flair": "sample_018_Flair.nii.gz"
+    },
+    "previews": {
+      "T1": "sample_018_T1_preview.png",
+      "T2": "sample_018_T2_preview.png",
+      "Flair": "sample_018_Flair_preview.png"
+    }
+  }
+}

demo_samples/AGE/TMFP/sample_020/sample_020_Flair.nii.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:85c84cf50563a38c9ef91a9d5ab2943a4ce14193cce042157cd0f9acf90fb880
+size 1110239

demo_samples/AGE/TMFP/sample_020/sample_020_PET.nii.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:546ec204487be18e75f950e214ffc8917aba2f7678c0db38a5a88f1547550d3c
+size 2198379

demo_samples/AGE/TMFP/sample_020/sample_020_T1.nii.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a490ea2db75a7fcd46cd8caf2d5d59be722036062265ccb705f79f404e68d22e
+size 2099660

demo_samples/AGE/TMFP/sample_020/sample_020_T2.nii.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9dd616c68eeb474a260b5cb98cdfe5d6fce9d6c2d28356fd8b4e35e0207f94d9
+size 1522248

demo_samples/AGE/TMFP/sample_020/sample_020_meta.json

@@ -0,0 +1,42 @@
+{
+  "task": "AGE",
+  "modality_combination": "TMFP",
+  "modalities": [
+    "T1",
+    "T2",
+    "Flair",
+    "PET"
+  ],
+  "modality_indices": [
+    1,
+    1,
+    1,
+    1
+  ],
+  "subject_id": "007_S_4516",
+  "sample_name": "sample_020",
+  "original_labels": {
+    "AGE": 73.4839151266256,
+    "MMSE": 28.0,
+    "DX": 1,
+    "Category": "CN"
+  },
+  "task_label": {
+    "value": 73.4839151266256,
+    "type": "regression"
+  },
+  "files": {
+    "nifti": {
+      "T1": "sample_020_T1.nii.gz",
+      "T2": "sample_020_T2.nii.gz",
+      "Flair": "sample_020_Flair.nii.gz",
+      "PET": "sample_020_PET.nii.gz"
+    },
+    "previews": {
+      "T1": "sample_020_T1_preview.png",
+      "T2": "sample_020_T2_preview.png",
+      "Flair": "sample_020_Flair_preview.png",
+      "PET": "sample_020_PET_preview.png"
+    }
+  }
+}

demo_samples/CN_vs_AD/T/sample_001/sample_001_T1.nii.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a490ea2db75a7fcd46cd8caf2d5d59be722036062265ccb705f79f404e68d22e
+size 2099660

demo_samples/CN_vs_AD/T/sample_001/sample_001_meta.json

@@ -0,0 +1,33 @@
+{
+  "task": "CN_vs_AD",
+  "modality_combination": "T",
+  "modalities": [
+    "T1"
+  ],
+  "modality_indices": [
+    1,
+    0,
+    0,
+    0
+  ],
+  "subject_id": "007_S_4516",
+  "sample_name": "sample_001",
+  "original_labels": {
+    "AGE": 73.4839151266256,
+    "MMSE": 28.0,
+    "DX": 1,
+    "Category": "CN"
+  },
+  "task_label": {
+    "value": 0.0,
+    "type": "classification"
+  },
+  "files": {
+    "nifti": {
+      "T1": "sample_001_T1.nii.gz"
+    },
+    "previews": {
+      "T1": "sample_001_T1_preview.png"
+    }
+  }
+}

demo_samples/CN_vs_AD/TF/sample_002/sample_002_Flair.nii.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:85c84cf50563a38c9ef91a9d5ab2943a4ce14193cce042157cd0f9acf90fb880
+size 1110239

demo_samples/CN_vs_AD/TF/sample_002/sample_002_T1.nii.gz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a490ea2db75a7fcd46cd8caf2d5d59be722036062265ccb705f79f404e68d22e
+size 2099660

demo_samples/CN_vs_AD/TF/sample_002/sample_002_meta.json

@@ -0,0 +1,36 @@
+{
+  "task": "CN_vs_AD",
+  "modality_combination": "TF",
+  "modalities": [
+    "T1",
+    "Flair"
+  ],
+  "modality_indices": [
+    1,
+    0,
+    1,
+    0
+  ],
+  "subject_id": "007_S_4516",
+  "sample_name": "sample_002",
+  "original_labels": {
+    "AGE": 73.4839151266256,
+    "MMSE": 28.0,
+    "DX": 1,
+    "Category": "CN"
+  },
+  "task_label": {
+    "value": 0.0,
+    "type": "classification"
+  },
+  "files": {
+    "nifti": {
+      "T1": "sample_002_T1.nii.gz",
+      "Flair": "sample_002_Flair.nii.gz"
+    },
+    "previews": {
+      "T1": "sample_002_T1_preview.png",
+      "Flair": "sample_002_Flair_preview.png"
+    }
+  }
+}