医学图像分割常用指标及代码（PyTorch实现）

一、引言

医学图像分割是计算机辅助诊断的核心技术，其评估指标直接决定模型性能的量化准确性。不同于自然图像分割，医学场景对分割边界的精确性、小目标检测能力有更高要求。本文将系统梳理Dice系数、IoU、Hausdorff距离等核心指标，提供基于PyTorch的高效实现方案，并分析各指标的适用场景与优化方向。

二、核心评估指标体系

1. Dice系数（F1-score）

数学定义：
Dice = (2 * |X ∩ Y|) / (|X| + |Y|)
其中X为预测掩码，Y为真实标签，取值范围[0,1]，1表示完美分割。

PyTorch实现：

import torch
def dice_coeff(pred: torch.Tensor, target: torch.Tensor, smooth=1e-6):
    """
    Args:
        pred: 模型输出logits或概率图 (B,C,H,W)
        target: 真实标签 (B,H,W) 或 (B,C,H,W) one-hot
        smooth: 数值稳定性常数
    Returns:
        dice系数张量 (B,)
    """
    if len(target.shape) == 3:  # (B,H,W) -> (B,C,H,W)
        target = torch.nn.functional.one_hot(target.long(), num_classes=2).permute(0,3,1,2).float()
    pred_flat = pred.softmax(dim=1)[:,1:,...].contiguous().view(-1)  # 取前景通道
    target_flat = target[:,1:,...].contiguous().view(-1)
    intersection = (pred_flat * target_flat).sum()
    union = pred_flat.sum() + target_flat.sum()
    return (2. * intersection + smooth) / (union + smooth)

优化技巧：

• 对类别不平衡数据，可采用加权Dice：
  weighted_dice = (β * dice_fg + (1-β) * dice_bg)
• 多类别扩展时，建议计算各类别Dice的均值（mDice）

2. 交并比（IoU/Jaccard指数）

数学定义：
IoU = |X ∩ Y| / |X ∪ Y|
与Dice呈正相关，但IoU对边界误差更敏感。

PyTorch实现：

def iou_score(pred: torch.Tensor, target: torch.Tensor, smooth=1e-6):
    """
    Args:
        pred: 二值化预测图 (B,H,W)
        target: 真实标签 (B,H,W)
    """
    pred_flat = pred.flatten()
    target_flat = target.flatten()
    intersection = (pred_flat * target_flat).sum()
    union = pred_flat.sum() + target_flat.sum() - intersection
    return (intersection + smooth) / (union + smooth)

应用场景：

• 适用于需要严格边界约束的任务（如器官轮廓分割）
• 在U-Net等模型中常作为辅助损失函数

3. Hausdorff距离（HD）

数学定义：
HD(X,Y) = max{sup x∈X inf y∈Y d(x,y), sup y∈Y inf x∈X d(x,y)}
衡量两组点集间的最大不匹配距离，对异常值敏感。

PyTorch实现优化：

def hausdorff_distance(pred: torch.Tensor, target: torch.Tensor, spacing=1.0):
    """
    Args:
        pred: 二值预测图 (1,H,W)
        target: 真实标签 (1,H,W)
        spacing: 物理空间分辨率（mm/pixel）
    Returns:
        HD95（95%分位数Hausdorff距离）
    """
    from scipy.spatial.distance import cdist
    import numpy as np
    # 提取轮廓点
    def get_contour_points(mask):
        from skimage.measure import find_contours
        contours = find_contours(mask.squeeze().cpu().numpy(), 0.5)
        points = []
        for cnt in contours:
            points.extend(cnt * spacing)  # 转换为物理坐标
        return np.array(points) if points else np.zeros((0,2))
    pred_points = get_contour_points(pred)
    target_points = get_contour_points(target)
    if len(pred_points) == 0 or len(target_points) == 0:
        return torch.tensor(float('inf'))
    # 计算距离矩阵（CPU计算更高效）
    dist_matrix = cdist(pred_points, target_points)
    hd1 = np.max(np.min(dist_matrix, axis=1))
    hd2 = np.max(np.min(dist_matrix, axis=0))
    hd_full = max(hd1, hd2)
    # 计算95%分位数（替代最大值以减少异常值影响）
    sorted_dists = np.sort(np.concatenate([
        np.min(dist_matrix, axis=1),
        np.min(dist_matrix, axis=0)
    ]))
    hd95 = sorted_dists[int(len(sorted_dists)*0.95)]
    return torch.tensor(hd95)

工程建议：

• 使用HD95替代原始HD以增强鲁棒性
• 对3D体积数据，建议先在各切片计算HD再取均值

4. 表面距离指标

ASD（平均表面距离）：

def average_surface_distance(pred, target, spacing=1.0):
    from medpy.metric.binary import hd, asd
    # 需安装medpy库：pip install medpy
    return asd(pred.squeeze().cpu().numpy(), 
               target.squeeze().cpu().numpy(), 
               voxelspacing=spacing)

