论文学习记录之一种具有边缘增强特点的医学图像分割网络

标题：一种具有边缘增强特点的医学图像分割网络

期刊：电子与信息学报 -（2022年5月出刊）

摘要：针对传统医学图像分割网络存在边缘分割不清晰、缺失值大等问题，该文提出一种具有边缘增强特点的医学图像分割网络(AS-UNet)。利用掩膜边缘提取算法得到掩膜边缘图，在UNet扩张路径的最后3层引入结合多尺度特征图的边缘注意模块(BAB)，并提出组合损失函数来提高分割精度；测试时通过舍弃BAB来减少参数。在3种不同类型的医学图像分割数据集Glas, DRIVE, ISIC2018 上进行实验，与其他分割方法相比，AS-UNet分割性能较优。

关键字：医学图像分割；注意力机制；边缘注意；组合损失函数

一、介绍

本文基于UNet提出一种具有边缘增强特点的医学图像分割网络(Add-andSubtract UNet, AS-UNet)。

Add是指训练时在UNet 的基础之上增加边缘注意模块(Boundary Attention Block, BAB),；
Subtract是指测试时舍弃 BAB。

创新点：

AS-UNet通过掩膜边缘提取算法得到掩膜边缘图，在UNet扩张路径的最后3层引入结合多尺度特征图的BAB，强化边缘，减少缺失值；
本文提出使用组合损失函数，结合基于区域的Dice损失(Dice Loss) 和基于边缘的Boundary损失 (Boundary Loss) ，在保证区域缺失值小的同时补充边缘信息，提高分割精度；
此外，由于组合损失函数包含了经过BAB前后的两种输出，在训练时经过前后向反馈不断更新AS-UNet中的网络参数，使得训练好的模型在测试时可以舍弃添加的BAB部分，保证分割精度的同时减少测试时参数。

二、具有边缘增强特点的医学图像分割网络 AS-UNet

2.1 网络结构

本文所提具有边缘增强特点的医学图像分割网络 AS-UNet 结构如图1所示。采用UNet网络的编码—解码结构，输入通过 UNet 得到的直接输出称为输出1(output1)。将UNet扩张路径中 $L_7,L_8,L_9$ 得到的特征图 $R_7,R_8,R_9$ 分别作为输入层送入BAB。 $R_7$ 经过BAB后采用 $2 \times 2$ 的上采样 (Up-conv $2 \times 2$ ) 达到与 $R_8$ 相同的尺寸，作为补充层传入 $L_8$ 的 BAB中。同理，R8经过BAB 后采用 $2 \times 2$ 的上采样达到与 $R_9$ 相同的尺寸，并作为补充层传入 $L_9$ 的 BAB中。而最后一层得到的特征图 $R_9$ 经过BAB后则直接进行 $1 \times1$ 的卷积得到输出2(output2)。此过程如式(1)所示：

$output2=d_1(B{p_2[B[p_2(B(R_7)),R_8]],R_9})$ (1)

其中， $d_s(*)$ 表示卷积核大小为 $s \times s$ 的卷积函数， $p_s(*)$ 表示卷积核大小为 $s \times s$ 的上采样函数， $B(*)$ 表示BAB函数。

训练时输入图像通过UNet生成输出1，同时将 $R_7,R_8,R_9$ 通过BAB还原图像边缘细节信息，生成输出2。将输出1和输出2分别与掩膜(Mask，即标签(Ground Truth, GT))进行损失计算，在前后向传播中更新参数，得到最终模型。

图1 AS-UNet网络结构

2.2 边缘注意力模块

在UNet中，边缘分割模糊以及部分区域缺失仍是最主要的不足。因此，本文所提AS-UNet网络通过在UNet网络中引入边缘注意模块BAB以强化边缘信息，补充缺失区域。BAB结构如图2所示。

图2 BAB结构

