【医学图像分割 2024】BEFUnet

论文题目：BEFUnet: A Hybrid CNN-Transformer Architecture for Precise Medical Image Segmentation

中文题目：BEFUnet： 用于精确医学图像的混合 CNN 变换器架构 分割

论文链接：https://arxiv.org/abs/2402.08793

论文代码：Omid-Nejati/BEFUnet: A Hybrid CNN-Transformer Architecture for Precise Medical Image Segmentation (github.com)

论文团队：伊朗科技大学电气工程学院，伊朗德黑兰

发表时间：

DOI:

引用：

引用数：

摘要

医学图像的准确分割对于各种医疗保健应用至关重要。卷积神经网络(cnn)，特别是像U-Net这样的全卷积网络(fcn)在医学图像分割任务中取得了显著的成功。然而，它们在捕获全局上下文和远程关系方面存在局限性，特别是对于形状、规模和纹理具有显著变化的对象。虽然Transformer在自然语言处理和图像识别方面取得了最先进的成果，但由于图像局部和平移不变性问题，它们在医学图像分割方面面临挑战。为了解决这些挑战，本文提出了一种创新的u型网络BEFUnet，该网络增强了身体和边缘信息的融合，以实现精确的医学图像分割。BEFUnet包括三个主要模块，包括新颖的局部交叉注意特征(LCAF)融合模块，新颖的双电平融合(DLF)模块和双支路编码器。双支路编码器由边缘编码器和主体编码器组成。边缘编码器使用PDC块进行有效的边缘信息提取，而主体编码器使用SwinTransformer捕获全局关注的语义信息。LCAF模块通过选择性地对两种模式之间空间接近的特征进行局部交叉关注，有效地融合了边缘和身体特征。与全局交叉注意相比，这种局部方法显著降低了计算复杂度，同时保证了准确的特征匹配。BEFUnet在医学图像分割数据集的各种评估指标上展示了优于现有方法的性能。BEFUnet中CNN和transformer架构的协同作用在处理不规则和具有挑战性的边界方面表现出色，使其成为推进医学图像分割任务的一种有前途的方法。我们的代码在GitHub上是公开的。

1. 介绍

医学图像分割在计算机视觉中至关重要，它为全面分析提供了对解剖区域的重要见解，并帮助医疗保健专业人员进行损伤可视化、疾病监测和治疗计划。cnn，尤其是像U-Net[36]这样的全卷积网络(fully convolutional networks, fcn)[32]已经成为主流方法，在各种医学应用中取得了显著的成功，如心脏、器官和息肉的分割。自动医学图像分割领域因其减轻放射科医生劳动密集型任务的潜力而获得了极大的关注。FCNs及其变体，包括SegNet[2]、U-Net[36]、DenseNet[4]和DeepLab[9]，通过从MRI、CT扫描、PET/CT扫描和其他医学成像方式中进行准确分割，在加强医学诊断和治疗方面发挥了至关重要的作用。

基于cnn的图像分割方法由于其强大的局部平移不变性、表示学习能力和滤波器共享特性而受到广泛青睐。然而，这些方法在捕获远程关系和明确的全局上下文方面存在局限性，特别是对于在纹理、规模和形状上具有较大患者间差异的对象。各种策略，如扩展卷积[51]、图像金字塔[52]、先验引导方法[11]、多尺度融合[18]和自注意机制[35]，都试图解决医学图像分割中的这些局限性，但在提取全局上下文特征方面存在弱点。在医学图像分割的成功架构中有FCN和U-Net，它们都在编码器和解码器模块上使用卷积层。U-Net[36]的跳跃连接已被证明在分割任务中是通用的，但卷积层的局域性限制了其在捕获医学图像分割中至关重要的形状和结构信息方面的表示能力。CNN模型在建立远程依赖关系和全局上下文方面的有限性能进一步阻碍了其在图像分割方面的有效性。在医学图像分割的CNN架构中，解决这些挑战仍然是一个持续的追求。

最近在自然语言处理(NLP)中Transformer的成功[44]导致了视觉Transformer的发展，以克服CNN在图像识别任务中的局限性[14]。视觉Transformer利用多头自我注意(MSA)来建立长期依赖关系并捕获全局上下文。然而，普通的视觉Transformer需要大量的数据来泛化，并且具有二次复杂度。各种方法，如DeiT[40]、Swin Transformer和金字塔视觉Transformer旨在解决这些限制。虽然Transformer在图像分类和语义分割方面取得了最先进的(SOTA)结果，但它们的高计算能力要求阻碍了实时应用，如放疗。为了解决CNN模型的局限性，Vision Transformer (ViT)模型利用了MSA机制，与基于卷积的方法相比，实现了SOTA性能。混合Transformer架构将Transformer与基于cnn的网络集成在一起，而纯Transformer(如swan - unet和其他Transformer)将Transformer应用于编码器和解码器，以在多个级别上进行全局特征表示。尽管Transformer器具有捕获全局依赖关系的能力，但它们在图像局部性和平移不变性方面存在问题，影响了器官边界的准确分割

