深入理解Boundary损失函数：从原理到实战中的图像分割优化

1次阅读

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

在图像分割任务中，边界区域的精确识别一直是难点。传统损失函数如交叉熵（Cross-Entropy, CE）和 Dice 损失存在明显局限性：

交叉熵损失：逐像素计算分类误差，但对边界像素和内部像素一视同仁，导致模型更关注易分类的大面积区域。
Dice 损失：虽然能缓解类别不平衡问题，但对边界轻微偏移惩罚不足，容易产生 ” 模糊边缘 ”。

这种现象在医学图像（如器官分割）和自动驾驶（如道路边缘检测）中尤为明显——边界精度直接决定应用效果。

Boundary 损失的核心思想是将边界误差转化为空间距离的度量。给定真实边界 $\partial G$ 和预测边界 $\partial S$，定义距离变换图 $\phi_G(p)$ 表示像素 $p$ 到 $\partial G$ 的最近距离：

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

损失函数则定义为预测区域 $S$ 与真实区域 $G$ 的边界距离差异：

$$\mathcal{L}{boundary} = \int \phi_G(p) \cdot s_p \, dp$$

其中 $s_p$ 是像素 $p$ 的预测概率。这个设计使得：

当预测边界在真实边界外侧（$\phi_G(p)<0$），增加 $s_p$ 会增大损失
当预测边界在真实边界内侧（$\phi_G(p)>0$），减少 $s_p$ 会增大损失

# 依赖库
import torch
import numpy as np
import cv2  # OpenCV 4.5+

def get_distance_map(mask: np.ndarray) -> np.ndarray:
    """
    生成二值 mask 的距离变换图
    Args:
        mask: [H,W] 值为 0 / 1 的 numpy 数组
    Returns:
        [H,W]距离变换图，边界处为 0，内部为正，外部为负
    """
    # 寻找轮廓（OpenCV 版本差异处理）contours, _ = cv2.findContours(mask.astype(np.uint8), 
        cv2.RETR_EXTERNAL, 
        cv2.CHAIN_APPROX_SIMPLE
    )

    # 生成边界图
    boundary = np.zeros_like(mask)
    cv2.drawContours(boundary, contours, -1, 1, thickness=1)

    # 计算距离变换（内部用正距离，外部用负距离）dist_inner = cv2.distanceTransform(mask, cv2.DIST_L2, 3)
    dist_outer = -cv2.distanceTransform(1-mask, cv2.DIST_L2, 3)

    return np.where(boundary>0, 0, dist_inner + dist_outer)

class BoundaryLoss(torch.nn.Module):
    def __init__(self, theta=10.0):
        super().__init__()
        self.theta = theta  # 控制距离变换的缩放系数

    def forward(self, pred: torch.Tensor, gt_dist: torch.Tensor):
        """
        Args:
            pred: [B,C,H,W] 模型输出的概率图
            gt_dist: [B,H,W] 预处理好的距离变换图
        """
        # 对多类别情况取最大概率类
        if pred.shape[1] > 1:
            pred = torch.softmax(pred, dim=1)
            pred = pred.max(dim=1)[0]
        else:
            pred = torch.sigmoid(pred).squeeze(1)

        # 核心计算（注意梯度传播）loss = torch.mean(pred * gt_dist)
        return loss / self.theta  # 缩放损失范围