基于3D ERA U-Net的MRI海马体分割

MRI Hippocampus Segmentation Based on 3D ERA U-Net

摘要: 海马体萎缩率是痴呆症等神经系统类疾病的早期诊断标志，因此精确地对海马体进行分割有利于准确测算海马体萎缩率，从而更高效地辅助医生预诊断病人的健康状态。传统图像分割技术在分割体积小、形状复杂的海马体时效果并不理想，因此本文提出了一种3D Efficient Residual Attention U-Net模型(3D ERA U-Net)。首先，利用改进的残差注意力机制将高层特征信息与低层特征信息融合，丰富二者的信息交流，使模型高质量地学习海马体细节信息。其次，在编码器和解码器过渡区加入一个扩张模块，在解码之前强化模型的感受野，对MR图像这样的三维图像起到了联系全局信息的作用。最后，改进了一种复合损失函数，加强关注边界分割。将该模型在ADNI数据集上进行实验，相比较其他海马体模型，整体精度提高，左右海马体平均Dice系数为89.38%、Precision为88.44%、Recall为89.29%，优化效果提升明显。

Abstract: The rate of hippocampal atrophy is a marker of early diagnosis of neurological diseases such as de-mentia. Therefore, accurate segmentation of the hippocampus is beneficial to accurately measure the rate of hippocampal atrophy, so as to more efficiently assist doctors to predict the health status of patients. Traditional image segmentation techniques are not effective in segmenting the small volume and complex shape of the hippocampus, so this paper proposes a 3D Efficient Residual At-tention U-Net model (3D ERA U-Net). Firstly, the improved residual attention mechanism was used to integrate the high-level feature information with the low-level feature information to enrich the information exchange between the two, so that the model could learn the details of the hippocam-pus in high quality. Secondly, a Dilation Module is added to the transition area between encoder and decoder to strengthen the receptive field of the model before decoding, which plays a role in con-necting global information for 3D images such as MR Images. Finally, a compound loss function is improved to pay more attention to boundary segmentation. The model was tested on the ADNI data set. Compared with other models, the overall accuracy was improved, with the average Dice coeffi-cient of the left and right hippocampus being 89.38%, the Precision 88.44%, and the Recall 89.29%. The optimization effect was significantly improved.

1. 引言

海马体是人脑重要组成部分，负责记忆相关工作，以及辅助边缘系统来调节情绪、学习等 [1] 。右侧海马体参与环境位置记忆处理，左侧海马体负责上下文记忆及自传性记忆 [2] 。海马体作为最早受影响区域之一 [3] ，其损失程度往往是判断痴呆相关神经系统疾病的依据，如：自闭症和阿兹海默症(AD)等。AD临床试验中，海马体萎缩率 [4] 是诊断和预后的标志物，轻度认知障碍(MCI)患者海马体萎缩率约10%~15%，早期AD患者萎缩率约15%~30%，中度AD患者萎缩率可达50% [5] 。因此，通过医学辅助手段准确、高效地诊断出病患海萎缩率，有助于医生在早期诊断中预防或者减缓由于海马体萎缩引起的健康问题。

从生物组织上说，海马体包含灰质和白质，像素信息杂乱。从物理形态上说，海马体体积小，边缘形状复杂。传统图像分割算法有图谱法，形变模型法等。Somasundaram [6] 提出了一种基于图谱的方法，通过识别海马体感兴趣区域后，进行分割。Ge [7] 提出了一个灰质体积协方差算法，提高了对海马体结构分析的效果。但是由于海马体的复杂性，这些传统方法分割后的精确度不高。