视觉Transformer的最新进展强调了多尺度特征表示的重要性。DSTransUNet[29]、CrossViT[7]和HRViT[18]引入双支路Transformer架构来提取上下文信息，增强细粒度特征，改进语义分割。在医学图像处理中，研究人员探索多分支架构来解决挑战，例如Valanarasu等人[43]使用欠完全和过完全编码器，Lin等人[29]使用双分支SwinTransformer，Zhu等人[55]利用不同的分支来处理不同的数据模式并实现卓越的性能。这些多分支模型在处理医学图像中的复杂特征方面显示出良好的前景。

尽管视觉Transformer有能力捕捉全局上下文表征，但它们的自注意机制可能忽略低级特征。为了解决这一问题，TransUnet[8]和LeVitUnet[49]等混合CNN-transformer方法将cnn的局域性与Transformer的远程依赖性相结合，对医学图像分割中的全局和局部特征进行编码。然而，这些方法面临着如何有效地结合高阶和低阶特征，同时保持特征一致性和正确利用分层编码器产生的多尺度信息的挑战。此外，许多现有的基于深度学习的医学图像分割方法往往只关注身体特征，而忽略了边缘信息的重要性。一些研究，如Kuang等[26]和Yang[50]，试图通过将边缘特征与身体特征分离或将其作为附加约束来纳入边缘特征。然而，这些方法并没有充分发挥边缘特征在医学图像分割中的潜力。

在本文中，我们引入了一种名为Body and Edge Fusion Unet (BEFUnet)的创新网络，该网络旨在通过增强边缘和身体信息的融合来实现精确的医学图像分割。BEFUnet由三个关键模块组成:双支路编码器、双电平融合(DLF)模块和局部交叉注意特征融合(LCAF)模块。双分支编码器使用轻量级的cnn分支和基于SwinTransformer的分层Transformer分支同时提取边缘和主体信息。LCAF通过考虑位于两个分支之间的位置的特征，有效地融合了跨模态特征，从而提高了精度，同时降低了计算复杂度。此外，DLF模块是一个多尺度视觉Transformer，采用交叉注意机制融合两个获得的特征映射。提出的BEFUnet不仅解决了上述问题，而且在各种评价指标方面优于其他方法。我们的主要贡献:

将CNN的边缘局部语义信息与Transformer的身体上下文交互相结合，增强互补特征的集成的新型混合方法。在医学图像分割中，这种方法在处理不规则和具有挑战性的边界时特别有利。

• 我们引入了一个双级融合模块，它有效地融合了粗粒度和细粒度的特征表示。
• 我们的实验结果强调了BEFUnet的显著有效性，因为我们在三个不同的医学图像分割数据集上精心训练和评估了我们的模型。
• 在我们的综合比较中，BEFUnet始终优于各种最先进的模型，证实了其在处理不同数据集的各种分割挑战方面的鲁棒性和优越性。

2. 相关工作

2.1 基于CNN的分割网络

在早期医学图像分割中，传统的机器学习算法被广泛使用[41]。然而，随着深度cnn、U-Net的出现，情况发生了变化[36]，其变体变成了强大的替代方案。U型结构的简单性和优越的性能使其成为二维和三维医学图像分割场景的热门选择。在医学图像分割领域，研究的重点主要转向了cnn及其变体的应用，以获得准确高效的分割结果。研究人员引入了各种方法来解决FCN的局限性[32]。这些方法包括融合不同层的输出，使用扩展卷积[51]，以及使用上下文建模[52]。此外，u型编码器-解码器结构，如U-Net及其变体，因其进一步增强了医学成像任务的性能而引起了极大的关注。在这些架构中，U-Net[36]已经成为医学图像分析的标准选择。它由一个对称的编码器和解码器网络组成，具有跳跃连接，具有很高的效率。在这一成功的基础上，一些类似U-net的架构被提出，如Res-UNet[47]、Dense-UNet[4]、Kiu-Net[43]和U-net++[54]。这些变体中的每一个都提供了改进和适应，以满足不同的医学成像任务。例如，unet++在模块之间引入了密集的跳过连接，从而改善了结果。因此，这些结构在各个医学领域都表现出出色的性能。尽管基于cnn的医学图像分割方法取得了显著进展，但由于卷积运算固有的局部性和复杂的数据访问模式，在分割精度和效率方面仍然存在局限性。人们将继续努力解决这些挑战，并进一步提高医学图像分割模型的能力

