边界损失函数(Boundary Loss)原理剖析与医学图像分割实战

1次阅读

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

医学影像中器官 / 病灶的边界往往呈现模糊性（fuzzy boundaries），主要表现为：

部分容积效应：由于成像分辨率限制，边界体素同时包含目标和背景信号
低对比度区域：如 CT 中软组织间的灰度差异小于 10HU
不规则形态：肿瘤浸润性生长导致边界呈毛刺状

传统损失函数面临的主要问题：

Cross-Entropy Loss：逐像素计算时，边界模糊区域的梯度信号会被非边界区域淹没
Dice Loss：基于区域重叠计算，对边界像素的几何位置不敏感

特性	Cross-Entropy	Dice Loss	Boundary Loss
梯度特性	点级梯度	区域级梯度	边界级梯度
边界敏感度(1-10)	3	5	9
小目标适应性	容易过拟合	梯度不稳定	稳定

边界损失函数定义为预测边界 $\partial S_p$ 与真实边界 $\partial S_g$ 的距离场积分：

$$
L_{boundary} = \int_{\partial S_p} \phi_G(q) dq
$$

其中 $\phi_G$ 是通过距离变换（Distance Transform）得到的真实边界距离图：

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

flowchart LR
    A[二值分割掩码] --> B[欧氏距离变换]
    B --> C[内部距离场]
    B --> D[外部距离场]
    C & D --> E[符号合并]

距离图的梯度 $\nabla\phi_G$ 始终指向最近的真实边界
反向传播时，网络参数会沿着梯度方向更新，使预测边界向真实边界 ” 吸附 ”

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

class BoundaryLoss(torch.nn.Module):
    """
    Args:
        pred: [B,C,H,W] probability map (after softmax)
        target: [B,H,W] ground truth (0,1,...,C-1)
    """
    def __init__(self):
        super().__init__()

    def _one_hot(self, target, num_classes):
        return torch.zeros(target.shape[0], num_classes, *target.shape[1:], 
                          device=target.device).scatter_(1, target.unsqueeze(1), 1)

    def forward(self, pred, target):
        # 转换为 one-hot 编码
        target_oh = self._one_hot(target.long(), pred.shape[1])

        # 生成距离变换图
        dist_maps = torch.zeros_like(target_oh)
        for b in range(pred.shape[0]):
            for c in range(pred.shape[1]):
                pos_mask = target_oh[b,c].bool().cpu().numpy()
                if pos_mask.any():
                    pos_dist = distance_transform_edt(~pos_mask)
                    neg_dist = distance_transform_edt(pos_mask)
                    dist_maps[b,c] = torch.from_numpy(pos_dist - neg_dist)

        # 计算边界损失               
        loss = (pred * dist_maps).sum() / (pred.shape[0] * pred.shape[2] * pred.shape[3])
        return loss