深度卷积神经网络(CNN)广泛应用于医学图像分割，2015年，U-Net [8] 提出后，极大地改善了医学图像分割数据集样本量小的难题 [9] [10] 。Zhu [11] 将扩张密集网络嵌入到U-Net中，在保持高分辨率前提下，生成多尺度特征，实现收缩路径中低层次特征与扩展路径中高层次特征的融合。Hazarika [12] 将U-Net网络中的卷积核替换成不同尺寸，充分融合不同邻域的图像信息。曹平 [13] 提出结合序列学习和U-Net结合的方法，将双向卷积长短期记忆网络引入相邻切片之间。刘昭 [14] 尝试在3D U-Net [15] 网络后添加CBAM注意力机制，微曲率边界收敛效果提升，不过大曲率边界收敛效果仍不理想。Hung [16] 通过特征切换注意力模块实现特征选择和反馈。以上方法在一定程度上提高了海马体整体分割效果，但边界区域部分仍然存在欠分割和过度分割问题。

为提高海马体分割的精准度，本文提出了3D ERA U-Net网络。该网络以3D Res U-Net [17] 为主干。首先，融合了一个高效残差注意力模块，强化了各层信息流动性。其次，添加了一个扩张模块，提高了全局信息完整性。最后，设计了一个复合损失函数，额外关注海马体边界训练效果。本文工作优化了对MR图像海马体在细节上的分割精准度。

2. 3D ERA U-Net网络

2.1. 3D ERA U-Net网络结构

Figure 1. Encoder Block residual structure of the second layer

图1. 第二层Encoder Block残差结构

Figure 2. 3D ERA U-Net structure

图2. 3D ERA U-Net 结构

本文针对海马体分割问题，改进了3D Res U-Net主干网络，提出3D ERA U-Net网络结构。海马体形态复杂的部分与其邻域形态之间是连贯统一存在的，跳跃连接可以增加网络深度，在高低层特征间建立恒等映射，提高训练准确率。本文在各Encoder Block和Decoder Block中嵌入残差结构。以前一个Block上采样(或下采样)后的输出为当前Block残差结构的输入(Input)，在当前Block上采样(或下采样)前相加，相加后的信息作为该Block残差结构的输出(Output)，再继续进行上采样(或下采样)，上采样(或下采样)后的特征信息再作为下一个Block残差结构的输入(Input)。跳跃连接时利用1 × 1 × 1卷积，调整输入的通道数量和输出通道数量一致，图1所示为第二层Encoder Block的残差结构(后称Block均为嵌入残差结构后)。

在主干网络的基础上，首先，将每个Encoder Block残差信息融合的相加节点，替换成融合高层特征信息和低层特征信息的高效残差注意力模块(efficient residual attention, ERA)，让低层信息即时矫正注意力机制学习到的重要特征，再将经过注意力机制学习后的特征信息下采样，作为下一个Block的输入。同时，将注意力机制学习到的特征信息与对应的Decoder Block级联。其次，在最后一个Encoder Block输出后，加入了一个扩张模块(dilation module, DM)，利用扩张卷积的不同扩张率注入不同感受野的信息，整体网络结构如图2所示。

2.2. 高效残差注意力模块

本文在3D Res U-Net的Encoder Block中间均嵌入了一种高效残差注意力模块(ERA)，该模块改进于ECA [18] 。残差结构在高低层特征之间建立起一个实时通道，当网络进行反馈传播时，高层特征可以和低层特征即时地进行信息交流。

ECA模块是一种精简、高效的注意力机制，其设计了一个自适应卷积核的非线性函数，令每个通道及其K领域进行局部跨通道交互信息，相比传统的注意力机制，增加了通道之间的信息交互覆盖率。首先，将ECA模块扩展成3D模块。然后，由于海马体本身既有白质，又有灰质，而ECA原本使用的全局最大池化层在保留多种信息方面的效果略微逊色，因此并行添加一个全局平均池化层，注入海马体的多样特征。接着，让两个特征信息相加，用1维卷积降维，令每个通道只包含1 × 1 × 1的信息。而后，利用ECA的自适应卷积核紧密覆盖邻近通道信息，覆盖率K = 3。将这个综合通道信息经过Sigmoid函数后与当前Encoder Block输出相乘，再利用残差跳跃，与当前Encoder Block输入相加。最后，将该模块的输出通过卷积核大小为2 × 2 × 2，步长为2的卷积降采样，代替了原来的池化下采样，高效残差注意力模块结构如图3所示。

Figure 3. Efficient residual attention module

