import numpy as np from tqdm import tqdm from torch.cuda.amp import autocast as autocast import torch import torch.nn.functional as F from sklearn.metrics import confusion_matrix from scipy.ndimage.morphology import binary_fill_holes, binary_opening from utils import test_single_volume import time def calculate_dice_per_class(pred, target, num_classes): """计算每个类别的dice系数""" dice_scores = [] pred_argmax = torch.argmax(pred, dim=1) for cls in range(1, num_classes): # 跳过背景类 pred_cls = (pred_argmax == cls).float() target_cls = (target == cls).float() intersection = (pred_cls * target_cls).sum() total = pred_cls.sum() + target_cls.sum() if total > 0: dice = (2.0 * intersection) / total dice_scores.append(dice.item()) # 这里需要.item()因为dice是tensor else: dice = 1.0 if intersection == 0 else 0.0 dice_scores.append(dice) # 这里不需要.item()因为dice已经是float return dice_scores def calculate_miou(pred, target, num_classes): """计算mIoU""" pred_argmax = torch.argmax(pred, dim=1) iou_scores = [] for cls in range(1, num_classes): # 跳过背景类 pred_cls = (pred_argmax == cls).float() target_cls = (target == cls).float() intersection = (pred_cls * target_cls).sum() union = pred_cls.sum() + target_cls.sum() - intersection if union > 0: iou = intersection / union iou_scores.append(iou.item()) # 这里需要.item()因为iou是tensor else: iou = 1.0 if intersection == 0 else 0.0 iou_scores.append(iou) # 这里不需要.item()因为iou已经是float return np.mean(iou_scores) def train_one_epoch(train_loader, model, criterion, optimizer, scheduler, epoch, logger, config, scaler=None): ''' train model for one epoch ''' stime = time.time() model.train() loss_list = [] dice_scores_all = [] miou_scores_all = [] for iter, data in enumerate(train_loader): optimizer.zero_grad() images, targets = data['image'], data['label'] images, targets = images.cuda(non_blocking=True).float(), targets.cuda(non_blocking=True).long() if config.amp: with autocast(): model_output = model(images) # 解包输出 - 处理多种模型输出格式 if isinstance(model_output, tuple): if len(model_output) == 3: # SMAFormerV2 with deep supervision: (seg_out, boundary_out, ds_outs) out, boundary_out, ds_outs = model_output # TODO: 可以添加boundary和deep supervision的损失 loss = criterion(out, targets) elif len(model_output) == 2: # SMAFormerV2 with boundary_head: (seg_out, boundary_out) # 或 其他模型: (out, intermediate_results) out, second_output = model_output # 检查是否支持中间监督 if hasattr(criterion, 'forward') and 'intermediate_preds' in criterion.forward.__code__.co_varnames: loss = criterion(out, targets, intermediate_preds=second_output) else: loss = criterion(out, targets) else: out = model_output[0] loss = criterion(out, targets) else: out = model_output loss = criterion(out, targets) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() else: model_output = model(images) # 解包输出 - 处理多种模型输出格式 if isinstance(model_output, tuple): if len(model_output) == 3: # SMAFormerV2 with deep supervision: (seg_out, boundary_out, ds_outs) out, boundary_out, ds_outs = model_output # TODO: 可以添加boundary和deep supervision的损失 loss = criterion(out, targets) elif len(model_output) == 2: # SMAFormerV2 with boundary_head: (seg_out, boundary_out) # 或 其他模型: (out, intermediate_results) out, second_output = model_output # 检查是否支持中间监督 if hasattr(criterion, 'forward') and 'intermediate_preds' in criterion.forward.__code__.co_varnames: loss = criterion(out, targets, intermediate_preds=second_output) else: loss = criterion(out, targets) else: out = model_output[0] loss = criterion(out, targets) else: out = model_output loss = criterion(out, targets) loss.backward() optimizer.step() loss_list.append(loss.item()) # 计算训练时的dice和mIoU with torch.no_grad(): dice_scores = calculate_dice_per_class(out, targets, config.num_classes) miou_score = calculate_miou(out, targets, config.num_classes) dice_scores_all.append(dice_scores) miou_scores_all.append(miou_score) now_lr = optimizer.state_dict()['param_groups'][0]['lr'] mean_loss = np.mean(loss_list) if iter % config.print_interval == 0 and iter != 0: log_info = f'train: epoch {epoch}, iter:{iter}, loss: {loss.item():.4f}, lr: {now_lr}' print(log_info) logger.info(log_info) scheduler.step() # 计算平均指标 mean_loss = np.mean(loss_list) mean_dice_per_class = np.mean(dice_scores_all, axis=0) mean_dice_avg = np.mean(mean_dice_per_class) mean_miou = np.mean(miou_scores_all) etime = time.time() log_info = f'Finish one epoch train: epoch {epoch}, loss: {mean_loss:.4f}, avg_dice: {mean_dice_avg:.4f}, mIoU: {mean_miou:.4f}, time(s): {etime-stime:.2f}' print(log_info) logger.info(log_info) return { 'loss': mean_loss, 'avg_dice': mean_dice_avg, 'dice_per_class': mean_dice_per_class, 'miou': mean_miou } def val_one_epoch(test_datasets, test_loader, model, epoch, logger, config, test_save_path, val_or_test=False): # switch to evaluate mode stime = time.time() model.eval() with torch.no_grad(): metric_list = 0.0 i_batch = 0 all_dice_scores = [] all_hd95_scores = [] for data in tqdm(test_loader): img, msk, case_name = data['image'], data['label'], data['case_name'][0] metric_i = test_single_volume(img, msk, model, classes=config.num_classes, patch_size=[config.input_size_h, config.input_size_w], test_save_path=test_save_path, case=case_name, z_spacing=config.z_spacing, val_or_test=val_or_test) metric_list += np.array(metric_i) # 收集每个case的dice和hd95分数 case_dice_scores = [metric[0] for metric in metric_i] # 每个类别的dice case_hd95_scores = [metric[1] for metric in metric_i] # 每个类别的hd95 all_dice_scores.append(case_dice_scores) all_hd95_scores.append(case_hd95_scores) logger.info('idx %d case %s mean_dice %f mean_hd95 %f' % (i_batch, case_name, np.mean(metric_i, axis=0)[0], np.mean(metric_i, axis=0)[1])) i_batch += 1 metric_list = metric_list / len(test_datasets) # 计算平均指标 avg_dice_per_class = np.mean(all_dice_scores, axis=0) # 每个类别的平均dice avg_hd95_per_class = np.mean(all_hd95_scores, axis=0) # 每个类别的平均hd95 performance = np.mean(avg_dice_per_class) # 所有类别的平均dice mean_hd95 = np.mean(avg_hd95_per_class) # 所有类别的平均hd95 # 计算mIoU (基于dice分数的近似) mean_miou = np.mean([dice / (2 - dice) if dice < 1.0 else 1.0 for dice in avg_dice_per_class]) for i in range(len(avg_dice_per_class)): logger.info('Mean class %d mean_dice %f mean_hd95 %f' % (i+1, avg_dice_per_class[i], avg_hd95_per_class[i])) etime = time.time() log_info = f'val epoch: {epoch}, mean_dice: {performance:.4f}, mean_hd95: {mean_hd95:.4f}, mIoU: {mean_miou:.4f}, time(s): {etime-stime:.2f}' print(log_info) logger.info(log_info) return { 'avg_dice': performance, 'dice_per_class': avg_dice_per_class, 'avg_hd95': mean_hd95, 'hd95_per_class': avg_hd95_per_class, 'miou': mean_miou } def val_one_epoch_slice(val_dataset, val_loader, model, epoch, logger, config, test_save_path, val_or_test=False): """ 使用slice-by-slice方式进行验证，与训练时的数据处理保持一致 """ stime = time.time() model.eval() all_dice_scores = [] all_miou_scores = [] with torch.no_grad(): for data in tqdm(val_loader, desc=f"Validation Epoch {epoch}"): images, targets = data['image'], data['label'] images = images.cuda(non_blocking=True).float() targets = targets.cuda(non_blocking=True).long() # 前向传播 model_output = model(images) # 解包输出 - AFFSegNet返回两个值 if isinstance(model_output, tuple) and len(model_output) == 2: out, intermediate_results = model_output else: out = model_output # 计算dice和mIoU dice_scores = calculate_dice_per_class(out, targets, config.num_classes) miou_score = calculate_miou(out, targets, config.num_classes) all_dice_scores.append(dice_scores) all_miou_scores.append(miou_score) # 计算平均指标 avg_dice_per_class = np.mean(all_dice_scores, axis=0) # 每个类别的平均dice avg_dice = np.mean(avg_dice_per_class) # 总体平均dice avg_miou = np.mean(all_miou_scores) # 平均mIoU # HD95暂时设为0，因为slice-by-slice计算HD95不太合理 avg_hd95_per_class = np.zeros_like(avg_dice_per_class) avg_hd95 = 0.0 # 打印每个类别的性能 organ_names = ['spleen', 'right_kidney', 'left_kidney', 'gallbladder', 'esophagus', 'liver', 'stomach', 'aorta'] for i, organ in enumerate(organ_names): if i < len(avg_dice_per_class): logger.info(f'Validation class {i+1} ({organ}): dice={avg_dice_per_class[i]:.4f}') etime = time.time() log_info = f'val_slice epoch: {epoch}, mean_dice: {avg_dice:.4f}, mIoU: {avg_miou:.4f}, slices: {len(all_dice_scores)}, time(s): {etime-stime:.2f}' print(log_info) logger.info(log_info) return { 'avg_dice': avg_dice, 'dice_per_class': avg_dice_per_class, 'avg_hd95': avg_hd95, 'hd95_per_class': avg_hd95_per_class, 'miou': avg_miou } def val_one_epoch_with_visualization(test_datasets, test_loader, model, epoch, logger, config, test_save_path, prediction_vis_dir=None, attention_vis_dir=None, activation_vis_dir=None, val_or_test=False, save_vis_every_n=5, save_prediction_comparison_func=None, extract_attention_maps_func=None, save_attention_heatmaps_func=None, extract_activation_maps_func=None, save_activation_heatmaps_func=None): """ 验证函数，支持生成预测对比图、注意力热图和激活热图 Args: save_vis_every_n: 每n个案例保存一次可视化图像 save_prediction_comparison_func: 保存预测对比图的函数 extract_attention_maps_func: 提取注意力图的函数 save_attention_heatmaps_func: 保存注意力热图的函数 extract_activation_maps_func: 提取激活图的函数 save_activation_heatmaps_func: 保存激活热图的函数 activation_vis_dir: 激活热图保存目录 """ stime = time.time() model.eval() with torch.no_grad(): metric_list = 0.0 i_batch = 0 for data in tqdm(test_loader): img, msk, case_name = data['image'], data['label'], data['case_name'][0] # 执行预测并计算指标 metric_i = test_single_volume(img, msk, model, classes=config.num_classes, patch_size=[config.input_size_h, config.input_size_w], test_save_path=test_save_path, case=case_name, z_spacing=config.z_spacing, val_or_test=val_or_test) metric_list += np.array(metric_i) # 每隔save_vis_every_n个案例生成可视化图像 if i_batch % save_vis_every_n == 0 and (prediction_vis_dir or attention_vis_dir or activation_vis_dir): try: # 生成预测结果用于可视化 img_np = img.squeeze(0).cpu().detach().numpy() msk_np = msk.squeeze(0).cpu().detach().numpy() # 创建模型预测 # 根据配置推断期望输入通道数 desired_in_ch = getattr(config, 'input_channels', None) if hasattr(config, 'model_config') and isinstance(config.model_config, dict): desired_in_ch = config.model_config.get('input_channels', desired_in_ch) if desired_in_ch is None: desired_in_ch = 1 def to_input_tensor(slice_img_np): # (H, W) -> 构造模型期望的 (1, C, H, W) t = torch.from_numpy(slice_img_np).unsqueeze(0).float().cuda() # (1, H, W) if desired_in_ch == 3: t = t.repeat(3, 1, 1).unsqueeze(0) # (1,H,W)->(3,H,W)->(1,3,H,W) else: t = t.unsqueeze(0) # (1,H,W)->(1,1,H,W) return t if len(img_np.shape) == 3: # 3D数据 prediction_np = np.zeros_like(msk_np) # 为3D数据生成注意力热图和激活热图 - 完整的3D数据处理 if attention_vis_dir and extract_attention_maps_func and save_attention_heatmaps_func: try: # 使用完整的3D数据生成注意力热图 mid_slice = img_np[img_np.shape[0] // 2] x, y = mid_slice.shape[0], mid_slice.shape[1] if x != config.input_size_h or y != config.input_size_w: from scipy.ndimage import zoom mid_slice = zoom(mid_slice, (config.input_size_h / x, config.input_size_w / y), order=3) input_tensor = to_input_tensor(mid_slice) attention_maps = extract_attention_maps_func(model, input_tensor) if attention_maps: save_attention_heatmaps_func( img_np, attention_maps, case_name, attention_vis_dir ) except Exception as e: logger.warning(f"生成注意力热图失败 {case_name}: {e}") # 为3D数据生成激活热图 if activation_vis_dir and extract_activation_maps_func and save_activation_heatmaps_func: try: # 使用完整的3D数据生成激活热图 mid_slice = img_np[img_np.shape[0] // 2] x, y = mid_slice.shape[0], mid_slice.shape[1] if x != config.input_size_h or y != config.input_size_w: from scipy.ndimage import zoom mid_slice = zoom(mid_slice, (config.input_size_h / x, config.input_size_w / y), order=3) input_tensor = to_input_tensor(mid_slice) activation_maps = extract_activation_maps_func(model, input_tensor) if activation_maps: save_activation_heatmaps_func( img_np, activation_maps, case_name, activation_vis_dir ) except Exception as e: logger.warning(f"生成激活热图失败 {case_name}: {e}") # 生成每个slice的预测 for ind in range(img_np.shape[0]): slice_img = img_np[ind, :, :] x, y = slice_img.shape[0], slice_img.shape[1] # 调整尺寸 if x != config.input_size_h or y != config.input_size_w: from scipy.ndimage import zoom slice_img = zoom(slice_img, (config.input_size_h / x, config.input_size_w / y), order=3) # 根据期望通道数构造输入 input_tensor = to_input_tensor(slice_img) # 获取预测结果 outputs = model(input_tensor) pred = torch.argmax(torch.softmax(outputs, dim=1), dim=1).squeeze(0).cpu().detach().numpy() # 调整回原尺寸 if x != config.input_size_h or y != config.input_size_w: pred = zoom(pred, (x / config.input_size_h, y / config.input_size_w), order=0) prediction_np[ind] = pred else: # 2D数据处理 input_tensor = to_input_tensor(img_np) outputs = model(input_tensor) prediction_np = torch.argmax(torch.softmax(outputs, dim=1), dim=1).squeeze(0).cpu().detach().numpy() # 生成注意力热图 if attention_vis_dir and extract_attention_maps_func and save_attention_heatmaps_func: try: attention_maps = extract_attention_maps_func(model, input_tensor) if attention_maps: save_attention_heatmaps_func( img_np, attention_maps, case_name, attention_vis_dir ) except Exception as e: logger.warning(f"生成注意力热图失败 {case_name}: {e}") # 生成激活热图 if extract_activation_maps_func and save_activation_heatmaps_func: try: activation_maps = extract_activation_maps_func(model, input_tensor) if activation_maps: save_activation_heatmaps_func( img_np, activation_maps, case_name, activation_vis_dir ) except Exception as e: logger.warning(f"生成激活热图失败 {case_name}: {e}") # 保存预测对比图 if prediction_vis_dir and save_prediction_comparison_func: try: save_prediction_comparison_func( img_np, msk_np, prediction_np, case_name, prediction_vis_dir ) logger.info(f"已保存预测对比图: {case_name}") except Exception as e: logger.warning(f"保存预测对比图失败 {case_name}: {e}") except Exception as e: logger.warning(f"可视化处理失败 {case_name}: {e}") logger.info('idx %d case %s mean_dice %f mean_hd95 %f' % (i_batch, case_name, np.mean(metric_i, axis=0)[0], np.mean(metric_i, axis=0)[1])) i_batch += 1 metric_list = metric_list / len(test_datasets) performance = np.mean(metric_list, axis=0)[0] mean_hd95 = np.mean(metric_list, axis=0)[1] # 提取每个器官的dice和hd95 dice_per_organ = [] hd95_per_organ = [] for i in range(1, config.num_classes): dice_score = metric_list[i-1][0] hd95_score = metric_list[i-1][1] dice_per_organ.append(dice_score) hd95_per_organ.append(hd95_score) logger.info('Mean class %d mean_dice %f mean_hd95 %f' % (i, dice_score, hd95_score)) etime = time.time() log_info = f'val epoch: {epoch}, mean_dice: {performance}, mean_hd95: {mean_hd95}, time(s): {etime-stime:.2f}' print(log_info) logger.info(log_info) return performance, mean_hd95, dice_per_organ, hd95_per_organ