2.2 ViT

近年来，基于Transformer的模型在各个领域，特别是在自然语言处理领域取得了显著的成功。这些模型在机器翻译等任务中实现了SOTA性能[44]。此外，随着视觉Transformer(vision Transformer, ViT)的引入，它们的应用已扩展到视觉任务[14]。ViT在图像识别方面展示了令人印象深刻的速度精度权衡;然而，它需要在大数据集上进行预训练。为了解决这一限制，已经努力提高ViT在ImageNet上的性能，从而产生了Deit[40]，其中包含了改善结果的训练策略。此外，另一个值得注意的视觉Transformer是Swin Transformer[31]，它是一种分层模型，可作为高效的视觉骨干。Swin Transformer在图像分类、目标检测和语义分割任务中达到了SOTA性能。

2.3 用于医学图像分割的Transformer

在医学分割的背景下，研究人员已经探索了基于transformer的模型的潜力[46,53,17,37]。已经引入了几种变体，例如MedT[42]，它将门控轴向Transformer层集成到现有架构中，以及TransUNet[8]，它将transformer和cnn结合起来，以利用各自的优势。此外，swin-unet[5]在u型编码器-解码器架构中采用纯transformer进行全局语义特征学习。这些基于transformer的方法有望克服CNN模型的局限性，CNN模型由于其受限制的接受域而难以进行长期依赖建模。通过结合全局和局部信息，这些混合模型旨在改进医学图像3D和2D分割的特征提取[55,3,28]。

为了进一步增强边缘和身体信息的融合，提出的BEFUnet架构脱颖而出。BEFUnet采用了一种新颖的基于transformer的融合方案，保证了二维医学图像分割的特征一致性和丰富性。该方法解决了其他混合模型中简单的特征融合机制的局限性，并保证了包含丰富信息的更有效的特征映射，从而改善了分割结果。

3. 方法

BEFUnet的框架如图1所示。它由一个双分支编码器、一个称为LCAF的交叉注意特征融合模块和一个DLF模块组成，DLF模块有效地融合了粗粒度和细粒度的特征表示。在接下来的部分中，我们将提供这些组件的详细描述。

3.1 双支路编码器

双支路编码器包括主体编码器和边缘编码器。边缘编码器利用cnn架构和PDC (Pixel Different Convolution, PDC)[39]块有效提取边缘信息。另一方面，正文编码器采用Transformer体系结构来使用全局注意力捕获语义信息。

3.1.1 边缘编码器

边缘特征的提取是图像分割的关键，为了解决传统医学分割网络边缘信息提取不足的问题，引入了专门的边缘检测分支。

如图1所示，边缘编码器由四个阶段组成，每个阶段包含四个PDC块用于特征检测。利用最大池化层对各阶段之间的特征映射进行降采样，得到分层特征。第一阶段将原始的3通道图像扩展到C通道，并将特征映射的大小减小到1/4，以确保与主体编码器的输出大小兼容。

PDC块包括一个深度卷积层、一个ReLU层和一个核大小为1的卷积层。残余连接也包括在模型训练中帮助。由于缺乏明确编码的梯度信息，传统卷积网络在提取边缘相关信息方面存在困难。为了克服这个问题，引入了PDC，将梯度信息集成到卷积运算中，增强了边缘特征。

Vanilla convolution计算卷积核内像素值的加权和，PDC计算内核内像素值差异的加权和。香草卷积和PDC的表达式分别由式(1)和式(2)给出:
$$y=f(x,\theta )=\sum _{i=1}^{k\times k}{w}_i\cdot x_i\text{(vanilla convolution)}$$

$$y=f(\nabla x,\theta )=\sum _{(x_i,x_i^{\prime })\in P}{w}_i\cdot (x_i-x_i^{\prime })\text{(PDC)}$$

其中，$w_i$表示k × k卷积核中的权重，$x_i$和$x_i^{\prime }$表示核所覆盖的像素，P表示核所覆盖的局部区域中所选像素对的集合。

为了进一步增强边缘提取，我们采用了一种监督策略[48]，即为每个阶段的输出特征生成一个边缘图，并计算生成的边缘图与地面真值之间的损失

3.1.2 主体编码器

主体编码器利用Transformer架构，特别是Swin Transformer[31]，对具有全局远程建模的高级特征表示进行编码。Swin Transformer采用滑动窗口机制来构建分层特征，非常适合对不规则形状的器官进行分割。

