LSegDiff:用于医学图像分割的潜在扩散模型
摘要
扩散模型最初是为图像生成而设计的,它也可以有效地应用于各种任务,包括语义分割。然而,现有的基于扩散的语义分割方法大多在高维像素空间中运行,在训练和推理过程中需要大量的计算和内存资源。本文首次尝试利用潜在扩散模型进行语义分割。具体而言,我们提出了一种快速有效的潜在扩散模型,并对其在医学图像分割任务中的应用进行了评价。首先,我们训练了一个变分自编码器(vae)网络,将二值图像掩码转换为紧凑的潜在向量。扩散过程可以在这个低维潜在空间中执行,从而大大加速。随后,我们使用V AE解码器从扩散过程产生的潜在输出向量重建精确的预测图。最后,我们通过使用形态学操作的简单后处理步骤来细化最终的分割结果。我们报告了两个公共数据集的结果,包括结肠息肉图像和皮肤癌图像。实验表明,与传统扩散模型相比,我们的方法在具有更好的训练和推理速度以及更有效的内存消耗的同时,获得了具有竞争力的准确性。
1 介绍
医学图像分割是一个专注于精确描绘医学图像中特定结构的专业领域。
特别是,我们的工作集中在两个关键领域,包括息肉和皮肤病变的分割。两特定域分割是指在医学图像中识别和描绘异常区域,如息肉和皮肤病变的过程,特别是在内窥镜和皮肤镜等医学成像背景下。息肉是一种异常生长的组织,可以出现在身体的各个部位,如结肠、胃或呼吸道。这个过程在医疗诊断和治疗计划中至关重要,因为准确地识别和测量息肉可以帮助医生发现疾病的早期迹象,如结肠直肠癌,并对患者护理做出明智的决定。此外,皮肤病变是指皮肤表面的异常和病理区域,包括痣、溃疡、皮疹和皮肤癌等多种情况。在本研究中,医学图像分割算法通常利用先进的图像处理技术、机器学习和深度学习方法来分析医学图像,准确定位内窥镜图像中息肉的边界和皮肤镜图像中病变皮肤的区域。
解决这一挑战的传统机器学习方法主要依赖于手工制作的特征来提取图像属性,如颜色、形状和纹理。然而,由于病变类型的多样性,这些方法在准确表示和检测息肉时经常遇到困难。近年来,深度神经网络,特别是UNet[19]架构,已经成为息肉分割的主要方法。这些网络可以捕获复杂和抽象的特征,从而产生令人印象深刻的性能结果。然而,深度神经网络的采用也引入了复杂性方面的权衡,因为模型可能非常大(包含数百万个参数),并可能导致显著的推理延迟。
去噪扩散模型在各种生成任务中取得了显著的成就,包括将医学图像细分为有意义的组件。MedSegDiff[23]利用去噪扩散模型对二维医学图像进行分割。这是通过利用名为
去噪-UNet,它将去噪原理与UNet神经网络的框架相结合。此外,利用傅里叶变换结合了结构间的洞察。这种技术组合使MedSegDiff能够有效地隔离和突出医学图像中感兴趣的区域,从而获得准确和详细的2D分割结果。另一项值得注意的研究[22]采用扩散模型来解决分割二维医学图像的挑战。在他们的方法中,采用扩散模型逐步细化分割结果。具体来说,每个扩散步骤的输出结果在测试阶段通过求和进行汇总。这种聚合策略增强了分割结果的弹性,提高了最终分割区域的一致性和可靠性。
虽然去噪扩散模型可以获得准确的分割结果,但它们在推理过程中使用的迭代细化过程可能导致与计算复杂性、内存使用以及准确性和速度之间的权衡相关的挑战。
稳定扩散[18]技术可以通过加速迭代细化过程的收敛来解决扩散图像模型去噪时推理缓慢的局限性。
这可以在不牺牲分割或去噪结果质量的情况下实现更快、更有效的推理。隐空间中去噪扩散模型的一般概念包括在应用迭代去噪过程之前将带噪声的输入图像转换为低维隐空间表示。
该方法结合了降维、降噪和稳定收敛的优点,实现了准确高效的去噪结果。
提出了一种利用潜在扩散模型进行语义分割的新方法。具体而言,我们提出了一种快速有效的潜在扩散模型LSegDiff,并评估了其在医学图像分割任务中的性能。
LSegDiff基于三个阶段:将图像从像素空间转换到低维潜在空间,在潜在空间上训练条件扩散模型,将扩散模型的输出从潜在空间转换回原始像素空间。具体而言,我们的主要贡献如下:
-据我们所知,我们提出的方法LSegDiff是第一次尝试使用潜在扩散模型进行语义分割。
我们训练了一个V AE模型,将像素级二值图像掩模映射到潜在空间向量。这个预训练的V AE允许我们在低维空间中执行扩散过程。
-在几个基准医学图像分割数据集上的实验表明,与传统的像素空间扩散模型相比,我们的方法在速度和计算资源方面具有优越的性能,同时保持相同的精度水平。
本文的其余部分组织如下。在第二节中,我们回顾了相关工作。第3节描述了我们提出的模型。
实验结果见第4节。最后,我们对本文进行了总结,并在第5节中提出了未来的工作。
2 相关工作
2.1 图像语义分割
图像语义分割是计算机视觉领域的一项重要任务,具有广泛的实际应用。分割模型应用于医学图像诊断、自动驾驶汽车、场景理解等多个任务。已经提出了许多不同的方法来解决这个任务。这些方法可以根据不同的网络架构分为四类:基于cnn的架构、基于变压器的架构、基于mlp的架构和其他。针对该任务提出的最著名的UNet模型[19]显著提高了性能。另一种利用变压器架构[24]或使用卷积核变体[4]的方法也可以在语义分割数据集上获得SoTA结果。一些基于扩散的架构,如SegDiff[1]、MedSegDiff[23]和DiffUNet[25],与传统方法相比,也取得了显著的效果。
对于特定的任务,如医学语义分割,一些基于扩散的架构,如MedSegDiff[23]或DiffUNet[25]表现非常好,并在医疗数据集上获得了显着的结果。
2.2 条件扩