医学图像分割实战：Boundary Loss损失函数原理与PyTorch实现

1次阅读

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

在医学图像分割任务中，边界模糊问题一直是影响分割精度的重要因素。传统交叉熵损失函数和 Dice Loss 在处理这类问题时存在明显局限性。本文将详细介绍 Boundary Loss 的原理及其在 PyTorch 中的实现，帮助开发者更好地理解和应用这一损失函数。

医学图像分割中，边界模糊问题主要由以下几个因素引起：

医学成像设备的固有噪声
组织边界的自然模糊
部分容积效应

传统交叉熵损失函数仅关注像素级别的分类准确性，而忽略了边界区域的几何特性。Dice Loss 虽然在一定程度上考虑了区域重叠，但对边界的敏感性仍然不足。

与 Hausdorff Distance 和 Active Contour 等边界优化方法相比，Boundary Loss 具有以下优势：

计算效率更高
更容易与深度学习框架集成
对噪声和初始边界位置的鲁棒性更好

Boundary Loss 的核心思想是通过距离变换图来优化分割边界。其数学表达式为：

$$L_{boundary} = \sum_{p\in\Omega} \phi_G(p) \cdot S(p)$$

其中：
– $\phi_G(p)$ 表示像素 p 到真实边界 G 的距离变换
– $S(p)$ 表示预测的分割概率图

距离变换图的计算公式为：

$$\phi_G(p) = \begin{cases}
0 & p\in G \
d(p,G) & p\in R_G \
-d(p,G) & p\notin R_G
\end{cases}$$

以下是 Boundary Loss 的完整 PyTorch 实现代码：

import torch
import torch.nn as nn
import numpy as np
from scipy.ndimage import distance_transform_edt

class BoundaryLoss(nn.Module):
    """Boundary Loss for medical image segmentation"""
    def __init__(self):
        super(BoundaryLoss, self).__init__()

    def one_hot2dist(self, seg):
        """
        Convert one-hot encoded segmentation to distance transform
        Args:
            seg: (B, C, H, W) one-hot encoded segmentation
        Returns:
            distance transform map
        """
        res = np.zeros_like(seg)
        for b in range(seg.shape[0]):
            for c in range(seg.shape[1]):
                posmask = seg[b, c].astype(np.bool)
                if posmask.any():
                    negmask = ~posmask
                    res[b, c] = distance_transform_edt(negmask) * negmask \
                              - (distance_transform_edt(posmask) - 1) * posmask
        return res

    def forward(self, pred, target):
        """
        Args:
            pred: (B, C, H, W) predicted probability map
            target: (B, C, H, W) one-hot encoded ground truth
        Returns:
            boundary loss
        """
        # Convert target to distance transform
        target_dist = torch.from_numpy(self.one_hot2dist(target.cpu().numpy())).float().to(pred.device)

        # Calculate boundary loss
        loss = (pred * target_dist).mean()

        return loss

在 BraTS2018 数据集上的实验结果表明：