图1用红色块表示Swin Transformer的总体结构，由四个阶段组成。第一阶段使用一个补丁嵌入层和两个Swin Transformer块对原始图像进行特征编码。图像被分割成大小为P × P的小块，重构成一维向量，投影到c维空间。添加位置参数以编码位置信息，然后将序列通过旋转Transformer块。随后的阶段对特征映射进行采样，并使用补丁合并和Swin Transformer块提取更高级别的特征

SwinTransformer的每一级由两种不同类型的块交替排列组成。第一个块由层归一化(LN)、基于窗口的多头自注意(W-MSA)、多层感知器(MLP)和残差连接组成。第二个块是类似的，但它取代了W-MSA与转移基于窗口的多头自注意(SW-MSA)。这些块内的操作可以用式(3)-(6)表示。
$$\begin{array}{c}{\hat{z}}^l=W-MSA(LN(z^{l-1}))+z^{l-1}\text{(3)}\\ z^l=MLP(LN({\hat{z}}^l))+{\hat{z}}^l\text{(4)}\\ {\hat{z}}^{l+1}=SW-MSA(LN(z^l))+z^l\text{(5)}\\ z^{l+1}=MLP(LN({\hat{z}}^{l+1}))+{\hat{z}}^{l+1}\text{(6)}\end{array}$$
在每个阶段之间，进行补丁合并以downsample feature map并收集上下文特征。

它将相邻的2×2补丁合并成一个更大的补丁，减少了补丁的数量，并将它们的尺寸连接起来，以最大限度地减少信息损失。补丁合并对特征进行2倍的下采样。假设输入图像尺寸为H × W × 3，每个阶段的输出特征尺寸分别为$\frac{H}{4}\times \frac{W}{4}\times C,\frac{H}{8}\times \frac{W}{8}\times 2C,\frac{H}{16}\times \frac{W}{16}\times 4C,\frac{H}{32}\times \frac{W}{32}\times 8C,$。

然后将提取的边缘和体特征输入LCAF模块进行进一步融合。

3.2 LCAF模块

局部交叉注意融合(LCAF)模块通过选择性地对两种图像模式之间空间上接近的特征进行交叉注意，准确有效地融合边缘和身体特征。该方法克服了传统的交叉关注方法需要存储每个patch的全局信息，不能利用模态之间像素的一对一对应进行精确匹配和全面融合的局限性

图2说明了LCAF模块的结构。它利用具有局部交叉注意的Transformer块来调制近距离的特征。将输入的边缘和体特征划分为局部区域，投影为查询向量、键向量和值向量，通过点积注意力计算注意力分数。对结果向量进行加权和，得到融合的特征表示。LCAF模块还集成了一个多头操作，将向量投影到不同的子空间上，并计算不同子空间中的交叉注意分数。LCAF模块的输出是通过将前馈网络(FFN)应用于调制特征，如如式(9)所示。残差连接用于保持特征完整性。

$$LCA(X_f)=softmax(\frac{X_{edge}{W}_Q(X_{body}{W}_K{)}^T}{\sqrt{d_k}})X_{body}{W}_v$$

$$M-LCA(X_f)=X_f+Concat[LCA(X_f{)}_1,...,LCA(X_f{)}_h]W_o$$

$$X_f=FFN(M\text{-}LCA(X_f))+M\text{-}LCA(X_f)(9)$$

在这里，Xbody、Xedge和Xf分别是指与边缘模态、体模态和融合模态相对应的特征。可学习矩阵Wk、WQ、WV和Wo分别表示键矩阵、查询矩阵、值矩阵和输出矩阵。与传统的全局交叉注意算法相比，LCAF算法通过选择性地降低了计算复杂度当地cross-attention。假设特征中有h和w个patch, LCA窗口中有hl和wl个patch，全局交叉注意(GCA)和局部交叉注意(LCA)的计算复杂度分别由式(10)和式(11)给出
$$\begin{array}{rlr}\Omega (M-GCA)& =4hw{C}^2+2(hw{)}^{2}C& (10)\\ \Omega (M-LCA)& =4hw{C}^2+2h_l{h}_w\cdot hwC& (11)\end{array}$$
LCAF模块输出融合的特征，然后将其馈送到下一阶段。

3.2.1 双级融合模块(DLF)

主要障碍在于如何有效地融合CNN和Swin Transformer的电平特征，同时保持特征的一致性。一种直接的方法是将CNN的电平和它们对应的Swin Transformer电平的和输入到解码器中，得到分割图。然而，这种方法无法确保级别之间的特征一致性，从而导致性能低于标准。因此，我们提出了一个名为DLF的新模块，通过将最小(P s)和最大(P l)水平作为输入，并采用交叉注意机制来融合跨尺度的信息，从而解决了这个问题