图3. 高效残差注意力模块

2.3. 扩张模块

在三维图像的分割任务中，除了学习重点信息以外，还需要整合全局上下文。扩张卷积通过不同的膨胀系数，灵活获得特征信息的不同视野。因此本文在编码和解码过渡区添加了一个扩张模块(DM)，其包含两层扩张卷积。根据Encoder阶段的计算，最后一个Encoder Block输出为128 × 8 × 8 × 8。在该特征信息后接上扩张模块，第一个扩张卷积层的扩张率为1，第二个扩张卷积层的扩张率为2。两个扩张卷积的感受野分别为3和7，第一个扩张卷积映射了每个海马体小邻域的特征，第二个扩张卷积映射了每个海马体大邻域的特征。最后再将输入在内的三个特征信息加在一起，上采样之后作为第一个Decoder Block的输入，扩张模块的结构如图4所示。

Figure 4. Dilation module

图4. 扩张模块

2.4. 损失函数

本文主要将海马体和其余脑组织的分割看作是一个二分类问题。海马体为前景，其余脑组织为背景。为了使训练更加关注边界信息，设计了一个区域损失和边界损失的复合损失函数，公式如下：

$Los{s}_{TB}=\alpha Los{s}_{TS}+\beta Los{s}_{BD}$ (1)

其中， $\alpha$ 、 $\beta$ 是两个函数的权重系数，通过训练测试，令 $\alpha =0.7$ ， $\alpha +\beta =1$ 。 $Los{s}_{TS}$ 是Tversky Loss [19] ，它衡量了预测结果P和真实结果G之间的重叠程度，并对识别错误的海马体、识别错误的背景进行了加权，优化了海马体在整个脑部占比小所带来的类不平衡问题。

$Los{s}_{TS}=\frac{P\cap G}{P\cap G+m\left|P-G\right|+n\left|G-P\right|}$ (2)

但仅仅使用TverskyLoss函数，在边界上的分割仍然不理想。Kervadec [20] 提出了一种通过形状上的距离度量的边界损失函数 $Los{s}_{BD}$ ，我们将其用于预测分割边界与海马体标签边界之间的距离，其公式如下：

$Los{s}_{BD}\left(\theta \right)={\displaystyle {\int }_{\Omega }{\varphi }_G\left(q\right)S_{\theta }\left(q\right)\text{d}q}$ (3)

其中， $\Omega \to R^3$ 代表训练的MR海马体图像，q是图像上的任意一点。G表示专家标注的海马体区域， $S_{\theta }\left(q\right)$ 是q点分类正确的概率， ${\varphi }_G\left(q\right)$ 是距离惩罚项，与边界的距离映射 $D_G\left(q\right)$ 有关。

