diff --git a/BrainAnytime/altas/AAL116_standard.nii.gz b/BrainAnytime/altas/AAL116_standard.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..822e27bbd07f8d5fcf532d6ebf7542e7e6d9cbd0 --- /dev/null +++ b/BrainAnytime/altas/AAL116_standard.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:1a1e8d16d6b7616f3904f4319683692dda934010a52ace85de082d12702d2732 +size 203563 diff --git a/BrainAnytime/anatomy_masking.py b/BrainAnytime/anatomy_masking.py new file mode 100644 index 0000000000000000000000000000000000000000..7873f2f8eb3b20f0e605f14e176b5184bf398800 --- /dev/null +++ b/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 diff --git a/BrainAnytime/downstream_dataloader.py b/BrainAnytime/downstream_dataloader.py new file mode 100644 index 0000000000000000000000000000000000000000..dfb812ef8607abf9c1f402fb926f5a525d7e8407 --- /dev/null +++ b/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) + + + diff --git a/BrainAnytime/finetune_main.py b/BrainAnytime/finetune_main.py new file mode 100644 index 0000000000000000000000000000000000000000..7867017ba78447c758eed1675ea97056855ff49f --- /dev/null +++ b/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() diff --git a/BrainAnytime/models/__init__.py b/BrainAnytime/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/BrainAnytime/models/multimae3d.py b/BrainAnytime/models/multimae3d.py new file mode 100644 index 0000000000000000000000000000000000000000..a74512b6daa3a34d6402c4e69174c4c660d60b1a --- /dev/null +++ b/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, + ) diff --git a/BrainAnytime/models/multimae3d_utils.py b/BrainAnytime/models/multimae3d_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d8b6d29f9f1fac130012844c01543736ca701561 --- /dev/null +++ b/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 diff --git a/BrainAnytime/pretrain_dataloader_v2.py b/BrainAnytime/pretrain_dataloader_v2.py new file mode 100644 index 0000000000000000000000000000000000000000..2813e66178e4f6aa7ff5fe7c7b2aede38e0e4b35 --- /dev/null +++ b/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) + diff --git a/BrainAnytime/test_main.py b/BrainAnytime/test_main.py new file mode 100644 index 0000000000000000000000000000000000000000..1792fc6d9d7b192c084a70f53d8a9290bdf01b6f --- /dev/null +++ b/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() diff --git a/README.md b/README.md index cf1e2f9797fcec95f1b4db37f8e96b24292d6c84..c4cab4ec48778e3b757f126282dfd56c857a4359 100644 --- a/README.md +++ b/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 diff --git a/app.py b/app.py index 0f229e4f476ed590c5d1c844b30e18ba5896c9dc..cc0b53eaffdd721b57021fb7541999a21836f63b 100644 --- a/app.py +++ b/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 -``` -""" +_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(""" + --- + +
+ + BrainAnytime Demo | Built with Gradio | + Model | + GitHub + +
+ """) + + 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, + ) diff --git a/demo_samples/AGE/T/sample_016/sample_016_T1.nii.gz b/demo_samples/AGE/T/sample_016/sample_016_T1.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..97e8ee1670ca4d848972f6cba49108762b4e27ea --- /dev/null +++ b/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 diff --git a/demo_samples/AGE/T/sample_016/sample_016_T1_preview.png b/demo_samples/AGE/T/sample_016/sample_016_T1_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..4478a91240ec5c23148819126dc19cba761b807a Binary files /dev/null and b/demo_samples/AGE/T/sample_016/sample_016_T1_preview.png differ diff --git a/demo_samples/AGE/T/sample_016/sample_016_meta.json b/demo_samples/AGE/T/sample_016/sample_016_meta.json new file mode 100644 index 0000000000000000000000000000000000000000..6ba5e393b95394dff30cb3ce46fa9ebd03b581b1 --- /dev/null +++ b/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" + } + } +} \ No newline at end of file diff --git a/demo_samples/AGE/TF/sample_017/sample_017_Flair.nii.gz b/demo_samples/AGE/TF/sample_017/sample_017_Flair.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..19c33f314e5c8eb55f47a0c08e215ed7ba027d16 --- /dev/null +++ b/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 diff --git a/demo_samples/AGE/TF/sample_017/sample_017_Flair_preview.png b/demo_samples/AGE/TF/sample_017/sample_017_Flair_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..64f9ad29dc116c91f2b6c68251d1fa44f040ea2f Binary files /dev/null and b/demo_samples/AGE/TF/sample_017/sample_017_Flair_preview.png differ diff --git a/demo_samples/AGE/TF/sample_017/sample_017_T1.nii.gz b/demo_samples/AGE/TF/sample_017/sample_017_T1.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..97e8ee1670ca4d848972f6cba49108762b4e27ea --- /dev/null +++ b/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 diff --git a/demo_samples/AGE/TF/sample_017/sample_017_T1_preview.png b/demo_samples/AGE/TF/sample_017/sample_017_T1_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..4478a91240ec5c23148819126dc19cba761b807a Binary files /dev/null and b/demo_samples/AGE/TF/sample_017/sample_017_T1_preview.png differ diff --git a/demo_samples/AGE/TF/sample_017/sample_017_meta.json b/demo_samples/AGE/TF/sample_017/sample_017_meta.json new file mode 100644 index 0000000000000000000000000000000000000000..12dcd4e3a5573c1ce80e045c1f48d56e2fb10c5e --- /dev/null +++ b/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" + } + } +} \ No newline at end of file diff --git a/demo_samples/AGE/TFP/sample_019/sample_019_Flair.nii.gz b/demo_samples/AGE/TFP/sample_019/sample_019_Flair.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..19c33f314e5c8eb55f47a0c08e215ed7ba027d16 --- /dev/null +++ b/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 diff --git a/demo_samples/AGE/TFP/sample_019/sample_019_Flair_preview.png b/demo_samples/AGE/TFP/sample_019/sample_019_Flair_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..64f9ad29dc116c91f2b6c68251d1fa44f040ea2f Binary files /dev/null and b/demo_samples/AGE/TFP/sample_019/sample_019_Flair_preview.png differ diff --git a/demo_samples/AGE/TFP/sample_019/sample_019_PET.nii.gz b/demo_samples/AGE/TFP/sample_019/sample_019_PET.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..30a260b830e45baad03aafee0adaa3bf499b617b --- /dev/null +++ b/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 diff --git a/demo_samples/AGE/TFP/sample_019/sample_019_PET_preview.png b/demo_samples/AGE/TFP/sample_019/sample_019_PET_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..f0547287323daa1a56e22447360295e6debe8d7c Binary files /dev/null and b/demo_samples/AGE/TFP/sample_019/sample_019_PET_preview.png differ diff --git a/demo_samples/AGE/TFP/sample_019/sample_019_T1.nii.gz b/demo_samples/AGE/TFP/sample_019/sample_019_T1.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..97e8ee1670ca4d848972f6cba49108762b4e27ea --- /dev/null +++ b/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 diff --git a/demo_samples/AGE/TFP/sample_019/sample_019_T1_preview.png b/demo_samples/AGE/TFP/sample_019/sample_019_T1_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..4478a91240ec5c23148819126dc19cba761b807a Binary files /dev/null and b/demo_samples/AGE/TFP/sample_019/sample_019_T1_preview.png differ diff --git a/demo_samples/AGE/TFP/sample_019/sample_019_meta.json b/demo_samples/AGE/TFP/sample_019/sample_019_meta.json new file mode 100644 index 0000000000000000000000000000000000000000..15f849798e9be1d4b1835f5f77ecb5f0274950ad --- /dev/null +++ b/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" + } + } +} \ No newline at end of file diff --git a/demo_samples/AGE/TMF/sample_018/sample_018_Flair.nii.gz b/demo_samples/AGE/TMF/sample_018/sample_018_Flair.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..19c33f314e5c8eb55f47a0c08e215ed7ba027d16 --- /dev/null +++ b/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 diff --git a/demo_samples/AGE/TMF/sample_018/sample_018_Flair_preview.png b/demo_samples/AGE/TMF/sample_018/sample_018_Flair_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..64f9ad29dc116c91f2b6c68251d1fa44f040ea2f Binary files /dev/null and b/demo_samples/AGE/TMF/sample_018/sample_018_Flair_preview.png differ diff --git a/demo_samples/AGE/TMF/sample_018/sample_018_T1.nii.gz b/demo_samples/AGE/TMF/sample_018/sample_018_T1.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..97e8ee1670ca4d848972f6cba49108762b4e27ea --- /dev/null +++ b/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 diff --git a/demo_samples/AGE/TMF/sample_018/sample_018_T1_preview.png b/demo_samples/AGE/TMF/sample_018/sample_018_T1_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..4478a91240ec5c23148819126dc19cba761b807a Binary files /dev/null and b/demo_samples/AGE/TMF/sample_018/sample_018_T1_preview.png differ diff --git a/demo_samples/AGE/TMF/sample_018/sample_018_T2.nii.gz b/demo_samples/AGE/TMF/sample_018/sample_018_T2.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..0c67cc72f21fa5c43b9a3d7136e64644a2ac466a --- /dev/null +++ b/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 diff --git a/demo_samples/AGE/TMF/sample_018/sample_018_T2_preview.png b/demo_samples/AGE/TMF/sample_018/sample_018_T2_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..822166d3ec8b6e8962ba70549509a17c023aac55 Binary files /dev/null and b/demo_samples/AGE/TMF/sample_018/sample_018_T2_preview.png differ diff --git a/demo_samples/AGE/TMF/sample_018/sample_018_meta.json b/demo_samples/AGE/TMF/sample_018/sample_018_meta.json new file mode 100644 index 0000000000000000000000000000000000000000..5984d9ef5dc223c71b13b9bef7bcf3a3e09a25f5 --- /dev/null +++ b/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" + } + } +} \ No newline at end of file diff --git a/demo_samples/AGE/TMFP/sample_020/sample_020_Flair.nii.gz b/demo_samples/AGE/TMFP/sample_020/sample_020_Flair.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..19c33f314e5c8eb55f47a0c08e215ed7ba027d16 --- /dev/null +++ b/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 diff --git a/demo_samples/AGE/TMFP/sample_020/sample_020_Flair_preview.png b/demo_samples/AGE/TMFP/sample_020/sample_020_Flair_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..64f9ad29dc116c91f2b6c68251d1fa44f040ea2f Binary files /dev/null and b/demo_samples/AGE/TMFP/sample_020/sample_020_Flair_preview.png differ diff --git a/demo_samples/AGE/TMFP/sample_020/sample_020_PET.nii.gz b/demo_samples/AGE/TMFP/sample_020/sample_020_PET.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..30a260b830e45baad03aafee0adaa3bf499b617b --- /dev/null +++ b/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 diff --git a/demo_samples/AGE/TMFP/sample_020/sample_020_PET_preview.png b/demo_samples/AGE/TMFP/sample_020/sample_020_PET_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..f0547287323daa1a56e22447360295e6debe8d7c Binary files /dev/null and b/demo_samples/AGE/TMFP/sample_020/sample_020_PET_preview.png differ diff --git a/demo_samples/AGE/TMFP/sample_020/sample_020_T1.nii.gz b/demo_samples/AGE/TMFP/sample_020/sample_020_T1.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..97e8ee1670ca4d848972f6cba49108762b4e27ea --- /dev/null +++ b/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 diff --git a/demo_samples/AGE/TMFP/sample_020/sample_020_T1_preview.png b/demo_samples/AGE/TMFP/sample_020/sample_020_T1_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..4478a91240ec5c23148819126dc19cba761b807a Binary files /dev/null and b/demo_samples/AGE/TMFP/sample_020/sample_020_T1_preview.png differ diff --git a/demo_samples/AGE/TMFP/sample_020/sample_020_T2.nii.gz b/demo_samples/AGE/TMFP/sample_020/sample_020_T2.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..0c67cc72f21fa5c43b9a3d7136e64644a2ac466a --- /dev/null +++ b/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 diff --git a/demo_samples/AGE/TMFP/sample_020/sample_020_T2_preview.png b/demo_samples/AGE/TMFP/sample_020/sample_020_T2_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..822166d3ec8b6e8962ba70549509a17c023aac55 Binary files /dev/null and b/demo_samples/AGE/TMFP/sample_020/sample_020_T2_preview.png differ diff --git a/demo_samples/AGE/TMFP/sample_020/sample_020_meta.json b/demo_samples/AGE/TMFP/sample_020/sample_020_meta.json new file mode 100644 index 0000000000000000000000000000000000000000..c906f59a5ff0d85d447fedb00a68890e33ec8d05 --- /dev/null +++ b/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" + } + } +} \ No newline at end of file diff --git a/demo_samples/CN_vs_AD/T/sample_001/sample_001_T1.nii.gz b/demo_samples/CN_vs_AD/T/sample_001/sample_001_T1.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..97e8ee1670ca4d848972f6cba49108762b4e27ea --- /dev/null +++ b/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 diff --git a/demo_samples/CN_vs_AD/T/sample_001/sample_001_T1_preview.png b/demo_samples/CN_vs_AD/T/sample_001/sample_001_T1_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..4478a91240ec5c23148819126dc19cba761b807a Binary files /dev/null and b/demo_samples/CN_vs_AD/T/sample_001/sample_001_T1_preview.png differ diff --git a/demo_samples/CN_vs_AD/T/sample_001/sample_001_meta.json b/demo_samples/CN_vs_AD/T/sample_001/sample_001_meta.json new file mode 100644 index 0000000000000000000000000000000000000000..d14b073c05980ce21d997d17119f66b882050d9a --- /dev/null +++ b/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" + } + } +} \ No newline at end of file diff --git a/demo_samples/CN_vs_AD/TF/sample_002/sample_002_Flair.nii.gz b/demo_samples/CN_vs_AD/TF/sample_002/sample_002_Flair.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..19c33f314e5c8eb55f47a0c08e215ed7ba027d16 --- /dev/null +++ b/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 diff --git a/demo_samples/CN_vs_AD/TF/sample_002/sample_002_Flair_preview.png b/demo_samples/CN_vs_AD/TF/sample_002/sample_002_Flair_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..64f9ad29dc116c91f2b6c68251d1fa44f040ea2f Binary files /dev/null and b/demo_samples/CN_vs_AD/TF/sample_002/sample_002_Flair_preview.png differ diff --git a/demo_samples/CN_vs_AD/TF/sample_002/sample_002_T1.nii.gz b/demo_samples/CN_vs_AD/TF/sample_002/sample_002_T1.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..97e8ee1670ca4d848972f6cba49108762b4e27ea --- /dev/null +++ b/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 diff --git a/demo_samples/CN_vs_AD/TF/sample_002/sample_002_T1_preview.png b/demo_samples/CN_vs_AD/TF/sample_002/sample_002_T1_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..4478a91240ec5c23148819126dc19cba761b807a Binary files /dev/null and b/demo_samples/CN_vs_AD/TF/sample_002/sample_002_T1_preview.png differ diff --git a/demo_samples/CN_vs_AD/TF/sample_002/sample_002_meta.json b/demo_samples/CN_vs_AD/TF/sample_002/sample_002_meta.json new file mode 100644 index 0000000000000000000000000000000000000000..8a9b545b48c45e23874f61370185931baf63c2ee --- /dev/null +++ b/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" + } + } +} \ No newline at end of file diff --git a/demo_samples/CN_vs_AD/TFP/sample_004/sample_004_Flair.nii.gz b/demo_samples/CN_vs_AD/TFP/sample_004/sample_004_Flair.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..19c33f314e5c8eb55f47a0c08e215ed7ba027d16 --- /dev/null +++ b/demo_samples/CN_vs_AD/TFP/sample_004/sample_004_Flair.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:85c84cf50563a38c9ef91a9d5ab2943a4ce14193cce042157cd0f9acf90fb880 +size 1110239 diff --git a/demo_samples/CN_vs_AD/TFP/sample_004/sample_004_Flair_preview.png b/demo_samples/CN_vs_AD/TFP/sample_004/sample_004_Flair_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..64f9ad29dc116c91f2b6c68251d1fa44f040ea2f Binary files /dev/null and b/demo_samples/CN_vs_AD/TFP/sample_004/sample_004_Flair_preview.png differ diff --git a/demo_samples/CN_vs_AD/TFP/sample_004/sample_004_PET.nii.gz b/demo_samples/CN_vs_AD/TFP/sample_004/sample_004_PET.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..30a260b830e45baad03aafee0adaa3bf499b617b --- /dev/null +++ b/demo_samples/CN_vs_AD/TFP/sample_004/sample_004_PET.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:546ec204487be18e75f950e214ffc8917aba2f7678c0db38a5a88f1547550d3c +size 2198379 diff --git a/demo_samples/CN_vs_AD/TFP/sample_004/sample_004_PET_preview.png b/demo_samples/CN_vs_AD/TFP/sample_004/sample_004_PET_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..f0547287323daa1a56e22447360295e6debe8d7c Binary files /dev/null and b/demo_samples/CN_vs_AD/TFP/sample_004/sample_004_PET_preview.png differ diff --git a/demo_samples/CN_vs_AD/TFP/sample_004/sample_004_T1.nii.gz b/demo_samples/CN_vs_AD/TFP/sample_004/sample_004_T1.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..97e8ee1670ca4d848972f6cba49108762b4e27ea --- /dev/null +++ b/demo_samples/CN_vs_AD/TFP/sample_004/sample_004_T1.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a490ea2db75a7fcd46cd8caf2d5d59be722036062265ccb705f79f404e68d22e +size 2099660 diff --git a/demo_samples/CN_vs_AD/TFP/sample_004/sample_004_T1_preview.png b/demo_samples/CN_vs_AD/TFP/sample_004/sample_004_T1_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..4478a91240ec5c23148819126dc19cba761b807a Binary files /dev/null and b/demo_samples/CN_vs_AD/TFP/sample_004/sample_004_T1_preview.png differ diff --git a/demo_samples/CN_vs_AD/TFP/sample_004/sample_004_meta.json b/demo_samples/CN_vs_AD/TFP/sample_004/sample_004_meta.json new file mode 100644 index 0000000000000000000000000000000000000000..13c1a51f0c03dfd9f67ad2119d9c4e5fc9bd2d58 --- /dev/null +++ b/demo_samples/CN_vs_AD/TFP/sample_004/sample_004_meta.json @@ -0,0 +1,39 @@ +{ + "task": "CN_vs_AD", + "modality_combination": "TFP", + "modalities": [ + "T1", + "Flair", + "PET" + ], + "modality_indices": [ + 1, + 0, + 1, + 1 + ], + "subject_id": "007_S_4516", + "sample_name": "sample_004", + "original_labels": { + "AGE": 73.4839151266256, + "MMSE": 28.0, + "DX": 1, + "Category": "CN" + }, + "task_label": { + "value": 0.0, + "type": "classification" + }, + "files": { + "nifti": { + "T1": "sample_004_T1.nii.gz", + "Flair": "sample_004_Flair.nii.gz", + "PET": "sample_004_PET.nii.gz" + }, + "previews": { + "T1": "sample_004_T1_preview.png", + "Flair": "sample_004_Flair_preview.png", + "PET": "sample_004_PET_preview.png" + } + } +} \ No newline at end of file diff --git a/demo_samples/CN_vs_AD/TMF/sample_003/sample_003_Flair.nii.gz b/demo_samples/CN_vs_AD/TMF/sample_003/sample_003_Flair.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..19c33f314e5c8eb55f47a0c08e215ed7ba027d16 --- /dev/null +++ b/demo_samples/CN_vs_AD/TMF/sample_003/sample_003_Flair.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:85c84cf50563a38c9ef91a9d5ab2943a4ce14193cce042157cd0f9acf90fb880 +size 1110239 diff --git a/demo_samples/CN_vs_AD/TMF/sample_003/sample_003_Flair_preview.png b/demo_samples/CN_vs_AD/TMF/sample_003/sample_003_Flair_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..64f9ad29dc116c91f2b6c68251d1fa44f040ea2f Binary files /dev/null and b/demo_samples/CN_vs_AD/TMF/sample_003/sample_003_Flair_preview.png differ diff --git a/demo_samples/CN_vs_AD/TMF/sample_003/sample_003_T1.nii.gz b/demo_samples/CN_vs_AD/TMF/sample_003/sample_003_T1.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..97e8ee1670ca4d848972f6cba49108762b4e27ea --- /dev/null +++ b/demo_samples/CN_vs_AD/TMF/sample_003/sample_003_T1.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a490ea2db75a7fcd46cd8caf2d5d59be722036062265ccb705f79f404e68d22e +size 2099660 diff --git a/demo_samples/CN_vs_AD/TMF/sample_003/sample_003_T1_preview.png b/demo_samples/CN_vs_AD/TMF/sample_003/sample_003_T1_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..4478a91240ec5c23148819126dc19cba761b807a Binary files /dev/null and b/demo_samples/CN_vs_AD/TMF/sample_003/sample_003_T1_preview.png differ diff --git a/demo_samples/CN_vs_AD/TMF/sample_003/sample_003_T2.nii.gz b/demo_samples/CN_vs_AD/TMF/sample_003/sample_003_T2.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..0c67cc72f21fa5c43b9a3d7136e64644a2ac466a --- /dev/null +++ b/demo_samples/CN_vs_AD/TMF/sample_003/sample_003_T2.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9dd616c68eeb474a260b5cb98cdfe5d6fce9d6c2d28356fd8b4e35e0207f94d9 +size 1522248 diff --git a/demo_samples/CN_vs_AD/TMF/sample_003/sample_003_T2_preview.png b/demo_samples/CN_vs_AD/TMF/sample_003/sample_003_T2_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..822166d3ec8b6e8962ba70549509a17c023aac55 Binary files /dev/null and b/demo_samples/CN_vs_AD/TMF/sample_003/sample_003_T2_preview.png differ diff --git a/demo_samples/CN_vs_AD/TMF/sample_003/sample_003_meta.json b/demo_samples/CN_vs_AD/TMF/sample_003/sample_003_meta.json new file mode 100644 index 0000000000000000000000000000000000000000..c2bf9c8448041bdd91fce982624ff862befc0350 --- /dev/null +++ b/demo_samples/CN_vs_AD/TMF/sample_003/sample_003_meta.json @@ -0,0 +1,39 @@ +{ + "task": "CN_vs_AD", + "modality_combination": "TMF", + "modalities": [ + "T1", + "T2", + "Flair" + ], + "modality_indices": [ + 1, + 1, + 1, + 0 + ], + "subject_id": "007_S_4516", + "sample_name": "sample_003", + "original_labels": { + "AGE": 73.4839151266256, + "MMSE": 28.0, + "DX": 1, + "Category": "CN" + }, + "task_label": { + "value": 0.0, + "type": "classification" + }, + "files": { + "nifti": { + "T1": "sample_003_T1.nii.gz", + "T2": "sample_003_T2.nii.gz", + "Flair": "sample_003_Flair.nii.gz" + }, + "previews": { + "T1": "sample_003_T1_preview.png", + "T2": "sample_003_T2_preview.png", + "Flair": "sample_003_Flair_preview.png" + } + } +} \ No newline at end of file diff --git a/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_Flair.nii.gz b/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_Flair.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..19c33f314e5c8eb55f47a0c08e215ed7ba027d16 --- /dev/null +++ b/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_Flair.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:85c84cf50563a38c9ef91a9d5ab2943a4ce14193cce042157cd0f9acf90fb880 +size 1110239 diff --git a/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_Flair_preview.png b/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_Flair_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..64f9ad29dc116c91f2b6c68251d1fa44f040ea2f Binary files /dev/null and b/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_Flair_preview.png differ diff --git a/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_PET.nii.gz b/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_PET.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..30a260b830e45baad03aafee0adaa3bf499b617b --- /dev/null +++ b/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_PET.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:546ec204487be18e75f950e214ffc8917aba2f7678c0db38a5a88f1547550d3c +size 2198379 diff --git a/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_PET_preview.png b/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_PET_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..f0547287323daa1a56e22447360295e6debe8d7c Binary files /dev/null and b/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_PET_preview.png differ diff --git a/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_T1.nii.gz b/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_T1.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..97e8ee1670ca4d848972f6cba49108762b4e27ea --- /dev/null +++ b/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_T1.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a490ea2db75a7fcd46cd8caf2d5d59be722036062265ccb705f79f404e68d22e +size 2099660 diff --git a/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_T1_preview.png b/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_T1_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..4478a91240ec5c23148819126dc19cba761b807a Binary files /dev/null and b/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_T1_preview.png differ diff --git a/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_T2.nii.gz b/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_T2.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..0c67cc72f21fa5c43b9a3d7136e64644a2ac466a --- /dev/null +++ b/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_T2.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9dd616c68eeb474a260b5cb98cdfe5d6fce9d6c2d28356fd8b4e35e0207f94d9 +size 1522248 diff --git a/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_T2_preview.png b/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_T2_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..822166d3ec8b6e8962ba70549509a17c023aac55 Binary files /dev/null and b/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_T2_preview.png differ diff --git a/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_meta.json b/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_meta.json new file mode 100644 index 0000000000000000000000000000000000000000..756781e4d320ad4693eaa9ebaf0941fbb2b3b7cc --- /dev/null +++ b/demo_samples/CN_vs_AD/TMFP/sample_005/sample_005_meta.json @@ -0,0 +1,42 @@ +{ + "task": "CN_vs_AD", + "modality_combination": "TMFP", + "modalities": [ + "T1", + "T2", + "Flair", + "PET" + ], + "modality_indices": [ + 1, + 1, + 1, + 1 + ], + "subject_id": "007_S_4516", + "sample_name": "sample_005", + "original_labels": { + "AGE": 73.4839151266256, + "MMSE": 28.0, + "DX": 1, + "Category": "CN" + }, + "task_label": { + "value": 0.0, + "type": "classification" + }, + "files": { + "nifti": { + "T1": "sample_005_T1.nii.gz", + "T2": "sample_005_T2.nii.gz", + "Flair": "sample_005_Flair.nii.gz", + "PET": "sample_005_PET.nii.gz" + }, + "previews": { + "T1": "sample_005_T1_preview.png", + "T2": "sample_005_T2_preview.png", + "Flair": "sample_005_Flair_preview.png", + "PET": "sample_005_PET_preview.png" + } + } +} \ No newline at end of file diff --git a/demo_samples/CN_vs_MCI/T/sample_006/sample_006_T1.nii.gz b/demo_samples/CN_vs_MCI/T/sample_006/sample_006_T1.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..97e8ee1670ca4d848972f6cba49108762b4e27ea --- /dev/null +++ b/demo_samples/CN_vs_MCI/T/sample_006/sample_006_T1.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a490ea2db75a7fcd46cd8caf2d5d59be722036062265ccb705f79f404e68d22e +size 2099660 diff --git a/demo_samples/CN_vs_MCI/T/sample_006/sample_006_T1_preview.png b/demo_samples/CN_vs_MCI/T/sample_006/sample_006_T1_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..4478a91240ec5c23148819126dc19cba761b807a Binary files /dev/null and b/demo_samples/CN_vs_MCI/T/sample_006/sample_006_T1_preview.png differ diff --git a/demo_samples/CN_vs_MCI/T/sample_006/sample_006_meta.json b/demo_samples/CN_vs_MCI/T/sample_006/sample_006_meta.json new file mode 100644 index 0000000000000000000000000000000000000000..5af3542a27de4b161d9d1d568a191700ecff9157 --- /dev/null +++ b/demo_samples/CN_vs_MCI/T/sample_006/sample_006_meta.json @@ -0,0 +1,33 @@ +{ + "task": "CN_vs_MCI", + "modality_combination": "T", + "modalities": [ + "T1" + ], + "modality_indices": [ + 1, + 0, + 0, + 0 + ], + "subject_id": "007_S_4516", + "sample_name": "sample_006", + "original_labels": { + "AGE": 73.4839151266256, + "MMSE": 28.0, + "DX": 1, + "Category": "CN" + }, + "task_label": { + "value": 0.0, + "type": "classification" + }, + "files": { + "nifti": { + "T1": "sample_006_T1.nii.gz" + }, + "previews": { + "T1": "sample_006_T1_preview.png" + } + } +} \ No newline at end of file diff --git a/demo_samples/CN_vs_MCI/TF/sample_007/sample_007_Flair.nii.gz b/demo_samples/CN_vs_MCI/TF/sample_007/sample_007_Flair.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..19c33f314e5c8eb55f47a0c08e215ed7ba027d16 --- /dev/null +++ b/demo_samples/CN_vs_MCI/TF/sample_007/sample_007_Flair.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:85c84cf50563a38c9ef91a9d5ab2943a4ce14193cce042157cd0f9acf90fb880 +size 1110239 diff --git a/demo_samples/CN_vs_MCI/TF/sample_007/sample_007_Flair_preview.png b/demo_samples/CN_vs_MCI/TF/sample_007/sample_007_Flair_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..64f9ad29dc116c91f2b6c68251d1fa44f040ea2f Binary files /dev/null and b/demo_samples/CN_vs_MCI/TF/sample_007/sample_007_Flair_preview.png differ diff --git a/demo_samples/CN_vs_MCI/TF/sample_007/sample_007_T1.nii.gz b/demo_samples/CN_vs_MCI/TF/sample_007/sample_007_T1.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..97e8ee1670ca4d848972f6cba49108762b4e27ea --- /dev/null +++ b/demo_samples/CN_vs_MCI/TF/sample_007/sample_007_T1.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a490ea2db75a7fcd46cd8caf2d5d59be722036062265ccb705f79f404e68d22e +size 2099660 diff --git a/demo_samples/CN_vs_MCI/TF/sample_007/sample_007_T1_preview.png b/demo_samples/CN_vs_MCI/TF/sample_007/sample_007_T1_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..4478a91240ec5c23148819126dc19cba761b807a Binary files /dev/null and b/demo_samples/CN_vs_MCI/TF/sample_007/sample_007_T1_preview.png differ diff --git a/demo_samples/CN_vs_MCI/TF/sample_007/sample_007_meta.json b/demo_samples/CN_vs_MCI/TF/sample_007/sample_007_meta.json new file mode 100644 index 0000000000000000000000000000000000000000..d72a7676363b47d80ba38bb8647be142bf1a8f1f --- /dev/null +++ b/demo_samples/CN_vs_MCI/TF/sample_007/sample_007_meta.json @@ -0,0 +1,36 @@ +{ + "task": "CN_vs_MCI", + "modality_combination": "TF", + "modalities": [ + "T1", + "Flair" + ], + "modality_indices": [ + 1, + 0, + 1, + 0 + ], + "subject_id": "007_S_4516", + "sample_name": "sample_007", + "original_labels": { + "AGE": 73.4839151266256, + "MMSE": 28.0, + "DX": 1, + "Category": "CN" + }, + "task_label": { + "value": 0.0, + "type": "classification" + }, + "files": { + "nifti": { + "T1": "sample_007_T1.nii.gz", + "Flair": "sample_007_Flair.nii.gz" + }, + "previews": { + "T1": "sample_007_T1_preview.png", + "Flair": "sample_007_Flair_preview.png" + } + } +} \ No newline at end of file diff --git a/demo_samples/CN_vs_MCI/TFP/sample_009/sample_009_Flair.nii.gz b/demo_samples/CN_vs_MCI/TFP/sample_009/sample_009_Flair.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..19c33f314e5c8eb55f47a0c08e215ed7ba027d16 --- /dev/null +++ b/demo_samples/CN_vs_MCI/TFP/sample_009/sample_009_Flair.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:85c84cf50563a38c9ef91a9d5ab2943a4ce14193cce042157cd0f9acf90fb880 +size 1110239 diff --git a/demo_samples/CN_vs_MCI/TFP/sample_009/sample_009_Flair_preview.png b/demo_samples/CN_vs_MCI/TFP/sample_009/sample_009_Flair_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..64f9ad29dc116c91f2b6c68251d1fa44f040ea2f Binary files /dev/null and b/demo_samples/CN_vs_MCI/TFP/sample_009/sample_009_Flair_preview.png differ diff --git a/demo_samples/CN_vs_MCI/TFP/sample_009/sample_009_PET.nii.gz b/demo_samples/CN_vs_MCI/TFP/sample_009/sample_009_PET.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..30a260b830e45baad03aafee0adaa3bf499b617b --- /dev/null +++ b/demo_samples/CN_vs_MCI/TFP/sample_009/sample_009_PET.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:546ec204487be18e75f950e214ffc8917aba2f7678c0db38a5a88f1547550d3c +size 2198379 diff --git a/demo_samples/CN_vs_MCI/TFP/sample_009/sample_009_PET_preview.png b/demo_samples/CN_vs_MCI/TFP/sample_009/sample_009_PET_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..f0547287323daa1a56e22447360295e6debe8d7c Binary files /dev/null and b/demo_samples/CN_vs_MCI/TFP/sample_009/sample_009_PET_preview.png differ diff --git a/demo_samples/CN_vs_MCI/TFP/sample_009/sample_009_T1.nii.gz b/demo_samples/CN_vs_MCI/TFP/sample_009/sample_009_T1.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..97e8ee1670ca4d848972f6cba49108762b4e27ea --- /dev/null +++ b/demo_samples/CN_vs_MCI/TFP/sample_009/sample_009_T1.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a490ea2db75a7fcd46cd8caf2d5d59be722036062265ccb705f79f404e68d22e +size 2099660 diff --git a/demo_samples/CN_vs_MCI/TFP/sample_009/sample_009_T1_preview.png b/demo_samples/CN_vs_MCI/TFP/sample_009/sample_009_T1_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..4478a91240ec5c23148819126dc19cba761b807a Binary files /dev/null and b/demo_samples/CN_vs_MCI/TFP/sample_009/sample_009_T1_preview.png differ diff --git a/demo_samples/CN_vs_MCI/TFP/sample_009/sample_009_meta.json b/demo_samples/CN_vs_MCI/TFP/sample_009/sample_009_meta.json new file mode 100644 index 0000000000000000000000000000000000000000..73b2b59fdb90bde0c9a287d444c1dcf635a2df71 --- /dev/null +++ b/demo_samples/CN_vs_MCI/TFP/sample_009/sample_009_meta.json @@ -0,0 +1,39 @@ +{ + "task": "CN_vs_MCI", + "modality_combination": "TFP", + "modalities": [ + "T1", + "Flair", + "PET" + ], + "modality_indices": [ + 1, + 0, + 1, + 1 + ], + "subject_id": "007_S_4516", + "sample_name": "sample_009", + "original_labels": { + "AGE": 73.4839151266256, + "MMSE": 28.0, + "DX": 1, + "Category": "CN" + }, + "task_label": { + "value": 0.0, + "type": "classification" + }, + "files": { + "nifti": { + "T1": "sample_009_T1.nii.gz", + "Flair": "sample_009_Flair.nii.gz", + "PET": "sample_009_PET.nii.gz" + }, + "previews": { + "T1": "sample_009_T1_preview.png", + "Flair": "sample_009_Flair_preview.png", + "PET": "sample_009_PET_preview.png" + } + } +} \ No newline at end of file diff --git a/demo_samples/CN_vs_MCI/TMF/sample_008/sample_008_Flair.nii.gz b/demo_samples/CN_vs_MCI/TMF/sample_008/sample_008_Flair.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..19c33f314e5c8eb55f47a0c08e215ed7ba027d16 --- /dev/null +++ b/demo_samples/CN_vs_MCI/TMF/sample_008/sample_008_Flair.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:85c84cf50563a38c9ef91a9d5ab2943a4ce14193cce042157cd0f9acf90fb880 +size 1110239 diff --git a/demo_samples/CN_vs_MCI/TMF/sample_008/sample_008_Flair_preview.png b/demo_samples/CN_vs_MCI/TMF/sample_008/sample_008_Flair_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..64f9ad29dc116c91f2b6c68251d1fa44f040ea2f Binary files /dev/null and b/demo_samples/CN_vs_MCI/TMF/sample_008/sample_008_Flair_preview.png differ diff --git a/demo_samples/CN_vs_MCI/TMF/sample_008/sample_008_T1.nii.gz b/demo_samples/CN_vs_MCI/TMF/sample_008/sample_008_T1.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..97e8ee1670ca4d848972f6cba49108762b4e27ea --- /dev/null +++ b/demo_samples/CN_vs_MCI/TMF/sample_008/sample_008_T1.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a490ea2db75a7fcd46cd8caf2d5d59be722036062265ccb705f79f404e68d22e +size 2099660 diff --git a/demo_samples/CN_vs_MCI/TMF/sample_008/sample_008_T1_preview.png b/demo_samples/CN_vs_MCI/TMF/sample_008/sample_008_T1_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..4478a91240ec5c23148819126dc19cba761b807a Binary files /dev/null and b/demo_samples/CN_vs_MCI/TMF/sample_008/sample_008_T1_preview.png differ diff --git a/demo_samples/CN_vs_MCI/TMF/sample_008/sample_008_T2.nii.gz b/demo_samples/CN_vs_MCI/TMF/sample_008/sample_008_T2.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..0c67cc72f21fa5c43b9a3d7136e64644a2ac466a --- /dev/null +++ b/demo_samples/CN_vs_MCI/TMF/sample_008/sample_008_T2.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9dd616c68eeb474a260b5cb98cdfe5d6fce9d6c2d28356fd8b4e35e0207f94d9 +size 1522248 diff --git a/demo_samples/CN_vs_MCI/TMF/sample_008/sample_008_T2_preview.png b/demo_samples/CN_vs_MCI/TMF/sample_008/sample_008_T2_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..822166d3ec8b6e8962ba70549509a17c023aac55 Binary files /dev/null and b/demo_samples/CN_vs_MCI/TMF/sample_008/sample_008_T2_preview.png differ diff --git a/demo_samples/CN_vs_MCI/TMF/sample_008/sample_008_meta.json b/demo_samples/CN_vs_MCI/TMF/sample_008/sample_008_meta.json new file mode 100644 index 0000000000000000000000000000000000000000..828f33283f167574105d1ec7c0cfdd94e509d480 --- /dev/null +++ b/demo_samples/CN_vs_MCI/TMF/sample_008/sample_008_meta.json @@ -0,0 +1,39 @@ +{ + "task": "CN_vs_MCI", + "modality_combination": "TMF", + "modalities": [ + "T1", + "T2", + "Flair" + ], + "modality_indices": [ + 1, + 1, + 1, + 0 + ], + "subject_id": "007_S_4516", + "sample_name": "sample_008", + "original_labels": { + "AGE": 73.4839151266256, + "MMSE": 28.0, + "DX": 1, + "Category": "CN" + }, + "task_label": { + "value": 0.0, + "type": "classification" + }, + "files": { + "nifti": { + "T1": "sample_008_T1.nii.gz", + "T2": "sample_008_T2.nii.gz", + "Flair": "sample_008_Flair.nii.gz" + }, + "previews": { + "T1": "sample_008_T1_preview.png", + "T2": "sample_008_T2_preview.png", + "Flair": "sample_008_Flair_preview.png" + } + } +} \ No newline at end of file diff --git a/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_Flair.nii.gz b/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_Flair.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..19c33f314e5c8eb55f47a0c08e215ed7ba027d16 --- /dev/null +++ b/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_Flair.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:85c84cf50563a38c9ef91a9d5ab2943a4ce14193cce042157cd0f9acf90fb880 +size 1110239 diff --git a/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_Flair_preview.png b/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_Flair_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..64f9ad29dc116c91f2b6c68251d1fa44f040ea2f Binary files /dev/null and b/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_Flair_preview.png differ diff --git a/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_PET.nii.gz b/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_PET.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..30a260b830e45baad03aafee0adaa3bf499b617b --- /dev/null +++ b/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_PET.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:546ec204487be18e75f950e214ffc8917aba2f7678c0db38a5a88f1547550d3c +size 2198379 diff --git a/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_PET_preview.png b/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_PET_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..f0547287323daa1a56e22447360295e6debe8d7c Binary files /dev/null and b/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_PET_preview.png differ diff --git a/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_T1.nii.gz b/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_T1.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..97e8ee1670ca4d848972f6cba49108762b4e27ea --- /dev/null +++ b/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_T1.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a490ea2db75a7fcd46cd8caf2d5d59be722036062265ccb705f79f404e68d22e +size 2099660 diff --git a/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_T1_preview.png b/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_T1_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..4478a91240ec5c23148819126dc19cba761b807a Binary files /dev/null and b/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_T1_preview.png differ diff --git a/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_T2.nii.gz b/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_T2.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..0c67cc72f21fa5c43b9a3d7136e64644a2ac466a --- /dev/null +++ b/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_T2.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9dd616c68eeb474a260b5cb98cdfe5d6fce9d6c2d28356fd8b4e35e0207f94d9 +size 1522248 diff --git a/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_T2_preview.png b/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_T2_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..822166d3ec8b6e8962ba70549509a17c023aac55 Binary files /dev/null and b/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_T2_preview.png differ diff --git a/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_meta.json b/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_meta.json new file mode 100644 index 0000000000000000000000000000000000000000..b5c550c5454f00a1d41598f64d5a338222a55cdc --- /dev/null +++ b/demo_samples/CN_vs_MCI/TMFP/sample_010/sample_010_meta.json @@ -0,0 +1,42 @@ +{ + "task": "CN_vs_MCI", + "modality_combination": "TMFP", + "modalities": [ + "T1", + "T2", + "Flair", + "PET" + ], + "modality_indices": [ + 1, + 1, + 1, + 1 + ], + "subject_id": "007_S_4516", + "sample_name": "sample_010", + "original_labels": { + "AGE": 73.4839151266256, + "MMSE": 28.0, + "DX": 1, + "Category": "CN" + }, + "task_label": { + "value": 0.0, + "type": "classification" + }, + "files": { + "nifti": { + "T1": "sample_010_T1.nii.gz", + "T2": "sample_010_T2.nii.gz", + "Flair": "sample_010_Flair.nii.gz", + "PET": "sample_010_PET.nii.gz" + }, + "previews": { + "T1": "sample_010_T1_preview.png", + "T2": "sample_010_T2_preview.png", + "Flair": "sample_010_Flair_preview.png", + "PET": "sample_010_PET_preview.png" + } + } +} \ No newline at end of file diff --git a/demo_samples/MMSE/T/sample_011/sample_011_T1.nii.gz b/demo_samples/MMSE/T/sample_011/sample_011_T1.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..97e8ee1670ca4d848972f6cba49108762b4e27ea --- /dev/null +++ b/demo_samples/MMSE/T/sample_011/sample_011_T1.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a490ea2db75a7fcd46cd8caf2d5d59be722036062265ccb705f79f404e68d22e +size 2099660 diff --git a/demo_samples/MMSE/T/sample_011/sample_011_T1_preview.png b/demo_samples/MMSE/T/sample_011/sample_011_T1_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..4478a91240ec5c23148819126dc19cba761b807a Binary files /dev/null and b/demo_samples/MMSE/T/sample_011/sample_011_T1_preview.png differ diff --git a/demo_samples/MMSE/T/sample_011/sample_011_meta.json b/demo_samples/MMSE/T/sample_011/sample_011_meta.json new file mode 100644 index 0000000000000000000000000000000000000000..23d31f9b4c362b2fb8a4c399b105049bcc56f5c0 --- /dev/null +++ b/demo_samples/MMSE/T/sample_011/sample_011_meta.json @@ -0,0 +1,33 @@ +{ + "task": "MMSE", + "modality_combination": "T", + "modalities": [ + "T1" + ], + "modality_indices": [ + 1, + 0, + 0, + 0 + ], + "subject_id": "007_S_4516", + "sample_name": "sample_011", + "original_labels": { + "AGE": 73.4839151266256, + "MMSE": 28.0, + "DX": 1, + "Category": "CN" + }, + "task_label": { + "value": 28.0, + "type": "regression" + }, + "files": { + "nifti": { + "T1": "sample_011_T1.nii.gz" + }, + "previews": { + "T1": "sample_011_T1_preview.png" + } + } +} \ No newline at end of file diff --git a/demo_samples/MMSE/TF/sample_012/sample_012_Flair.nii.gz b/demo_samples/MMSE/TF/sample_012/sample_012_Flair.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..19c33f314e5c8eb55f47a0c08e215ed7ba027d16 --- /dev/null +++ b/demo_samples/MMSE/TF/sample_012/sample_012_Flair.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:85c84cf50563a38c9ef91a9d5ab2943a4ce14193cce042157cd0f9acf90fb880 +size 1110239 diff --git a/demo_samples/MMSE/TF/sample_012/sample_012_Flair_preview.png b/demo_samples/MMSE/TF/sample_012/sample_012_Flair_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..64f9ad29dc116c91f2b6c68251d1fa44f040ea2f Binary files /dev/null and b/demo_samples/MMSE/TF/sample_012/sample_012_Flair_preview.png differ diff --git a/demo_samples/MMSE/TF/sample_012/sample_012_T1.nii.gz b/demo_samples/MMSE/TF/sample_012/sample_012_T1.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..97e8ee1670ca4d848972f6cba49108762b4e27ea --- /dev/null +++ b/demo_samples/MMSE/TF/sample_012/sample_012_T1.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a490ea2db75a7fcd46cd8caf2d5d59be722036062265ccb705f79f404e68d22e +size 2099660 diff --git a/demo_samples/MMSE/TF/sample_012/sample_012_T1_preview.png b/demo_samples/MMSE/TF/sample_012/sample_012_T1_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..4478a91240ec5c23148819126dc19cba761b807a Binary files /dev/null and b/demo_samples/MMSE/TF/sample_012/sample_012_T1_preview.png differ diff --git a/demo_samples/MMSE/TF/sample_012/sample_012_meta.json b/demo_samples/MMSE/TF/sample_012/sample_012_meta.json new file mode 100644 index 0000000000000000000000000000000000000000..ab242e8d3f298fa8b7385ef068a09e17a035c5b7 --- /dev/null +++ b/demo_samples/MMSE/TF/sample_012/sample_012_meta.json @@ -0,0 +1,36 @@ +{ + "task": "MMSE", + "modality_combination": "TF", + "modalities": [ + "T1", + "Flair" + ], + "modality_indices": [ + 1, + 0, + 1, + 0 + ], + "subject_id": "007_S_4516", + "sample_name": "sample_012", + "original_labels": { + "AGE": 73.4839151266256, + "MMSE": 28.0, + "DX": 1, + "Category": "CN" + }, + "task_label": { + "value": 28.0, + "type": "regression" + }, + "files": { + "nifti": { + "T1": "sample_012_T1.nii.gz", + "Flair": "sample_012_Flair.nii.gz" + }, + "previews": { + "T1": "sample_012_T1_preview.png", + "Flair": "sample_012_Flair_preview.png" + } + } +} \ No newline at end of file diff --git a/demo_samples/MMSE/TFP/sample_014/sample_014_Flair.nii.gz b/demo_samples/MMSE/TFP/sample_014/sample_014_Flair.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..19c33f314e5c8eb55f47a0c08e215ed7ba027d16 --- /dev/null +++ b/demo_samples/MMSE/TFP/sample_014/sample_014_Flair.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:85c84cf50563a38c9ef91a9d5ab2943a4ce14193cce042157cd0f9acf90fb880 +size 1110239 diff --git a/demo_samples/MMSE/TFP/sample_014/sample_014_Flair_preview.png b/demo_samples/MMSE/TFP/sample_014/sample_014_Flair_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..64f9ad29dc116c91f2b6c68251d1fa44f040ea2f Binary files /dev/null and b/demo_samples/MMSE/TFP/sample_014/sample_014_Flair_preview.png differ diff --git a/demo_samples/MMSE/TFP/sample_014/sample_014_PET.nii.gz b/demo_samples/MMSE/TFP/sample_014/sample_014_PET.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..30a260b830e45baad03aafee0adaa3bf499b617b --- /dev/null +++ b/demo_samples/MMSE/TFP/sample_014/sample_014_PET.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:546ec204487be18e75f950e214ffc8917aba2f7678c0db38a5a88f1547550d3c +size 2198379 diff --git a/demo_samples/MMSE/TFP/sample_014/sample_014_PET_preview.png b/demo_samples/MMSE/TFP/sample_014/sample_014_PET_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..f0547287323daa1a56e22447360295e6debe8d7c Binary files /dev/null and b/demo_samples/MMSE/TFP/sample_014/sample_014_PET_preview.png differ diff --git a/demo_samples/MMSE/TFP/sample_014/sample_014_T1.nii.gz b/demo_samples/MMSE/TFP/sample_014/sample_014_T1.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..97e8ee1670ca4d848972f6cba49108762b4e27ea --- /dev/null +++ b/demo_samples/MMSE/TFP/sample_014/sample_014_T1.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a490ea2db75a7fcd46cd8caf2d5d59be722036062265ccb705f79f404e68d22e +size 2099660 diff --git a/demo_samples/MMSE/TFP/sample_014/sample_014_T1_preview.png b/demo_samples/MMSE/TFP/sample_014/sample_014_T1_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..4478a91240ec5c23148819126dc19cba761b807a Binary files /dev/null and b/demo_samples/MMSE/TFP/sample_014/sample_014_T1_preview.png differ diff --git a/demo_samples/MMSE/TFP/sample_014/sample_014_meta.json b/demo_samples/MMSE/TFP/sample_014/sample_014_meta.json new file mode 100644 index 0000000000000000000000000000000000000000..199168c778e7b2f7f597ab526e1fea942fb23c77 --- /dev/null +++ b/demo_samples/MMSE/TFP/sample_014/sample_014_meta.json @@ -0,0 +1,39 @@ +{ + "task": "MMSE", + "modality_combination": "TFP", + "modalities": [ + "T1", + "Flair", + "PET" + ], + "modality_indices": [ + 1, + 0, + 1, + 1 + ], + "subject_id": "007_S_4516", + "sample_name": "sample_014", + "original_labels": { + "AGE": 73.4839151266256, + "MMSE": 28.0, + "DX": 1, + "Category": "CN" + }, + "task_label": { + "value": 28.0, + "type": "regression" + }, + "files": { + "nifti": { + "T1": "sample_014_T1.nii.gz", + "Flair": "sample_014_Flair.nii.gz", + "PET": "sample_014_PET.nii.gz" + }, + "previews": { + "T1": "sample_014_T1_preview.png", + "Flair": "sample_014_Flair_preview.png", + "PET": "sample_014_PET_preview.png" + } + } +} \ No newline at end of file diff --git a/demo_samples/MMSE/TMF/sample_013/sample_013_Flair.nii.gz b/demo_samples/MMSE/TMF/sample_013/sample_013_Flair.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..19c33f314e5c8eb55f47a0c08e215ed7ba027d16 --- /dev/null +++ b/demo_samples/MMSE/TMF/sample_013/sample_013_Flair.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:85c84cf50563a38c9ef91a9d5ab2943a4ce14193cce042157cd0f9acf90fb880 +size 1110239 diff --git a/demo_samples/MMSE/TMF/sample_013/sample_013_Flair_preview.png b/demo_samples/MMSE/TMF/sample_013/sample_013_Flair_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..64f9ad29dc116c91f2b6c68251d1fa44f040ea2f Binary files /dev/null and b/demo_samples/MMSE/TMF/sample_013/sample_013_Flair_preview.png differ diff --git a/demo_samples/MMSE/TMF/sample_013/sample_013_T1.nii.gz b/demo_samples/MMSE/TMF/sample_013/sample_013_T1.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..97e8ee1670ca4d848972f6cba49108762b4e27ea --- /dev/null +++ b/demo_samples/MMSE/TMF/sample_013/sample_013_T1.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a490ea2db75a7fcd46cd8caf2d5d59be722036062265ccb705f79f404e68d22e +size 2099660 diff --git a/demo_samples/MMSE/TMF/sample_013/sample_013_T1_preview.png b/demo_samples/MMSE/TMF/sample_013/sample_013_T1_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..4478a91240ec5c23148819126dc19cba761b807a Binary files /dev/null and b/demo_samples/MMSE/TMF/sample_013/sample_013_T1_preview.png differ diff --git a/demo_samples/MMSE/TMF/sample_013/sample_013_T2.nii.gz b/demo_samples/MMSE/TMF/sample_013/sample_013_T2.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..0c67cc72f21fa5c43b9a3d7136e64644a2ac466a --- /dev/null +++ b/demo_samples/MMSE/TMF/sample_013/sample_013_T2.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9dd616c68eeb474a260b5cb98cdfe5d6fce9d6c2d28356fd8b4e35e0207f94d9 +size 1522248 diff --git a/demo_samples/MMSE/TMF/sample_013/sample_013_T2_preview.png b/demo_samples/MMSE/TMF/sample_013/sample_013_T2_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..822166d3ec8b6e8962ba70549509a17c023aac55 Binary files /dev/null and b/demo_samples/MMSE/TMF/sample_013/sample_013_T2_preview.png differ diff --git a/demo_samples/MMSE/TMF/sample_013/sample_013_meta.json b/demo_samples/MMSE/TMF/sample_013/sample_013_meta.json new file mode 100644 index 0000000000000000000000000000000000000000..3de86a476cdc24b3e3797f6bb38df060f93f20d5 --- /dev/null +++ b/demo_samples/MMSE/TMF/sample_013/sample_013_meta.json @@ -0,0 +1,39 @@ +{ + "task": "MMSE", + "modality_combination": "TMF", + "modalities": [ + "T1", + "T2", + "Flair" + ], + "modality_indices": [ + 1, + 1, + 1, + 0 + ], + "subject_id": "007_S_4516", + "sample_name": "sample_013", + "original_labels": { + "AGE": 73.4839151266256, + "MMSE": 28.0, + "DX": 1, + "Category": "CN" + }, + "task_label": { + "value": 28.0, + "type": "regression" + }, + "files": { + "nifti": { + "T1": "sample_013_T1.nii.gz", + "T2": "sample_013_T2.nii.gz", + "Flair": "sample_013_Flair.nii.gz" + }, + "previews": { + "T1": "sample_013_T1_preview.png", + "T2": "sample_013_T2_preview.png", + "Flair": "sample_013_Flair_preview.png" + } + } +} \ No newline at end of file diff --git a/demo_samples/MMSE/TMFP/sample_015/sample_015_Flair.nii.gz b/demo_samples/MMSE/TMFP/sample_015/sample_015_Flair.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..19c33f314e5c8eb55f47a0c08e215ed7ba027d16 --- /dev/null +++ b/demo_samples/MMSE/TMFP/sample_015/sample_015_Flair.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:85c84cf50563a38c9ef91a9d5ab2943a4ce14193cce042157cd0f9acf90fb880 +size 1110239 diff --git a/demo_samples/MMSE/TMFP/sample_015/sample_015_Flair_preview.png b/demo_samples/MMSE/TMFP/sample_015/sample_015_Flair_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..64f9ad29dc116c91f2b6c68251d1fa44f040ea2f Binary files /dev/null and b/demo_samples/MMSE/TMFP/sample_015/sample_015_Flair_preview.png differ diff --git a/demo_samples/MMSE/TMFP/sample_015/sample_015_PET.nii.gz b/demo_samples/MMSE/TMFP/sample_015/sample_015_PET.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..30a260b830e45baad03aafee0adaa3bf499b617b --- /dev/null +++ b/demo_samples/MMSE/TMFP/sample_015/sample_015_PET.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:546ec204487be18e75f950e214ffc8917aba2f7678c0db38a5a88f1547550d3c +size 2198379 diff --git a/demo_samples/MMSE/TMFP/sample_015/sample_015_PET_preview.png b/demo_samples/MMSE/TMFP/sample_015/sample_015_PET_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..f0547287323daa1a56e22447360295e6debe8d7c Binary files /dev/null and b/demo_samples/MMSE/TMFP/sample_015/sample_015_PET_preview.png differ diff --git a/demo_samples/MMSE/TMFP/sample_015/sample_015_T1.nii.gz b/demo_samples/MMSE/TMFP/sample_015/sample_015_T1.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..97e8ee1670ca4d848972f6cba49108762b4e27ea --- /dev/null +++ b/demo_samples/MMSE/TMFP/sample_015/sample_015_T1.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a490ea2db75a7fcd46cd8caf2d5d59be722036062265ccb705f79f404e68d22e +size 2099660 diff --git a/demo_samples/MMSE/TMFP/sample_015/sample_015_T1_preview.png b/demo_samples/MMSE/TMFP/sample_015/sample_015_T1_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..4478a91240ec5c23148819126dc19cba761b807a Binary files /dev/null and b/demo_samples/MMSE/TMFP/sample_015/sample_015_T1_preview.png differ diff --git a/demo_samples/MMSE/TMFP/sample_015/sample_015_T2.nii.gz b/demo_samples/MMSE/TMFP/sample_015/sample_015_T2.nii.gz new file mode 100644 index 0000000000000000000000000000000000000000..0c67cc72f21fa5c43b9a3d7136e64644a2ac466a --- /dev/null +++ b/demo_samples/MMSE/TMFP/sample_015/sample_015_T2.nii.gz @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9dd616c68eeb474a260b5cb98cdfe5d6fce9d6c2d28356fd8b4e35e0207f94d9 +size 1522248 diff --git a/demo_samples/MMSE/TMFP/sample_015/sample_015_T2_preview.png b/demo_samples/MMSE/TMFP/sample_015/sample_015_T2_preview.png new file mode 100644 index 0000000000000000000000000000000000000000..822166d3ec8b6e8962ba70549509a17c023aac55 Binary files /dev/null and b/demo_samples/MMSE/TMFP/sample_015/sample_015_T2_preview.png differ diff --git a/demo_samples/MMSE/TMFP/sample_015/sample_015_meta.json b/demo_samples/MMSE/TMFP/sample_015/sample_015_meta.json new file mode 100644 index 0000000000000000000000000000000000000000..df79e15cbed38daae73e1c3786fbfde8029b7d67 --- /dev/null +++ b/demo_samples/MMSE/TMFP/sample_015/sample_015_meta.json @@ -0,0 +1,42 @@ +{ + "task": "MMSE", + "modality_combination": "TMFP", + "modalities": [ + "T1", + "T2", + "Flair", + "PET" + ], + "modality_indices": [ + 1, + 1, + 1, + 1 + ], + "subject_id": "007_S_4516", + "sample_name": "sample_015", + "original_labels": { + "AGE": 73.4839151266256, + "MMSE": 28.0, + "DX": 1, + "Category": "CN" + }, + "task_label": { + "value": 28.0, + "type": "regression" + }, + "files": { + "nifti": { + "T1": "sample_015_T1.nii.gz", + "T2": "sample_015_T2.nii.gz", + "Flair": "sample_015_Flair.nii.gz", + "PET": "sample_015_PET.nii.gz" + }, + "previews": { + "T1": "sample_015_T1_preview.png", + "T2": "sample_015_T2_preview.png", + "Flair": "sample_015_Flair_preview.png", + "PET": "sample_015_PET_preview.png" + } + } +} \ No newline at end of file diff --git a/demo_samples/samples.json b/demo_samples/samples.json new file mode 100644 index 0000000000000000000000000000000000000000..5b4ed94dfe02a5a41032b7cb862aabf9800c7551 --- /dev/null +++ b/demo_samples/samples.json @@ -0,0 +1,774 @@ +{ + "total_samples": 20, + "tasks": [ + "CN_vs_AD", + "CN_vs_MCI", + "MMSE", + "AGE" + ], + "modality_combinations": [ + "T", + "TF", + "TMF", + "TFP", + "TMFP" + ], + "samples": [ + { + "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" + } + } + }, + { + "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" + } + } + }, + { + "task": "CN_vs_AD", + "modality_combination": "TMF", + "modalities": [ + "T1", + "T2", + "Flair" + ], + "modality_indices": [ + 1, + 1, + 1, + 0 + ], + "subject_id": "007_S_4516", + "sample_name": "sample_003", + "original_labels": { + "AGE": 73.4839151266256, + "MMSE": 28.0, + "DX": 1, + "Category": "CN" + }, + "task_label": { + "value": 0.0, + "type": "classification" + }, + "files": { + "nifti": { + "T1": "sample_003_T1.nii.gz", + "T2": "sample_003_T2.nii.gz", + "Flair": "sample_003_Flair.nii.gz" + }, + "previews": { + "T1": "sample_003_T1_preview.png", + "T2": "sample_003_T2_preview.png", + "Flair": "sample_003_Flair_preview.png" + } + } + }, + { + "task": "CN_vs_AD", + "modality_combination": "TFP", + "modalities": [ + "T1", + "Flair", + "PET" + ], + "modality_indices": [ + 1, + 0, + 1, + 1 + ], + "subject_id": "007_S_4516", + "sample_name": "sample_004", + "original_labels": { + "AGE": 73.4839151266256, + "MMSE": 28.0, + "DX": 1, + "Category": "CN" + }, + "task_label": { + "value": 0.0, + "type": "classification" + }, + "files": { + "nifti": { + "T1": "sample_004_T1.nii.gz", + "Flair": "sample_004_Flair.nii.gz", + "PET": "sample_004_PET.nii.gz" + }, + "previews": { + "T1": "sample_004_T1_preview.png", + "Flair": "sample_004_Flair_preview.png", + "PET": "sample_004_PET_preview.png" + } + } + }, + { + "task": "CN_vs_AD", + "modality_combination": "TMFP", + "modalities": [ + "T1", + "T2", + "Flair", + "PET" + ], + "modality_indices": [ + 1, + 1, + 1, + 1 + ], + "subject_id": "007_S_4516", + "sample_name": "sample_005", + "original_labels": { + "AGE": 73.4839151266256, + "MMSE": 28.0, + "DX": 1, + "Category": "CN" + }, + "task_label": { + "value": 0.0, + "type": "classification" + }, + "files": { + "nifti": { + "T1": "sample_005_T1.nii.gz", + "T2": "sample_005_T2.nii.gz", + "Flair": "sample_005_Flair.nii.gz", + "PET": "sample_005_PET.nii.gz" + }, + "previews": { + "T1": "sample_005_T1_preview.png", + "T2": "sample_005_T2_preview.png", + "Flair": "sample_005_Flair_preview.png", + "PET": "sample_005_PET_preview.png" + } + } + }, + { + "task": "CN_vs_MCI", + "modality_combination": "T", + "modalities": [ + "T1" + ], + "modality_indices": [ + 1, + 0, + 0, + 0 + ], + "subject_id": "007_S_4516", + "sample_name": "sample_006", + "original_labels": { + "AGE": 73.4839151266256, + "MMSE": 28.0, + "DX": 1, + "Category": "CN" + }, + "task_label": { + "value": 0.0, + "type": "classification" + }, + "files": { + "nifti": { + "T1": "sample_006_T1.nii.gz" + }, + "previews": { + "T1": "sample_006_T1_preview.png" + } + } + }, + { + "task": "CN_vs_MCI", + "modality_combination": "TF", + "modalities": [ + "T1", + "Flair" + ], + "modality_indices": [ + 1, + 0, + 1, + 0 + ], + "subject_id": "007_S_4516", + "sample_name": "sample_007", + "original_labels": { + "AGE": 73.4839151266256, + "MMSE": 28.0, + "DX": 1, + "Category": "CN" + }, + "task_label": { + "value": 0.0, + "type": "classification" + }, + "files": { + "nifti": { + "T1": "sample_007_T1.nii.gz", + "Flair": "sample_007_Flair.nii.gz" + }, + "previews": { + "T1": "sample_007_T1_preview.png", + "Flair": "sample_007_Flair_preview.png" + } + } + }, + { + "task": "CN_vs_MCI", + "modality_combination": "TMF", + "modalities": [ + "T1", + "T2", + "Flair" + ], + "modality_indices": [ + 1, + 1, + 1, + 0 + ], + "subject_id": "007_S_4516", + "sample_name": "sample_008", + "original_labels": { + "AGE": 73.4839151266256, + "MMSE": 28.0, + "DX": 1, + "Category": "CN" + }, + "task_label": { + "value": 0.0, + "type": "classification" + }, + "files": { + "nifti": { + "T1": "sample_008_T1.nii.gz", + "T2": "sample_008_T2.nii.gz", + "Flair": "sample_008_Flair.nii.gz" + }, + "previews": { + "T1": "sample_008_T1_preview.png", + "T2": "sample_008_T2_preview.png", + "Flair": "sample_008_Flair_preview.png" + } + } + }, + { + "task": "CN_vs_MCI", + "modality_combination": "TFP", + "modalities": [ + "T1", + "Flair", + "PET" + ], + "modality_indices": [ + 1, + 0, + 1, + 1 + ], + "subject_id": "007_S_4516", + "sample_name": "sample_009", + "original_labels": { + "AGE": 73.4839151266256, + "MMSE": 28.0, + "DX": 1, + "Category": "CN" + }, + "task_label": { + "value": 0.0, + "type": "classification" + }, + "files": { + "nifti": { + "T1": "sample_009_T1.nii.gz", + "Flair": "sample_009_Flair.nii.gz", + "PET": "sample_009_PET.nii.gz" + }, + "previews": { + "T1": "sample_009_T1_preview.png", + "Flair": "sample_009_Flair_preview.png", + "PET": "sample_009_PET_preview.png" + } + } + }, + { + "task": "CN_vs_MCI", + "modality_combination": "TMFP", + "modalities": [ + "T1", + "T2", + "Flair", + "PET" + ], + "modality_indices": [ + 1, + 1, + 1, + 1 + ], + "subject_id": "007_S_4516", + "sample_name": "sample_010", + "original_labels": { + "AGE": 73.4839151266256, + "MMSE": 28.0, + "DX": 1, + "Category": "CN" + }, + "task_label": { + "value": 0.0, + "type": "classification" + }, + "files": { + "nifti": { + "T1": "sample_010_T1.nii.gz", + "T2": "sample_010_T2.nii.gz", + "Flair": "sample_010_Flair.nii.gz", + "PET": "sample_010_PET.nii.gz" + }, + "previews": { + "T1": "sample_010_T1_preview.png", + "T2": "sample_010_T2_preview.png", + "Flair": "sample_010_Flair_preview.png", + "PET": "sample_010_PET_preview.png" + } + } + }, + { + "task": "MMSE", + "modality_combination": "T", + "modalities": [ + "T1" + ], + "modality_indices": [ + 1, + 0, + 0, + 0 + ], + "subject_id": "007_S_4516", + "sample_name": "sample_011", + "original_labels": { + "AGE": 73.4839151266256, + "MMSE": 28.0, + "DX": 1, + "Category": "CN" + }, + "task_label": { + "value": 28.0, + "type": "regression" + }, + "files": { + "nifti": { + "T1": "sample_011_T1.nii.gz" + }, + "previews": { + "T1": "sample_011_T1_preview.png" + } + } + }, + { + "task": "MMSE", + "modality_combination": "TF", + "modalities": [ + "T1", + "Flair" + ], + "modality_indices": [ + 1, + 0, + 1, + 0 + ], + "subject_id": "007_S_4516", + "sample_name": "sample_012", + "original_labels": { + "AGE": 73.4839151266256, + "MMSE": 28.0, + "DX": 1, + "Category": "CN" + }, + "task_label": { + "value": 28.0, + "type": "regression" + }, + "files": { + "nifti": { + "T1": "sample_012_T1.nii.gz", + "Flair": "sample_012_Flair.nii.gz" + }, + "previews": { + "T1": "sample_012_T1_preview.png", + "Flair": "sample_012_Flair_preview.png" + } + } + }, + { + "task": "MMSE", + "modality_combination": "TMF", + "modalities": [ + "T1", + "T2", + "Flair" + ], + "modality_indices": [ + 1, + 1, + 1, + 0 + ], + "subject_id": "007_S_4516", + "sample_name": "sample_013", + "original_labels": { + "AGE": 73.4839151266256, + "MMSE": 28.0, + "DX": 1, + "Category": "CN" + }, + "task_label": { + "value": 28.0, + "type": "regression" + }, + "files": { + "nifti": { + "T1": "sample_013_T1.nii.gz", + "T2": "sample_013_T2.nii.gz", + "Flair": "sample_013_Flair.nii.gz" + }, + "previews": { + "T1": "sample_013_T1_preview.png", + "T2": "sample_013_T2_preview.png", + "Flair": "sample_013_Flair_preview.png" + } + } + }, + { + "task": "MMSE", + "modality_combination": "TFP", + "modalities": [ + "T1", + "Flair", + "PET" + ], + "modality_indices": [ + 1, + 0, + 1, + 1 + ], + "subject_id": "007_S_4516", + "sample_name": "sample_014", + "original_labels": { + "AGE": 73.4839151266256, + "MMSE": 28.0, + "DX": 1, + "Category": "CN" + }, + "task_label": { + "value": 28.0, + "type": "regression" + }, + "files": { + "nifti": { + "T1": "sample_014_T1.nii.gz", + "Flair": "sample_014_Flair.nii.gz", + "PET": "sample_014_PET.nii.gz" + }, + "previews": { + "T1": "sample_014_T1_preview.png", + "Flair": "sample_014_Flair_preview.png", + "PET": "sample_014_PET_preview.png" + } + } + }, + { + "task": "MMSE", + "modality_combination": "TMFP", + "modalities": [ + "T1", + "T2", + "Flair", + "PET" + ], + "modality_indices": [ + 1, + 1, + 1, + 1 + ], + "subject_id": "007_S_4516", + "sample_name": "sample_015", + "original_labels": { + "AGE": 73.4839151266256, + "MMSE": 28.0, + "DX": 1, + "Category": "CN" + }, + "task_label": { + "value": 28.0, + "type": "regression" + }, + "files": { + "nifti": { + "T1": "sample_015_T1.nii.gz", + "T2": "sample_015_T2.nii.gz", + "Flair": "sample_015_Flair.nii.gz", + "PET": "sample_015_PET.nii.gz" + }, + "previews": { + "T1": "sample_015_T1_preview.png", + "T2": "sample_015_T2_preview.png", + "Flair": "sample_015_Flair_preview.png", + "PET": "sample_015_PET_preview.png" + } + } + }, + { + "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" + } + } + }, + { + "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" + } + } + }, + { + "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" + } + } + }, + { + "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" + } + } + }, + { + "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" + } + } + } + ] +} \ No newline at end of file diff --git a/inference_engine.py b/inference_engine.py new file mode 100644 index 0000000000000000000000000000000000000000..9af6648b9c2999dc3b45a059e4929ccf2dff0cf6 --- /dev/null +++ b/inference_engine.py @@ -0,0 +1,562 @@ +#!/usr/bin/env python3 +""" +BrainAnytime Inference Engine for Hugging Face Space + +提供模型加载、推理和结果处理功能。 +支持4个任务：CN vs AD, CN vs MCI, MMSE, AGE +支持5种模态组合：T, TF, TMF, TFP, TMFP +""" + +import os +import sys +import warnings +from pathlib import Path +from typing import Dict, List, Tuple, Optional, Union +import json + +import numpy as np +import torch +import torch.nn as nn +import nibabel as nib + +warnings.filterwarnings("ignore") + +# 添加 BrainAnytime 代码路径 +BASE_DIR = Path(__file__).parent +BRAINANYTIME_DIR = BASE_DIR / "BrainAnytime" +if BRAINANYTIME_DIR.exists(): + sys.path.insert(0, str(BRAINANYTIME_DIR)) + +# 导入 BrainAnytime 模型 +from models.multimae3d import create_multimae3d + + +# ============================================================================= +# 配置常量 +# ============================================================================= + +# 模态配置 +MODALITY_ORDER = ['T1', 'T2', 'Flair', 'PET'] +MODALITY_SHORT = {'T1': 'T', 'T2': 'M', 'Flair': 'F', 'PET': 'P'} +SHORT_TO_FULL = {'T': 'T1', 'M': 'T2', 'F': 'Flair', 'P': 'PET'} + +# 任务配置 +TASKS = { + 'CN_vs_AD': { + 'type': 'classification', + 'display_name': 'CN vs AD', + 'num_classes': 2, + 'classes': ['CN', 'AD'], + }, + 'CN_vs_MCI': { + 'type': 'classification', + 'display_name': 'CN vs MCI', + 'num_classes': 2, + 'classes': ['CN', 'MCI'], + }, + 'MMSE': { + 'type': 'regression', + 'display_name': 'MMSE Score', + 'min': 10.0, + 'max': 30.0, + }, + 'AGE': { + 'type': 'regression', + 'display_name': 'Age', + 'min': 50.0, + 'max': 100.0, + }, +} + +# 模型默认参数（与训练时一致） +DEFAULT_MODEL_ARGS = { + 'img_size': 128, + 'patch_size': 16, + 'embed_dim': 768, + 'depth': 12, + 'num_heads': 12, + 'decoder_embed_dim': 512, + 'decoder_depth': 8, + 'decoder_num_heads': 16, + 'pool': 'mean', + 'dropout': 0.5, +} + +# 归一化参数 +MMSE_MIN, MMSE_MAX = 10.0, 30.0 + + +# ============================================================================= +# 下游任务模型头 +# ============================================================================= + +class MultiMAE3DForDownstream(nn.Module): + """ + MultiMAE3D + 下游任务头 (分类/回归) + 从 BrainAnytime/finetune_main.py 复制 + """ + def __init__( + self, + encoder, + embed_dim: int = 768, + num_outputs: int = 1, + pool: str = 'mean', + dropout: float = 0.5, + ): + super().__init__() + self.encoder = encoder + self.pool = pool + self.num_outputs = num_outputs + + # 预测头 + self.head = nn.Sequential( + nn.Dropout(dropout), + nn.Linear(embed_dim, num_outputs) + ) + + def forward(self, images, observed, mc=None): + """ + Args: + images: [B, 4, D, H, W] - 4 modalities + observed: [B, 4] - 0/1 mask for available modalities + mc: [B] - modality combination index (optional) + + Returns: + logits: [B, num_outputs] + """ + # Encode + x = self.encoder.forward_encoder(images, observed) + + # Pool + if self.pool == 'mean': + x = x.mean(dim=1) # [B, embed_dim] + elif self.pool == 'cls': + x = x[:, 0] # Use CLS token + else: + x = x[:, 0] + + # Head + logits = self.head(x) # [B, num_outputs] + + return logits + + +# ============================================================================= +# 推理引擎 +# ============================================================================= + +class BrainAnytimeInference: + """ + BrainAnytime 推理引擎 + + 用法: + engine = BrainAnytimeInference(checkpoints_dir) + engine.load_model('CN_vs_AD') + result = engine.predict(nifti_data, modalities) + """ + + def __init__( + self, + checkpoints_dir: Optional[str] = None, + device: Optional[str] = None, + ): + """ + Args: + checkpoints_dir: 模型检查点目录，默认从 HF Hub 下载 + device: 'cuda' 或 'cpu'，默认自动选择 + """ + self.checkpoints_dir = checkpoints_dir + self.device = torch.device( + device if device else ('cuda' if torch.cuda.is_available() else 'cpu') + ) + print(f"Inference device: {self.device}") + + # 缓存加载的模型 + self.models: Dict[str, MultiMAE3DForDownstream] = {} + + def _get_checkpoint_path(self, task: str) -> str: + """获取任务对应的 checkpoint 路径""" + checkpoint_files = { + 'CN_vs_AD': 'CN_vs_AD_seed_0_best.pth', + 'CN_vs_MCI': 'CN_vs_MCI_seed_0_best.pth', + 'MMSE': 'MMSE_seed_0_best.pth', + 'AGE': 'AGE_seed_0_best.pth', + } + + if task not in checkpoint_files: + raise ValueError(f"Unknown task: {task}. Available: {list(checkpoint_files.keys())}") + + if self.checkpoints_dir: + path = os.path.join(self.checkpoints_dir, checkpoint_files[task]) + if os.path.exists(path): + return path + + # 从 Hugging Face Hub 下载 + try: + from huggingface_hub import hf_hub_download + path = hf_hub_download( + repo_id="Simmonstt/BrainAnytime", + filename=checkpoint_files[task], + ) + return path + except Exception as e: + raise RuntimeError( + f"Cannot load checkpoint for {task}. " + f"Please provide checkpoints_dir or ensure HF Hub access. Error: {e}" + ) + + def load_model(self, task: str) -> MultiMAE3DForDownstream: + """ + 加载指定任务的模型（懒加载） + + Args: + task: 任务名称 ('CN_vs_AD', 'CN_vs_MCI', 'MMSE', 'AGE') + + Returns: + 加载好的模型 + """ + if task in self.models: + return self.models[task] + + print(f"Loading model for task: {task}") + + # 创建编码器 + encoder = create_multimae3d( + img_size=DEFAULT_MODEL_ARGS['img_size'], + patch_size=DEFAULT_MODEL_ARGS['patch_size'], + embed_dim=DEFAULT_MODEL_ARGS['embed_dim'], + depth=DEFAULT_MODEL_ARGS['depth'], + num_heads=DEFAULT_MODEL_ARGS['num_heads'], + decoder_embed_dim=DEFAULT_MODEL_ARGS['decoder_embed_dim'], + decoder_depth=DEFAULT_MODEL_ARGS['decoder_depth'], + decoder_num_heads=DEFAULT_MODEL_ARGS['decoder_num_heads'], + ) + + # 创建下游任务模型 + model = MultiMAE3DForDownstream( + encoder=encoder, + embed_dim=DEFAULT_MODEL_ARGS['embed_dim'], + num_outputs=1, + pool=DEFAULT_MODEL_ARGS['pool'], + dropout=DEFAULT_MODEL_ARGS['dropout'], + ).to(self.device) + + # 加载 checkpoint + checkpoint_path = self._get_checkpoint_path(task) + print(f" Loading checkpoint: {checkpoint_path}") + + ckpt = torch.load(checkpoint_path, map_location=self.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', '?')})") + + model.eval() + self.models[task] = model + + return model + + def preprocess_nifti( + self, + nifti_path: str, + target_size: Tuple[int, int, int] = (128, 128, 128), + ) -> Optional[np.ndarray]: + """ + 加载并预处理 NIfTI 文件 + + Args: + nifti_path: NIfTI 文件路径 + target_size: 目标尺寸 (D, H, W) + + Returns: + 预处理后的数据 [D, H, W]，失败返回 None + """ + try: + # 加载 + nii = nib.load(nifti_path) + data = nii.get_fdata().astype(np.float32) + + # 确保 3D + if data.ndim == 4: + data = data[..., 0] + + # 检查尺寸 + if data.shape != target_size: + print(f" Warning: Size mismatch {data.shape} != {target_size}") + # 如果需要，可以在这里添加 resize 逻辑 + # 但目前假设输入已经是 128x128x128 + return None + + # Min-Max 归一化 + data_min, data_max = data.min(), data.max() + if data_max > data_min: + data = (data - data_min) / (data_max - data_min) + else: + print(f" Warning: Empty image (max == min)") + return None + + return data + + except Exception as e: + print(f" Error loading {nifti_path}: {e}") + return None + + def prepare_input( + self, + nifti_files: Dict[str, str], + modality_combo: str, + ) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: + """ + 准备模型输入 + + Args: + nifti_files: {模态名: 文件路径}，如 {'T1': 'path/to/T1.nii.gz'} + modality_combo: 模态组合字符串，如 'T', 'TF' + + Returns: + (images, observed) 或 None + - images: [1, 4, 128, 128, 128] + - observed: [1, 4] + """ + expected_modalities = [SHORT_TO_FULL[c] for c in modality_combo] + + # 加载所有模态 + images_list = [] + observed_list = [] + + for mod in MODALITY_ORDER: + if mod in expected_modalities and mod in nifti_files: + data = self.preprocess_nifti(nifti_files[mod]) + if data is not None: + images_list.append(data) + observed_list.append(1.0) + else: + return None + else: + # 缺失模态用零填充 + images_list.append(np.zeros((128, 128, 128), dtype=np.float32)) + observed_list.append(0.0) + + # 转换为张量 + images = np.stack(images_list, axis=0) # [4, 128, 128, 128] + images = torch.from_numpy(images).unsqueeze(0).to(self.device) # [1, 4, 128, 128, 128] + + observed = torch.tensor(observed_list, dtype=torch.float32).unsqueeze(0).to(self.device) + + return images, observed + + @torch.no_grad() + def predict( + self, + nifti_files: Dict[str, str], + task: str, + modality_combo: str, + ) -> Optional[Dict]: + """ + 执行推理 + + Args: + nifti_files: {模态名: 文件路径} + task: 任务名称 + modality_combo: 模态组合字符串 + + Returns: + 推理结果字典 + """ + # 加载模型 + model = self.load_model(task) + task_config = TASKS[task] + + # 准备输入 + input_data = self.prepare_input(nifti_files, modality_combo) + if input_data is None: + return None + + images, observed = input_data + + # 推理 + logits = model(images, observed) + + # 后处理 + if task_config['type'] == 'classification': + # 分类：sigmoid → probability + prob = torch.sigmoid(logits).item() + pred_class = 1 if prob > 0.5 else 0 + pred_label = task_config['classes'][pred_class] + confidence = prob if pred_class == 1 else 1 - prob + + result = { + 'task': task, + 'task_type': 'classification', + 'prediction': pred_label, + 'probability': prob, + 'confidence': confidence, + 'classes': task_config['classes'], + 'logits': logits.item(), + } + else: + # 回归：反归一化 + normalized_pred = logits.item() + + if task == 'MMSE': + # 反归一化到 [10, 30] + pred = normalized_pred * (MMSE_MAX - MMSE_MIN) + MMSE_MIN + pred = max(MMSE_MIN, min(MMSE_MAX, pred)) + unit = 'points' + reference_range = [MMSE_MIN, MMSE_MAX] + else: # AGE + # 假设 AGE 是 z-score，这里简化处理 + pred = normalized_pred * 20 + 75 # 近似反归一化 + unit = 'years' + reference_range = [50, 100] + + result = { + 'task': task, + 'task_type': 'regression', + 'prediction': pred, + 'unit': unit, + 'reference_range': reference_range, + 'normalized_value': normalized_pred, + 'logits': logits.item(), + } + + # 添加输入信息 + result['input'] = { + 'modality_combo': modality_combo, + 'modalities': [SHORT_TO_FULL[c] for c in modality_combo], + 'files': nifti_files, + } + + return result + + def predict_from_sample( + self, + sample_dir: str, + task: str, + modality_combo: str, + ) -> Optional[Dict]: + """ + 从样本目录执行推理 + + Args: + sample_dir: 样本目录路径 (包含 .nii.gz 文件) + task: 任务名称 + modality_combo: 模态组合 + + Returns: + 推理结果 + """ + sample_dir = Path(sample_dir) + + # 查找 NIfTI 文件 + expected_modalities = [SHORT_TO_FULL[c] for c in modality_combo] + nifti_files = {} + + for mod in expected_modalities: + # 查找匹配的文件 + pattern = f"*{mod}*.nii.gz" + matches = list(sample_dir.glob(pattern)) + + if not matches: + # 尝试更宽松的匹配 + pattern = f"*.nii.gz" + for f in sample_dir.glob(pattern): + if mod.lower() in f.name.lower(): + matches.append(f) + break + + if matches: + nifti_files[mod] = str(matches[0]) + else: + print(f" Error: Cannot find {mod} in {sample_dir}") + return None + + return self.predict(nifti_files, task, modality_combo) + + +# ============================================================================= +# 辅助函数 +# ============================================================================= + +def create_inference_engine(checkpoints_dir: Optional[str] = None) -> BrainAnytimeInference: + """ + 创建推理引擎实例 + + Args: + checkpoints_dir: 检查点目录，默认从 HuggingFace Hub 下载 + + Returns: + BrainAnytimeInference 实例 + """ + return BrainAnytimeInference(checkpoints_dir=checkpoints_dir) + + +def format_result(result: Dict) -> str: + """格式化推理结果为可读字符串""" + if result is None: + return "Inference failed" + + task = result['task'] + task_type = result['task_type'] + + lines = [ + f"Task: {task}", + f"Type: {task_type}", + ] + + if task_type == 'classification': + lines.extend([ + f"Prediction: {result['prediction']}", + f"Probability: {result['probability']:.4f}", + f"Confidence: {result['confidence']:.2%}", + ]) + else: + lines.extend([ + f"Prediction: {result['prediction']:.2f} {result['unit']}", + f"Reference Range: {result['reference_range']}", + ]) + + lines.append(f"Modalities: {', '.join(result['input']['modalities'])}") + + return '\n'.join(lines) + + +# ============================================================================= +# 测试 +# ============================================================================= + +if __name__ == '__main__': + import argparse + + parser = argparse.ArgumentParser(description='Test inference engine') + parser.add_argument('--checkpoints_dir', type=str, + default='/home/23037125r/code/random/multimae_freeze_then_finetune/', + help='Checkpoints directory') + parser.add_argument('--sample_dir', type=str, + default='/home/23037125r/code/Downstream_tasks/hf_space/demo_samples/CN_vs_AD/TMFP/sample_005', + help='Sample directory for testing') + parser.add_argument('--task', type=str, default='CN_vs_AD', + choices=list(TASKS.keys()), + help='Task to test') + parser.add_argument('--combo', type=str, default='TMFP', + choices=['T', 'TF', 'TMF', 'TFP', 'TMFP'], + help='Modality combination') + + args = parser.parse_args() + + # 创建引擎 + engine = create_inference_engine(args.checkpoints_dir) + + # 执行推理 + print(f"\nTesting inference:") + print(f" Task: {args.task}") + print(f" Combo: {args.combo}") + print(f" Sample: {args.sample_dir}") + + result = engine.predict_from_sample(args.sample_dir, args.task, args.combo) + + if result: + print("\nResult:") + print(format_result(result)) + else: + print("\nInference failed!") diff --git a/requirements.txt b/requirements.txt index a3655ed536f54c844bba83c4ae4199af23ade61a..35413cee726e74e8dc726e89880883d38cf63022 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,30 @@ -# Space runtime: Gradio is pre-installed by Hugging Face (see README sdk_version). -# Training/eval dependencies are listed in requirements-train.txt for local use. +# BrainAnytime Hugging Face Space Requirements + +# Core dependencies +# Note: Gradio is pre-installed by Hugging Face (version controlled via README.md sdk_version) +torch>=2.0.0 +torchvision>=0.15.0 +numpy>=1.24.0 + +# Medical imaging +nibabel>=4.0.0 +torchio>=0.18.95 + +# Model dependencies (from BrainAnytime) +timm>=0.9.0 +einops>=0.6.0 + +# Hugging Face Hub for model download +huggingface-hub>=0.16.0 + +# Data processing +pandas>=2.0.0 +scikit-learn>=1.3.0 +scipy>=1.10.0 + +# Visualization +matplotlib>=3.7.0 +Pillow>=10.0.0 + +# Optional: for advanced features +# tensorboardX>=2.6 # Only needed if logging