深入解析CLDice损失函数：原理、实现与医学图像分割应用

1次阅读

没有评论

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

医学图像分割常面临小目标稀疏分布的难题，比如视网膜血管、神经元突触等结构仅占图像的 1%-5%。传统 Dice 损失函数定义为：

$$\text{Dice} = \frac{2|X \cap Y|}{|X| + |Y|}$$

其中 $X$ 为预测掩膜，$Y$ 为真实标签。当处理稀疏目标时会出现两个典型问题：

拓扑结构敏感度不足 ：Dice 只关注体素重叠度，无法感知血管分支断裂等拓扑错误
梯度消失风险 ：小目标区域的正负样本极度不平衡时梯度信号微弱

CLDice 通过引入中心线约束改进传统 Dice，其核心公式分两部分：

$$\text{CLDice} = \frac{\text{TopDice} + \text{VolDice}}{2}$$

其中 TopDice 确保预测中心线与真实中心线重合：

$$\text{TopDice}(V,S) = \frac{2|T(V) \cap S|}{|T(V)| + |S|}$$

$T(V)$ 表示预测体积 $V$ 的中心线变换，$S$ 是真实中心线。VolDice 则保持传统 Dice 的体素重叠计算。

import torch
import numpy as np
from skimage.morphology import skeletonize

class CLDiceLoss(torch.nn.Module):
    def __init__(self, smooth=1e-5):
        super(CLDiceLoss, self).__init__()
        self.smooth = smooth

    def centerline_transform(self, y_pred):
        """GPU 加速的中心线提取"""
        skeletons = []
        for batch in y_pred:
            skel = skeletonize(batch.cpu().numpy() > 0.5)
            skeletons.append(torch.from_numpy(skel).float())
        return torch.stack(skeletons).to(y_pred.device)

    def forward(self, y_pred, y_true):
        # 计算传统 VolDice
        intersection = (y_pred * y_true).sum()
        vol_dice = (2. * intersection + self.smooth) / 
                  (y_pred.sum() + y_true.sum() + self.smooth)

        # 计算 TopDice
        pred_centerline = self.centerline_transform(y_pred)
        true_centerline = self.centerline_transform(y_true)
        top_intersect = (pred_centerline * true_centerline).sum()
        top_dice = (2. * top_intersect + self.smooth) / 
                  (pred_centerline.sum() + true_centerline.sum() + self.smooth)

        return 1 - (vol_dice + top_dice) / 2

在 DRIVE 视网膜血管数据集上的评测结果：