nnU-Net：自配置医学分割框架，让UNet发挥极致

def extract_fingerprint(dataset):
    fingerprint = {
        # 基本信息
        'modality': get_modality(dataset),  # CT, MRI, etc.
        'num_classes': count_classes(dataset),
        'num_samples': len(dataset),
        
        # 空间特性
        'median_spacing': np.median([img.spacing for img in dataset], axis=0),
        'median_shape': np.median([img.shape for img in dataset], axis=0),
        'size_reduction': compute_size_reduction_by_cropping(dataset),
        
        # 强度特性
        'intensity_properties': {
            'mean': np.mean(...),
            'std': np.std(...),
            'percentiles': np.percentile(..., [0.5, 50, 99.5])
        },
        
        # 类别特性
        'class_locations': get_class_locations(dataset),
        'foreground_ratio': compute_foreground_ratio(dataset)
    }
    return fingerprint

class UNet2D(nn.Module):
    """标准2D UNet"""
    def __init__(self, in_channels, num_classes, base_num_features=32):
        super().__init__()
        
        # 编码器
        self.conv1 = StackedConvLayers(in_channels, base_num_features)
        self.conv2 = StackedConvLayers(base_num_features, base_num_features * 2)
        self.conv3 = StackedConvLayers(base_num_features * 2, base_num_features * 4)
        self.conv4 = StackedConvLayers(base_num_features * 4, base_num_features * 8)
        
        # Bottleneck
        self.bottleneck = StackedConvLayers(base_num_features * 8, base_num_features * 16)
        
        # 解码器（省略）
        # ...
    
    def forward(self, x):
        # ... 标准UNet前向传播
        pass


class StackedConvLayers(nn.Module):
    """nnU-Net的标准卷积块"""
    def __init__(self, in_ch, out_ch, kernel_size=3, num_convs=2):
        super().__init__()
        layers = []
        for i in range(num_convs):
            layers.append(nn.Conv2d(
                in_ch if i == 0 else out_ch,
                out_ch,
                kernel_size,
                padding=kernel_size // 2
            ))
            layers.append(nn.InstanceNorm2d(out_ch))  # 使用Instance Norm
            layers.append(nn.LeakyReLU(0.01, inplace=True))
        
        self.convs = nn.Sequential(*layers)
    
    def forward(self, x):
        return self.convs(x)

class UNet3D(nn.Module):
    """3D UNet（全分辨率）"""
    def __init__(self, in_channels, num_classes, base_num_features=32):
        super().__init__()
        
        # 编码器（3D卷积）
        self.conv1 = StackedConvLayers3D(in_channels, base_num_features)
        self.pool1 = nn.Conv3d(base_num_features, base_num_features, 
                               kernel_size=3, stride=2, padding=1)
        
        # ... 更深的编码器
        
    def forward(self, x):
        # 3D前向传播
        pass

class Cascade3D:
    def __init__(self):
        self.stage1 = UNet3D(...)  # 低分辨率
        self.stage2 = UNet3D(...)  # 高分辨率
    
    def predict(self, image):
        # 阶段1：下采样预测
        image_lowres = downsample(image, factor=2)
        mask_lowres = self.stage1(image_lowres)
        mask_upsampled = upsample(mask_lowres, target_size=image.shape)
        
        # 阶段2：精细化
        input_stage2 = torch.cat([image, mask_upsampled], dim=1)
        mask_final = self.stage2(input_stage2)
        
        return mask_final

def determine_patch_size(median_shape, median_spacing, target_spacing):
    """根据图像尺寸和spacing确定patch size"""
    
    # 目标：使用约128×128×128的patch（GPU内存平衡）
    reference_patch = np.array([128, 128, 128])
    
    # 计算重采样后的尺寸
    new_shape = np.round(median_shape * median_spacing / target_spacing).astype(int)
    
    # 调整patch size，确保不超过图像尺寸
    patch_size = np.minimum(reference_patch, new_shape)
    
    # 确保是2的倍数（方便pooling）
    patch_size = np.round(patch_size / 16) * 16
    
    return patch_size.astype(int)

# 示例
median_shape = [512, 512, 130]
median_spacing = [0.7, 0.7, 5.0]
target_spacing = [1.0, 1.0, 1.0]  # 重采样到1mm

patch_size = determine_patch_size(median_shape, median_spacing, target_spacing)
# 输出：[192, 192, 128]