适用场景：

• 评估分割表面与真实表面的平均偏差
• 特别适用于需要精确测量器官体积的任务

三、评估体系构建实践

1. 多指标融合评估

class SegmentationMetrics:
    def __init__(self, num_classes=2):
        self.num_classes = num_classes
        self.dice_scores = []
        self.iou_scores = []
        self.hd95_scores = []
    def update(self, pred: torch.Tensor, target: torch.Tensor):
        # pred: (B,C,H,W) logits
        # target: (B,H,W) labels
        prob = torch.softmax(pred, dim=1)
        pred_mask = torch.argmax(prob, dim=1)
        for b in range(pred.shape[0]):
            # 计算各类别Dice
            for c in range(self.num_classes):
                pred_c = (pred_mask[b] == c).float()
                target_c = (target[b] == c).float()
                dice = dice_coeff(pred_c.unsqueeze(0), target_c.unsqueeze(0))
                self.dice_scores.append(dice.item())
                if c > 0:  # 通常只计算前景类别的IoU/HD
                    iou = iou_score(pred_c.unsqueeze(0), target_c.unsqueeze(0))
                    hd95 = hausdorff_distance(pred_c.unsqueeze(0), target_c.unsqueeze(0))
                    self.iou_scores.append(iou.item())
                    self.hd95_scores.append(hd95.item())
    def compute(self):
        metrics = {
            'mean_dice': np.mean(self.dice_scores),
            'mean_iou': np.mean(self.iou_scores) if self.iou_scores else 0,
            'mean_hd95': np.mean(self.hd95_scores) if self.hd95_scores else float('inf')
        }
        return metrics

2. 3D医学图像评估优化

对于CT/MRI体积数据，建议采用分层评估：

def volumetric_metrics(pred_vol: torch.Tensor, 
                      target_vol: torch.Tensor, 
                      slice_spacing: float = 1.0):
    """
    Args:
        pred_vol: (D,H,W) 预测体积
        target_vol: (D,H,W) 真实标签
        slice_spacing: 层间物理间距（mm）
    """
    metrics = {'dice': [], 'hd95': []}
    for d in range(pred_vol.shape[0]):
        pred_slice = pred_vol[d].unsqueeze(0)
        target_slice = target_vol[d].unsqueeze(0)
        dice = dice_coeff(pred_slice, target_slice)
        hd95 = hausdorff_distance(pred_slice, target_slice, spacing=slice_spacing)
        metrics['dice'].append(dice.item())
        metrics['hd95'].append(hd95.item())
    # 计算加权平均（考虑层间物理距离）
    weights = torch.linspace(0, 1, pred_vol.shape[0])
    weighted_dice = np.average(metrics['dice'], weights=weights)
    weighted_hd95 = np.average(metrics['hd95'], weights=weights)
    return {
        'mean_dice': np.mean(metrics['dice']),
        'weighted_dice': weighted_dice,
        'mean_hd95': np.mean(metrics['hd95']),
        'weighted_hd95': weighted_hd95
    }

四、工程优化建议

1. 内存管理：

   • 对大体积数据，采用分块计算策略
   • 使用torch.cuda.amp进行混合精度计算

2. 计算加速：

   • 将距离计算移至CPU（使用NumPy/SciPy）
   • 对HD计算，可采用近似算法（如基于KDTree的快速搜索）

3. 可视化验证：

   import matplotlib.pyplot as plt
   def plot_segmentation(img, pred, target):
       fig, (ax1, ax2, ax3) = plt.subplots(1,3, figsize=(15,5))
       ax1.imshow(img, cmap='gray')
       ax1.set_title('Input Image')
       ax2.imshow(pred.squeeze(), cmap='jet')
       ax2.set_title('Prediction')
       ax3.imshow(target.squeeze(), cmap='jet')
       ax3.set_title('Ground Truth')
       plt.show()

4. 多GPU评估：

   def distributed_evaluate(model, dataloader, device):
       model.eval()
       metrics = SegmentationMetrics()
       with torch.no_grad():
           for batch in dataloader:
               images, targets = batch
               images = images.to(device)
               targets = targets.to(device)
               outputs = model(images)
               metrics.update(outputs, targets)
       # 使用torch.distributed进行全局同步（需初始化进程组）
       if torch.distributed.is_initialized():
           # 实现聚合逻辑
           pass
       return metrics.compute()

五、结论

医学图像分割评估需要构建多维度指标体系，Dice系数适合整体分割质量评估，IoU对边界敏感，Hausdorff距离关注极端误差。在实际工程中，建议：

1. 根据任务特点选择主评估指标（如器官分割优先Dice，肿瘤分割关注HD）
2. 实现分层评估（2D切片/3D体积）
3. 结合可视化工具进行定性验证
4. 采用混合精度计算优化性能

本文提供的PyTorch实现方案经过工程优化，可直接集成至医学图像分析流水线，为模型迭代提供可靠的量化依据。