医学图像分割实战：如何用Boundary Loss损失函数解决边缘模糊问题

1次阅读

共计 2506 个字符，预计需要花费 7 分钟才能阅读完成。

医学图像分割的质量直接影响后续诊断和分析的准确性，但传统损失函数在处理边缘细节时常常力不从心。交叉熵损失函数（Cross-Entropy Loss）主要关注像素级别的分类准确性，而 Dice Loss 虽然能够处理类别不平衡问题，但它们都存在一个共同的缺陷：无法直接优化分割边界的位置。

交叉熵的局限：过度关注像素级分类，容易导致边界模糊，特别是当前景和背景像素分布不均匀时。
Dice Loss 的问题：虽然对类别不平衡更鲁棒，但仍然是通过区域重叠率来优化，无法直接捕捉边界距离的差异。

这两种损失函数在医学图像分割中常常导致边缘不连续、细节丢失，尤其是对于小目标（如细胞、血管）的分割任务。

Boundary Loss 的核心思想是通过计算预测分割边界与真实边界之间的距离来优化模型。其数学表达如下：

$$
L_{boundary} = \sum_{p \in \partial G} \phi_S(p) + \sum_{p \in \partial S} \phi_G(p)
$$

其中：
– (\partial G) 和 (\partial S) 分别是真实边界和预测边界。
– (\phi_S(p)) 是预测边界到真实边界的距离变换（Distance Transform），表示预测边界点 (p) 到最近真实边界点的距离。
– (\phi_G(p)) 是类似的距离变换，表示真实边界点到预测边界的距离。

通过最小化这个损失函数，模型能够直接优化边界位置，而不是像 Dice Loss 那样仅仅关注区域重叠。

以下是 Boundary Loss 的 PyTorch 实现代码，关键步骤已添加注释：

import torch
import torch.nn as nn
import torch.nn.functional as F
from scipy.ndimage import distance_transform_edt

class BoundaryLoss(nn.Module):
    def __init__(self, alpha=0.5):
        super(BoundaryLoss, self).__init__()
        self.alpha = alpha  # 权重系数，用于混合损失

    def forward(self, pred, target):
        # 输入 pred 为模型输出的概率图，target 为二值化的真实标签
        # 1. 将 pred 通过 sigmoid 激活，并二值化为 0 / 1 掩码
        pred_sigmoid = torch.sigmoid(pred)
        pred_mask = (pred_sigmoid > 0.5).float()

        # 2. 计算真实边界和预测边界的距离变换
        # 注意：距离变换需要在 CPU 上计算（PyTorch 暂无原生支持）target_np = target.cpu().numpy()
        pred_mask_np = pred_mask.detach().cpu().numpy()

        # 计算真实边界的距离变换（到预测边界）phi_target = torch.from_numpy(distance_transform_edt(1 - target_np)).float().to(pred.device)
        # 计算预测边界的距离变换（到真实边界）phi_pred = torch.from_numpy(distance_transform_edt(1 - pred_mask_np)).float().to(pred.device)

        # 3. 计算边界损失
        loss_target = (target * phi_pred).mean()
        loss_pred = (pred_mask * phi_target).mean()
        loss = loss_target + loss_pred

        return loss * self.alpha  # 返回加权后的损失

为了直观展示 Boundary Loss 的效果，我们在 ISBI 细胞分割数据集上对比了 Dice Loss 和 Boundary Loss 的分割结果：