将UNet扩张路径中 $L_7,L_8,L_9$ 的最后特征图 $R_7,R_8,R_9$ 分别作为输入层送入BAB。具体地，假设现在是第 $i$ 层，模块的输入特征维度为 $R_i \in \Re ^{w^i \times h^i \times c^i}$ ，经过 $1 \times 1$ 的卷积后，将通过掩膜边缘提取算法得到的掩膜边缘图 $\hat{M}_i \in \Re ^{w^i \times h^i \times c^i}$ 改变大小(Resize)至与输入相同分辨率，随后将两者按通道连接(Concat)，再将连接后的特征图进行 $3 \times 3$ 的卷积，得到全局信息。同时，为了结合多尺度信息，对于 $L_8,L_9$ ，将 $3 \times 3$ 卷积后得到的特征图与上一层经过BAB并上采样后得到的补充层特征图 $f_{i-1}$ 按通道连接，执行完连接操作后再进行 $3 \times 3$ 的卷积，而 $L_7$ 则无须连接补充层，直接进行 $3 \times 3$ 的卷积。最后经过注意力模块得到输出。此过程如式(2)所示：

$F_i = AB (d_3\{c[d_3[c(d_1(R_i),\hat{M}^i)],f_{i-1} ] \})$ (2)

其中， $d_s(*)$ 表示卷积核大小为 $s \times s$ 的卷积函数， $c(*)$ 表示连接操作， $AB(*)$ 表示注意力模块函数， $U \in \Re ^{w \times h \times c}$ 表示输出。

在BAB中，改变掩膜边缘图大小并按通道连接至各特征图之后，有效地提供了精确的图形边缘，增强图像对于目标区域的边缘信息获取，使得后续网络更重视这一部分的信息，从而弥补UNet 对边缘信息的忽视。

2.2.1 掩膜边缘提取算法

为了引入像素级别的精确边缘信息，本文通过掩膜边缘提取算法将训练图像的掩膜导出为掩膜边缘图，以此作为边缘信息的重要补充。掩膜边缘提取算法过程可表述为：遍历掩膜的每一个像素点 $(i,j)$ ，当遍历的像素点值为0，且以该像素点为中心的九宫格内其余像素点不都为0时，将该像素点记作0，直到掩膜的所有像素点都遍历结束后，生成掩膜边缘图。图3为掩膜与经过掩膜边缘提取算法导出的掩膜边缘图。

图3 图像掩模与对应边缘图

2.2.2 注意力模块

受空间通道压缩与激励(spatial-channel Sequeeze & Excitation, scSE) 注意力模块的启发，本文提出一种新的注意力模块，具体结构如图4所示。对于输入特征图 $U \in \Re ^{w \times h \times c}$ ，分别在通道和空间上进行压缩得到特征图 $\hat{U}_{sCE} \in \Re ^{w \times h \times 1}$ 和向量 $\hat{U}_{sCE} \in \Re ^{1 \times 1 \times c}$ ，两者相乘得到新的权重 $W \in \Re ^{w \times h \times c}$ ，再将其与输入特征图 $U$ 逐像素相乘得到最终输出 $\hat{U}$ 。整个过程如式(3)所示：

$\hat{U} = (\hat{U}_{sCE} \times \hat{U}_{cSE}) \odot U$ (3)

其中， $\times$ 代表两图扩展维度后直接相乘， $\odot$ 代表逐像素相乘。

cSE和sCE顺着空间和通道维度分别进行特征压缩，将得到的特征权重向量分别加权到输入特征图上，完成不同维度上对原始特征的重标定。而 scSE的工作仅仅是将cSE和sCE的输出结果逐像素相加，同时完成空间以及通道上的压缩与激励。而本文所提出的注意力模块首先将空间和通道上压缩得到的特征图进行相乘，得到与输入相同大小的权重W，如此便可以为输入特征图的每个像素都提供各自的权重，且相乘能够让重要的位置更加突出，而价值较小的位置则可以被抑制。随后再将这一权重与输入逐像素相乘得到最终输出。此方法相较于scSE同样简单且不增加任何参数，但能够更多地关注重要位置的信息，获取感兴趣区域及边缘特征信息。本文将其应用于边缘注意模块3×3的卷积层之后，帮助模型取得更好的分割效果。

图4 注意力模块

2.3 组合损失函数

在图像分割领域，主要采用的损失函数包括基于分布的交叉熵损失函数(Cross-Entropy Loss)以及基于区域的Dice损失函数。

