BrainAnytime/downstream_dataloader.py

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)