def determine_batch_size(network_type, patch_size, gpu_memory=11):
    """根据GPU内存和patch size确定batch size"""
    
    # 估算单个样本的内存占用
    if network_type == '2D':
        memory_per_sample = patch_size[0] * patch_size[1] * 4 / (1024**2)  # MB
        max_batch_size = 12
    elif network_type == '3D':
        memory_per_sample = np.prod(patch_size) * 4 / (1024**2)
        max_batch_size = 2
    
    # 根据GPU内存调整
    available_memory = gpu_memory * 1024  # GB -> MB
    estimated_batch_size = int(available_memory * 0.6 / memory_per_sample)
    
    batch_size = min(estimated_batch_size, max_batch_size)
    batch_size = max(batch_size, 2)  # 至少2
    
    return batch_size

def get_augmentation_pipeline(dataset_properties):
    """自适应数据增强"""
    
    transforms = []
    
    # 基础增强（总是使用）
    transforms.append(A.RandomRotate90(p=0.5))
    transforms.append(A.Flip(p=0.5))
    
    # 弹性变形（3D数据）
    if dataset_properties['is_3d']:
        transforms.append(A.ElasticTransform(
            alpha=30,
            sigma=5,
            p=0.3
        ))
    
    # 旋转（根据模态调整）
    if dataset_properties['modality'] == 'CT':
        # CT对旋转敏感度低
        transforms.append(A.Rotate(limit=30, p=0.5))
    elif dataset_properties['modality'] == 'MRI':
        # MRI更敏感
        transforms.append(A.Rotate(limit=15, p=0.3))
    
    # 强度增强
    transforms.append(A.RandomBrightnessContrast(p=0.3))
    transforms.append(A.RandomGamma(p=0.3))
    
    # 缩放（根据spacing各向异性）
    anisotropy = np.max(dataset_properties['spacing']) / np.min(dataset_properties['spacing'])
    if anisotropy > 2:
        # 各向异性大，限制Z轴缩放
        transforms.append(A.RandomScale(scale_limit=(0.7, 1.3), p=0.2))
    else:
        # 各向同性，自由缩放
        transforms.append(A.RandomScale(scale_limit=(0.5, 1.5), p=0.5))
    
    return A.Compose(transforms)

# nnU-Net自动集成多个模型
predictions = []

# 3种配置
for config in ['2d', '3d_fullres', '3d_cascade']:
    # 5折交叉验证
    for fold in range(5):
        model = load_model(f"nnUNet/{config}/fold_{fold}")
        pred = model.predict(image)
        predictions.append(pred)

# 平均15个预测（3配置 × 5折）
final_pred = np.mean(predictions, axis=0)

from nnunet.training.network_training.nnUNetTrainerV2 import nnUNetTrainerV2

class CustomTrainer(nnUNetTrainerV2):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        
        # 自定义学习率
        self.initial_lr = 5e-4
        
        # 自定义损失函数权重
        self.loss = DC_and_CE_loss({'batch_dice': True, 'smooth': 1e-5}, {})
    
    def initialize_optimizer_and_scheduler(self):
        """自定义优化器"""
        self.optimizer = torch.optim.SGD(
            self.network.parameters(),
            self.initial_lr,
            weight_decay=self.weight_decay,
            momentum=0.99,
            nesterov=True
        )
        
        self.lr_scheduler = torch.optim.lr_scheduler.PolyLR(
            self.optimizer,
            self.num_epochs,
            power=0.9
        )

# nnU-Net自动学习后处理策略
# 在验证集上测试不同后处理的效果

def custom_postprocessing(pred_mask, remove_small_lesions_threshold=100):
    """移除小连通域"""
    import cc3d
    
    # 连通域分析
    labels_out = cc3d.connected_components(pred_mask, connectivity=26)
    
    # 计算每个连通域的大小
    stats = cc3d.statistics(labels_out)
    
    # 移除小连通域
    for label_id in range(1, stats['voxel_counts'].shape[0]):
        if stats['voxel_counts'][label_id] < remove_small_lesions_threshold:
            pred_mask[labels_out == label_id] = 0
    
    return pred_mask

class ImbalancedTrainer(nnUNetTrainerV2):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
    
    def get_train_transforms(self):
        """增强小类别的采样"""
        transforms = super().get_train_transforms()
        
        # 添加类别平衡采样
        transforms.append(OverSampleForegroundClasses(
            classes_to_oversample=[2],  # 肿瘤类别
            oversample_factor=2.0
        ))
        
        return transforms