通常，较浅的级别包含更精确的定位信息，而较深的级别包含更适合解码器的更多语义信息。为了权衡计算成本和中间层特征映射对模型精度的边际效应，我们选择不将中间层纳入特征融合过程，从而节省计算资源。因此，我们鼓励在保留定位信息的同时对最浅(P s)和最后(P l)层进行多尺度表示

在提议的DLF模块中，类标记起着重要的作用，因为它总结了来自输入特征的所有信息。每个级别都被分配了一个类令牌，该类令牌来自于该级别规范之上的全局平均池(GAP)。

类令牌的获取方法如下:
$$\begin{array}{rl}CL{S}^s& =GAP(Norm(P^s))\\ CL{S}^l& =GAP(Norm(P^l))\end{array}$$
其中cls∈R4D ’ ×1, CLSl∈RD ’ ×1。然后，在将类标记传递到转换器编码器之前，将其与相应的级别嵌入连接起来。

小电平后面跟着S个Transformer编码器，大电平后面跟着L个Transformer编码器来计算全局自关注。值得注意的是，还在两个级别的每个标记中添加了可学习的位置嵌入，以将位置信息合并到Transformer编码器的学习过程中

将嵌入信息通过Transformer编码器后，利用交叉关注模块融合各层特征。在融合之前，两个级别的类标记被交换，这意味着一个级别的类标记与另一个级别的标记连接在一起。然后，每个新的嵌入都单独地通过模块进行融合，并最终反投影到其原始水平。

与来自其他级别的令牌的交互允许类令牌跨级别共享丰富的信息

特别是，对于小水平，这种相互作用如图3所示。f s(.)首先将CLS投影到p1的维度上，得到的输出记为CLS。CLS’s与p1相连接，作为交叉注意计算的键和值，同时也独立作为查询。由于只查询类令牌，所以交叉注意机制在线性时间内运行。最终输出Z s可以用数学形式表示如下:
$$\begin{array}{lc}{y}^s=f^s(CL{S}^s)+MCA(LN([f^s(CL{S}^s)\parallel P^l]))& \\ Z^s=[P^s\parallel g^s(y^s)]& (13)\end{array}$$

3.3 损失函数

由于编码器被设计为同时检测边缘和体特征，我们使用两个损失分量:Ledge和Lbody。

3.3.1 边缘监督损失

我们将这个损失函数应用于编码器的每个阶段生成的边缘映射。对于第j个边缘图中的第i个像素，记为y j i，损耗按以下条件计算:
$$l_i^j=\{\begin{array}{ll}\alpha \cdot \log (1-y_i^j)& \text{if}{y}_i=0\\ 0& \text{if}0<y_i<\eta \\ \beta \cdot \log y_i^j& \text{otherwise}\end{array}\text{(14)}$$
其中，y j i表示第j个边缘图中第i个像素的预测值，η为预定义阈值。被标注者标注为正数且比例小于η的像素被认为是负样本。负样本在数据集中的比例用β表示。此外，我们定义α = λ·(1−β)，其中λ是用于平衡正负样本的超参数。通过对每个像素的损失求和，得到整体的Ledge损失:
$$L_{edge}=\sum _{i,j}{l}_i^j$$

3.3.2 整体边缘损失

医学图像分割面临严重的类不平衡问题，我们采用二元交叉熵损失(Binary Cross-Entropy Loss)和骰子损失(Dice Loss)[34]作为体损失来解决这一问题。

二值交叉熵损失测量像素级预测误差，适用于大多数语义分割场景。可以表示为:
$$L_{Bce}=-\sum _{i=0}^N[(1-{\hat{y}}_i)\log (1-y_i)+{\hat{y}}_i\log (y_i)]\text{(16)}$$
Dice Loss是医学图像分割中广泛使用的一种损失函数，可以有效地处理不平衡样本。

由下式给出:
$$L_{Dice}=1-2\times \frac{2{\sum }_{i=0}^N{y}_i\hat{y}i}{\sum i=0^N(y_i+{\hat{y}}_i)}\text{(17)}$$
Lbody损失可以表示为二元交叉熵损失和骰子损失的组合:
$$L_{body}={\lambda }_1{L}_{Bce}+{\lambda }_2{L}_{Dice}$$
最后，结合Ledge和Lbody以及一个加权超参数γ，得到最终损失L:
$$L=L_{body}+\gamma \cdot L_{edge}$$