${\varphi }_G\left(q\right)=\{\begin{array}{ll}{D}_G\left(q\right),\hfill & q\in G\hfill \\ -D_G\left(q\right),\hfill & q\notin G\hfill \end{array}$ (4)

$D_G\left(q\right)$ 计算了 $q\in \Omega$ 到边界 $\partial G$ 最近点 $z_{\partial G}\left(q\right)$ 的距离。

$D_G\left(q\right)=\Vert q-z_{\partial G}\left(q\right)\Vert$ (5)

当图像中两个点m、n预测出来的概率接近时，假设m点距离海马体标签较远，那么m点会被赋予较大的距离惩罚权重 ${\varphi }_G\left(m\right)$ ，这样m比n对 $Los{s}_{BD}$ 贡献就更大了。因此，通过对不同体素点赋予不同的距离惩罚权重，对海马体的边界形状进行了约束。

3. 实验设置

3.1. ADNI数据集

本研究使用的海马体数据集来源于阿尔茨海默病神经成像计划数据库(ADNI)下的EADC项目(http://www.hippocampal-protocol.net)。该数据集包括135个T1加权脑部MR图像以及由专家标注的掩码图像。具体临床统计学特征参见表1。所有MR图像尺寸均为198 × 233 × 183，体素大均为1 × 1 × 1 mm，采用NIFTI格式保存，本文使用三维医学处理软件ITK-SNAP可视化，见图5所示。

Table 1. Clinical statistical characteristics

表1. 临床统计学特征

(a) 矢状面(b) 冠状面(c) 横截面

Figure 5. Visualization of dataset

图5. 数据集可视化

由于海马体在整个大脑中占据的体积很小且相对位置比较固定，实验先将数据粗裁剪，以降低参数量和计算量。首先，利用MATLAB中的SPM12工具箱对所有MR图像进行AC校准。接着，通过对ADNI数据集中掩码图像的测试发现，64 × 64 × 64这个尺寸的空间盒可以充分容纳每个海马体。因此根据AC校准后的新原点所在的位置信息，向三个方向各拓展64个体素，图6以右海马体为示。

(a) 矢状面 (b) 冠状面 (c) 横截面

Figure 6. ADNI dataset preprocessing

图6. ADNI数据集预处理

3.2. 实验环境和参数

实验环境为Windows操作系统，在Pycharm平台上，使用Python3.8和Pytorch1.11.0框架，GPU为RTX3060，采用CUDA11.3加速计算。学习率为0.001，使用Adam优化器，batchsize大小设为2，预设的训练轮数为100。

3.3. 评价指标

为精确度量海马体的分割效果，本文将Dice系数、Precision和Recall作为评价指标。

1) Dice系数

Dice系数在分割任务中衡量了两个样本相似程度。其中P表示预测值集合，G表示真实值集合。

$\text{Dice}=\frac{2\left|P\cap G\right|}{\left|P\right|+\left|G\right|}$ (6)

2) Precision和Recall

Precision反映了海马体预测的可信度，Precision越高表示预测结果越准；Recall用以评价算法对海马体预测的查全率，Recall值越高表示漏测的概率越小。

$P=\frac{TP}{TP+FP}$ (7)

$R=\frac{TP}{TP+FN}$ (8)

其中，TP表示正确预测出海马体的体素数，FP表示将海马体预测为背景的体素数，FN表示将背景预测为海马体的体素数。

4. 实验结果与分析

Figure 7. Segmentation results of different models

图7. 不同模型的分割结果

图7为各模型分割后的三视图，分别为横断面、矢状面、冠状面。由上往下分别是3D U-Net、3D Res U-Net、3D ERA U-Net。为了对比各个模型在细节上的分割效果，特意挑选(29, 27, 36)位置的三视图。该位置所呈现的海马体形态完整，边缘细节清晰，复杂程度相对较高。从图7可见，几种分割模型都存在一定误判。3D U-Net在海马体边界处几乎没有分割能力，大面积的将海马体周围相似的灰质组织错判了，过分割情况严重，这是因为3D U-Net特征提取能力不足。3D Res U-Net利用跳跃连接加深了网络深度，虽然在海马体微曲率边界处分割效果明显提升，但是对于相似灰度的区域的区分度特征提取能力不够，在曲率较大处仍然出现了明显错判。而3D ERA U-Net模型通过高低语义层的特征信息融合以及边缘损失函数的约束，在海马体大曲率边界处和微曲率边界处分割效果都有了明显提升，特别是在海马体狭长处，更少的脑部周围组织被错判。

(a) 训练集 (b) 验证集

Figure 8. Loss curves of different models

图8. 不同模型的损失函数曲线

(a) 训练集 (b) 验证集

Figure 9. Loss curves of different loss functions

图9. 不同损失函数的损失曲线

1) 损失函数

为验证模型性能，将3D ERA U-Net和3D U-Net、3D Res U-Net比较。图8给出了以Tversky Loss为损失函数的三种模型损失曲线。当模型迭代到40左右时，三个模型训练集收敛效果基本稳定。其中，训练集(图8(a))中3D ERA U-Net损失值在0.0815上下浮动，效果明显优于另外两个模型。验证集(图8(b))前期曲线震荡明显，当模型迭代到46左右时，3D ERA U-Net损失曲线的震荡减弱，逐渐稳定在0.09左右。图9以3D ERA U-Net为标准，对比了单独使用Tversky Loss损失函数和本文提出的以边界约束为辅的复合损失函数损失曲线。复合损失函数在多个体素预测正确且概率相近的情况下，权衡了这些体素离最近边界点的距离，将距离作为惩罚项，从而筛选出最接近标签位置同时又预测正确的边界点。经过该算法的优化，模型训练集(图9(a))损失函数降到了0.059左右，验证集(图9(b))损失函数降到了0.078左右。

