皮肤病灶分割的医学边界扩散模型
摘要
由于多尺度边界关注和特征增强模块的进步,皮肤镜图像中的皮肤病变分割最近取得了成功。然而,现有的方法依赖于端到端学习范式,直接输入图像和输出分割图,经常与极其困难的边界作斗争,例如在特别小或特别大的病变中发现的边界。出现这种限制是因为任何有限模型的接受域和局部上下文提取能力都不可避免地受到限制,并且获取大型模型所需的额外专家标记数据的成本很高。由于将图像合成视为参数化链过程的扩散模型取得了令人印象深刻的进展,我们引入了一种将皮肤病变分割作为边界进化过程的新方法,以彻底研究边界知识。
具体而言,我们提出了医学边界扩散模型(MBDiff),该模型从随机采样的高斯噪声开始,并在有限时间内进行边界演化以获得清晰的分割图。
首先,我们提出了一种高效的多尺度图像引导模块来约束边界演化,使演化方向符合我们想要的病灶;其次,我们提出了一种基于进化不确定性的融合策略,以改进进化结果并产生更精确的损伤边界。我们在两个流行的皮肤病变分割数据集上评估了我们的模型的性能,并将我们的模型与最新的CNN和transformer模型进行了比较。我们的结果表明,我们的模型在所有指标上都优于现有的方法,并在极具挑战性的皮肤病变上取得了卓越的性能。该方法有可能显著提高皮肤病变分割的准确性和可靠性,为诊断和治疗提供关键信息。所有资源将在https://github上公开提供。
com/jcwang123/MBDiff。
1 介绍
从皮肤镜图像中分割皮肤病变是皮肤癌疾病诊断和治疗计划的关键任务[17]。人工病灶分割是费时的,而且容易在观察者之间和观察者内部发生变化。为了提高临床工作流程的效率和准确性,多年来已经开发了许多自动皮肤病变分割模型[1,2,7,10,18,19,21]。这些模型专注于使用各种技术增强特征表示,如多尺度特征融合[10]、注意机制[1,7]、自注意机制[18,19]和边界感知注意[2,18,19],从而显著提高了皮肤病变分割性能。尽管取得了这些进展,但边界模糊的皮肤病变分割,特别是在极具挑战性的尺度上,仍然是需要解决的瓶颈问题。在这种情况下,即使是最先进的分割模型也难以获得准确和一致的结果。
图1所示。边界演化过程。可以看出,通过将分割分成连续的时间步(t),可以准确地分割出各种病变,本工作称之为边界进化。
图1显示了两个具有代表性的边界,其中显示了一个极小的病变和一个特别大的病变。小图像空间占1.03%,大图像空间占72.96%。如前所述,解决这两类病灶的分割问题有不同的策略。(1)对于小病变,将较低深度的特征转换到较高深度的卷积层,可以避免丢失局部上下文[10]。(2)对于较大的病灶,通过扩张卷积(expanded convolution)扩大感受野[1],甚至全局关注(global attention)[18]都可以捕捉到远程依赖关系,从而改善边界决策。除了如何在不同的尺度下产生稳定的表示的挑战之外,多尺度的病变会引起训练波动,即小的病变通常会导致大的Dice损失。给予更多的边界意识监管可以在一定程度上减少这些负面影响[2,19]。最新的变压器Xbound-Former通过跨尺度边界学习全面解决了多尺度边界问题,无论大小病变都能达到更高的性能。
然而,目前的皮肤病变分割模型仍然在与临床实践中经常遇到的极具挑战性的病例作斗争。
虽然一些方法旨在通过结合局部和全局上下文以及多任务监督来优化模型架构,而另一些方法则试图通过收集更多标记数据和构建更大的模型来提高性能,但这两种策略都是昂贵的,并且可能受到皮肤病变边界固有复杂性的限制。因此,我们提出了一种新的方法,将焦点从仅仅分割病变边界转移到预测其演变。我们的方法受到扩散概率模型在图像合成方面的最新进展的启发[6,9,14,15],扩散概率模型在一系列有限步骤中从随机抽样的高斯分布生成合成样本。我们采用这一过程将皮肤病变边界的演变建模为一个参数化的链过程,从高斯噪声开始,经过一系列去噪步骤,产生具有明确定义的病变边界的清晰分割图。通过预测链过程的下一步而不是最终的分割图,我们的方法能够比以前的模型更准确地分割具有挑战性的病变。我们在图1中说明了边界演化的过程,其中每一行对应于演化过程中的不同步骤,最终形成具有明确边界的清晰分割图。
在本文中,我们提出了一种医学边界扩散模型(MB-Diff)来改进皮肤病变的分割,特别是在病变边界模糊和尺寸极大或极小的情况下。MB-Diff模型遵循平面扩散模型的基本设计,使用顺序去噪过程生成病灶掩模。然而,它也包括两个关键的创新:首先,我们开发了一个高效的多尺度图像引导模块,该模块使用预训练的变压器编码器从先验图像中提取多尺度特征。然后将这些特征与演化特征融合以约束演化方向。其次,我们实现了一种基于进化不确定性的融合策略,该策略考虑了不同初始化的不确定性,对进化结果进行了细化,得到了更精确的损伤边界。我们在两个流行的皮肤病变分割数据集(ISIC-2016和PH2数据集)上对我们的模型进行了评估,发现它的性能明显优于现有模型。
2 方法
MB-Diff的关键目标是通过一系列级联步骤而不是单个步骤来学习边界进化,从而改善模糊边界的表示。在本节中,我们介绍了级联边界进化学习过程的细节以及进化过程的参数化架构。我们还介绍了基于进化的不确定性估计和边界集合技术,这些技术在提高进化边界的精度和可靠性方面具有重要的潜力。
2.1 边界演化过程
我们从最近的扩散概率模型(dpm)中获得灵感,采用逐步去噪过程来模拟MBDiff中的边界演化。
具体来说,给定图像和边界掩模分布为(X, Y),假设进化总共由T步组成,第T步(yT)处的边界为随机初始化噪声,第0步(y0)处的边界为精确结果。我们将边界演化过程表述为:
式中p(yT) = N (yT;0, I)为初始化的高斯分布,pθ(yt−1|yt)为每个可学习的进化步长,表示为高斯跃迁,表示为:
需要注意的是,预测函数以输入图像为条件,使得进化的边界能够准确拟合相应的病变。MB-Diff将边界演化建模为一个逐步去噪的过程,可以有效捕获边界模糊的皮肤病变的复杂结构,从而在病变分割方面具有优异的性能。
为了优化模型参数θ,我们使用进化目标作为每个进化步骤的后验近似值。假设分割标签y为y0,通过高斯噪声逐渐添加标签为:
其中{βt}Tt=1是一组从0到1的常数。然后,我们用贝叶斯规则计算后验q(yt−1|yt, y0)。利用MSE损失函数来度量高斯过渡分布的预测均值和协方差与演化目标q(yt−1|yt, y0)之间的距离。
2.2 具有图像先验的参数化体系结构
所提出的模型是一个参数化的链式过程,可预测μ∗t−1和?在图像x和先验进化y * t的先验条件下,在每个进化步骤t上。为了捕获这些条件的深层语义并进行有效的融合,我们采用了受普通DPM启发的基本U-Net[16]架构,并引入了新的条件融合设计,即高效的多尺度图像引导模块。
该架构由多层卷积编码器和对称解码器组成,两者之间有短连接层。为了将变量t纳入模型,我们首先将其嵌入到潜在空间中。然后,在每次卷积之前将先验进化y * t加到潜在t上。在瓶颈层,我们将进化特征与图像引导融合,以约束进化,并确保最终边界符合条件图像。
为了实现这一点,先验者与进化模型同时训练分割模型,并使用基于注意力的解析器来翻译图像特征在分割分支中分为进化分支[22]。由于分割模型的训练速度比进化模型快得多,因此我们采用预训练的金字塔视觉变压器(PVT)[20]作为图像特征提取器来获得多尺度图像特征。让{fl}4l=1表示在四个级别上提取的特征,分别是原始输入的2x, 4x, 8x, 16x。使用Adaptive A平均池化层调整三个较低级别的每个特征的大小以匹配f4的比例。然后,将这四个特征拼接并馈送到一个全连接层,将图像特征空间映射到进化空间中。然后,我们执行映射图像特征和编码的先验进化特征的简单而有效的加法,类似于时间嵌入的融合,以避免冗余计算。
2.3 进化不确定性
与典型的进化算法相似,边界进化的最终结果很大程度上受到初始种群的影响。边界演化过程是一个随机链式过程,由于每一步的高斯样本是随机的,可能会产生不同的终点。当处理边界区域中较大的模糊性时,这种差异尤其明显。原因是这些模糊区域的图像特征可能无法为进化提供判别性指导,导致不同进化时间的差异很大。而不是减少差异,我们惊讶地发现这些差异可以代表分割的不确定性。在实践中,基于进化的不确定性估计使得分割结果更加准确可信[4,5,12]。
不确定性估计:为了估计不确定性,模型参数θ是固定的,进化从随机采样的高斯噪声y∗T ~ N (0,i)开始。L e T {y∗,I T}ni=1表示总共N个初始化。演化完成后,得到{μ∗,i}ni=1, {?*,i}ni=1用于样本最终病变映射:y∗,i = μ∗,i +exp(1 2 ?∗,i)N (0, i)。与传统分割模型通常将预测缩放到0到1的范围不同,MB-Diff生成的进化映射由于随机抽样而具有不固定的分布。
由于最终结果主要是由平均值μ决定的,因此an d是预测的?有一个有限的范围[6],我们计算不确定性为:
进化集成:MB-Diff允许多次运行推理并融合获得的进化,而不是训练多个网络或参数来进行集成。然而,简单地从多个演化中平均预测恒等式是无效的,因为没有激活的MSE损失将预测恒等式限制在0或1左右,而不像Sigmoid函数将恒等式限制在0到1之间的范围。因此,我们使用Max投票算法来获得最终的分割图。在该算法中,只有当每个像素在所有n次进化中的恒等和大于阈值τ时,才能将其分类为病变。最后,分割映射生成为y * = (?N i=1 y *,i)≥τ。
表1。ISIC2016和PH2数据集上不同方法对皮肤病灶分割的比较分别给出两组的平均得分。
3 实验
3.1 数据集和评估指标
数据集:我们在实验中使用了来自不同机构的两个公开可用的皮肤病变分割数据集:ISIC-2016数据集和PH2数据集。ISIC-2016数据集[8]由国际皮肤成像协作(International Skin Imaging Collaboration, ISIC)档案馆提供,包括900个公共训练集样本和379个公共验证集样本。由于其公共测试集的注释目前不可用,我们另外收集了PH2数据集[13],其中包含200个标记样本,用于评估我们方法的泛化性能。
评价指标:为了全面比较分割结果,特别是边界划定,我们采用了四个常用的指标来定量评估我们的分割方法的性能。这些指标包括Dice分数、IoU分数、平均对称表面距离(ASSD)和边界的Hausdorff距离(第95百分位;HD95)。
为了确保公平的比较,在计算这些分数之前,所有标签和预测都被调整为(512×512),遵循先前研究[18]的方法。
3.2 实现细节
对于扩散模型超参数,我们使用普通扩散模型的默认设置,可以在补充资料中找到。关于对于训练参数,我们将所有图像的大小调整为(256 × 256),以提高内存利用率和计算效率。我们使用一组随机增强,包括垂直翻转、水平翻转和随机尺度变化(限制在0.9 ~ 1.1),来增强训练数据。我们将批大小设置为4,并训练我们的模型进行总共200,000次迭代。在训练过程中,我们使用初始学习率为1e-4的AdamW优化器。对于推理,考虑到速度,我们设置n = 4和τ = 2。
图2所示。我们的方法与SOTAs的视觉比较。前三行是来自ISIC-2016验证集的样本,后三行来自PH2数据集。我们在第三行用虚线框标出小病变。
3.3 与艺术状态的比较
我们主要将我们的方法与最新的皮肤病变分割模型进行比较,包括基于cnn和基于变压器的模型,即U-Net++ [24], CANet [7], TransFuse [23], TransUNet[3],特别是边界增强方法X-BoundFormer[18]。此外,我们针对MedSegDiff[22]评估了我们的方法,MedSegDiff是最近发布的一个基于扩散的模型,我们重新训练了200,000步以确保公平的比较。
定量结果如表1所示,其中报告了两个数据集的四个评估分数。尽管在ISIC-2016验证集上选择了性能最好的cnn和transformer参数,并且通过完成200,000次迭代选择了我们方法的参数,MB-Diff仍然实现了1.18%的IoU改进和0.7%的Dice改进。此外,我们的预测边界更接近注释,正如ASSD和HD95指标所证明的那样,它们分别减少了1.02和1.93像素。与MedSegDiff相比,MB-Diff在
所有指标。此外,我们的方法在p H2数据集上的泛化能力比p H2数据集的泛化能力有较大的提高,表明它具有更好的处理新数据的能力。
我们在图2中展示了具有挑战性的样本的视觉比较,包括来自ISIC-2016验证集的三个样本和来自PH2数据集的三个样本。
这些样本代表了社区中目前正在研究的边缘情况,包括大小变化、边界模糊和邻居混淆。
我们的视觉比较揭示了几个关键发现:(1)由于MB-Diff对边界演化的彻底学习,它在大小病变上始终如一地取得了更好的分割性能,如第3、5和6行所示。(2) MBDiff即使在人类感知几乎无法区分的情况下也能产生正确的边界,无需进一步的人工调整,具有重要的实用价值。(3) MB-Diff产生更少的假阳性分割,从而产生更清晰的预测,增强用户体验。
此外,我们在图2中提供了进化不确定性的可视化,其中更深的橙色表示更大的不确定性。很明显,大多数具有高不确定性的区域对应于错误的预测。这些信息可以在实际应用中指导人类对分割进行细化,最终提高人工智能的可信度。
图3所示。详细分析了我们的方法,包括烧蚀分析(a)和与其他基于扩散的方法的比较(b, c)。
3.4 详细的演化分析
在本小节中,我们将对我们方法中每个组件的性能进行全面分析,并将其与基于扩散的模型MedSegDiff进行比较。我们消融研究的结果如图3(a)所示,其中“w/o Evo”是指利用图像特征直接训练FPN[11]架构的分割模型,“w/o Fusion”是指不使用进化融合。为了确保公平的比较,我们对多个进化的得分进行平均,以表示“w/o融合”的性能。结果表明,我们的进化方法可以显著提高性能,基于进化不确定性的融合策略进一步提高了性能。与MedSegDiff方法相比,图3(b)的训练损失曲线显示,我们的方法收敛速度更快,损失更小,说明我们的多尺度图像制导比MedSegDiff方法更有效。此外,我们评估我们的方法使用不同迭代保存的参数的性能,如图3(c)所示。我们的结果表明,我们的方法在5万次迭代时与MedSegDiff在20万次迭代时具有竞争力的性能,并且我们的方法在10万次迭代时已经优于训练有素的MedSegDiff。
4 结论
本文介绍了医学边界扩散(MB-Diff)模型,这是一种新的皮肤损伤分割方法。我们提出的方法将病变分割作为一个有限时间步长的边界演化过程,从而可以高效、准确地分割皮肤病变。为了引导边界向病灶方向演化,我们引入了高效的多尺度图像引导模块。此外,我们提出了一种基于进化不确定性的融合策略,以获得更准确的分割。我们的方法在两个已知的皮肤病变分割数据集上进行了评估,结果表明该方法在未知领域具有优异的性能和泛化能力。通过对我们的训练方案的详细分析,我们发现我们的模型与其他基于扩散的模型相比具有更快的收敛速度和更好的性能。
总的来说,我们提出的MB-Diff模型提供了一种很有前途的解决方案,可以准确地分割皮肤病变,并有可能在临床环境中应用。