交叉熵损失函数单独评估每个像素矢量的类预测，然后对所有像素求平均值，使得图像中的像素能够被平等学习。但是，医学图像中常出现类别不均衡的问题，由此导致训练会被像素较多的类主导，对于像素较少的物体很难学习到其特征，从而降低网络的有效性。
而 Dice损失函数，本质上是衡量两个样本的重叠部分，虽然解决了类别不均衡的情况，但是未关注到图像的边缘信息。

由于医学图像的边缘信息尤为重要，因此本文提出结合基于区域的Dice 损失和基于边缘的Boundary 损失的组合损失函数L，在两个不同的侧重维度上进行监督。Dice损失和Boundary 损失的定义如式(4)、式(5)所示：

$L_{Dice} = 1- \frac{ 2 \sum_{i=1}^N \sum_{c=1}^{C} g_{i}^{c} s_{i}^{c} } { \sum_{i=1}^N \sum_{c=1}^{C} g_{i}^{c^{2}} + \sum_{i=1}^N \sum_{c=1}^{C} s_{i}^{c^{2}}}$ (4)

其中， $i$ 表示每个像素点， $c$ 表示分类， $g_{i}^{c}$ 表示是否分类正确， $s_{i}^{c}$ 表示分为每个类的概率。

$L_{BD} = \int_{\Omega }^{} \varnothing_{G} (\xi ) s_{\theta} (\xi) d\xi$ (5)

其中， $\Omega_{G}$ 为边界水平集表示，若 $\xi \in G$ ，则 $\varnothing_{G}(\xi)$ 为该点与GT距离的负值，反之为正值， $s_{\theta}(\xi)$ 是网络的概率输出。

损失函数 $L$ 的定义如式（6）所示：

$L = \alpha L_{Dice} + \beta L_{BD}$ （6）

其中，参数 α 和 β 为平衡系数，用于平衡区域损失及边缘损失对最终结果的影响。 $L_{Dice}$ 与 $L_{BD}$ 分别为Dice损失和Boundary损失，均需结合输出1和输出2计算，计算方式如式(7)、式(8)所示：

$L_{Dice}=L_{Dice}(y_{gt},y_1)+L_{Dice}(y_{gt},y_2)$ (7)

$L_{BD}=L_{BD}(y_{gt},y_1)+L_{BD}(y_{gt},y_2)$ (8)

其中， $y_{gt}$ 为训练集的标签 GT； $y_1$ 为训练时AS-UNet 输出1的预测结果； $y_2$ 为训练时AS-UNet输出2的预测结果。

损失函数L结合了基于区域的Dice损失和基于边缘的Boundary损失，使得网络同时关注区域与边缘信息，在保证区域缺失值小的同时补充边缘信息，从而提高分割精度。随着神经网络的不断迭代，平衡系数 $\alpha$ 和 $\beta$ 自学习调整更新，促使在前半部分的UNet网络中，Dice损失占据比重较大，相对来说更加关注区域的信息，在后半部分的边缘注意模块中，Boundary损失占据比重较大，相对来说更加关注边缘信息。本文使用组合损失函数发挥边缘注意模块的作用，达到关注区域信息的同时不损失边缘信息，解决目前医学图像分割中存在的分割缺失值大、边缘不清晰的问题。

BAB的引入提高了模型的分割精度，但也不可避免地增加了网络参数，本文方法AS-UNet在测试时利用神经网络结构的冗余性进行模型精简，只需要选择模型的输出1，达到不增加网络参数的目的。

这种方法可执行的依据是使用新的组合损失函数L，L中的Dice损失和Boundary损失都结合了 y1和y2两种输出，因此，在将L总体最小化时必然会将针对y1部分的Loss最小化，如式(9)、式(10)所示。训练时经过前后向反馈不断更新AS-UNet中的网络参数，修复未留意的边缘信息。因此可实现在最终测试时即使舍弃边缘注意网络，仅采用y1作为测试结果，也不影响分割效果，在提高分割精度的同时极大地减少了模型参数，为实际应用提供了更大的可能性。

$min(L_{Dice}) = min(L_{Dice}(y_{gt},y_1)) + min(L_{Dice}(y_{gt},y_2))$ （9）