2) 评价指标对比

对训练好的模型进行定量评估是评价模型的重要标准。表2展示了3D U-Net [16] 、3D Res U-Net [17] 以及HGM-cNet [21] 、ViT U-Net [22] 的分割海马体的各项指标。3D U-Net在海马体分割中表现一般。这是因为3D U-Net编解码器采用了普通卷积和简单的级联方式，存在特征提取能力不足和特征融合不足问题。3D Res U-Net在3D U-Net的基础上设计了跳跃连接结构，进一步聚合了高级语义信息与低级语义信息，但是没有考虑到类别不平衡会产生很多冗余特征。HGM-cNet和ViT U-Net是近几年新提出的海马体分割模型且均是在3D UNet的模型上进行改进。HGM-cNet融合了海马体灰质概率图，矫正了灰质部分的错判，但是与大脑其他组织灰度相似处的特征提取仍然不足。ViT U-Net在nnU-Net中整合了transformer模块，帮助模型保持了记忆，但是训练参数很大且边缘处损失仍相对严重。3D ERA U-Net的左右海马体平均Dice为89.38%、Precision为88.44%、Recall为89.29%，指标均高于其他模型，说明本文提出的模型分割精度更高。

Table 2. Segmentation of different models of quantitative evaluation indicators

表2. 不同模型的分割定量评估指标

3) 消融实验

了验证本模型各模块的有效性，对各个模块进行了消融实验，如图10所示。将3D Res U-Net模型作为backbone，用热力图体现各模块对特征提取优化的性能。仅使用3D Res U-Net时，网络对海马体的边缘分割和窄段分割效果不佳，注意力特征值较小。加入高效残差注意力模块后，网络对海马体的边界部分判别效果有所提升，注意力特征值有所提高，但是存在过分割的情况。加入扩张模块后，边界过分割情况被矫正，但是狭窄区存在欠分割的问题。改进损失函数后，海马的狭窄处和边界处出都得到了更加精确的关注，过分割和欠分割情况都被有效的改善，分割效果大大提升，消融实验评价指标表3。

Figure 10. Heat map of Ablation experiment

图10. 消融实验热力图

Table 3. Ablation experiment based on 3D Res U-Net

表3. 基于3D Res U-Net的消融实验

5. 结论

本文针对MRI海马体分割问题，提出了3D ERA U-Net。该网络以3D Res U-Net作为主干。在编码结构中融合了高效残差注意力模块，通过灵活的通道覆盖率，实时帮助特征信息矫正，并行增加全局池化后，提高了模型识别类似海马体组织的能力。在编码–解码过渡区添加了扩张模块，通过学习不同感受野下的全局信息，进一步提高了模型对海马体和类似周围组织的分辨力，为解码提供了丰富的信息。两者综合改善了欠分割和过分割情况，缓解了类别不平衡问题。为训练设计的复合损失函数针对边界预测结果做了加权处理，在预测概率正确且相似的情况下，权衡了与最近边界点的距离因素，大大提高了正确边界点的筛选能力。通过在ADNI数据集上实验表明，该模型分割左右海马体平均Dice系数为89.38%、Precision为88.44%、Recall为89.29%，均高于其他海马体分割模型。对网络进行消融实验，将输出特征以热力图的形式可视化，可以清晰地观察到，高效残差注意力模块和扩张模块对海马体大面积区域的识别效果明显提升，而复合损失函数在添加距离惩罚项后，精细化了边界部分以及海马体的狭窄区域的判别，由此证明，本文提出的模块对于模型的特征提取能力确实起到了优化作用。因此，该模型具有更精准的分割性能，尤其是边界分割效果提升明显。