LSKA-Unet:结合大型可分离核注意力机制与Unet的细胞分割模型
LSKA-Unet: Large Separable Kernel Attention Mechanism and Unet Modeling Combined for Cell Segmentation
摘要: 医学图像中的细胞分割是生物医学研究和临床诊断的重要步骤。为了解决当前细胞分割方法存在细胞边界分割不连续、精度低等问题,提出了一种基于Unet的改进模型LSKA-Unet。首先,针对医学图像受噪声、伪影和其他干扰的影响导致图像质量低的问题,引入了大型可分离核注意力模块(Large Separable Kernel Attention, LSKA)来减少背景噪声对分割结果的干扰,通过自适应地学习空间注意力权重,强调重要特征,增强特征的表达能力;其次,使用Leaky ReLU激活函数替换ReLU激活函数,通过允许小的负输出,Leaky ReLU能保持一定的特征信息,有助于更好地捕捉医学图像中的细微变化;最后,使用BCE Loss (Binary Cross Entropy Loss)替换Dice Loss,针对细胞分割二分类这类问题,能够逐像素地评估细胞和背景的分类,从而提高模型在边界区域的准确性。实验结果表明,LSKA-Unet网络在细胞数据集Drosophila EM上,Precision、Dice系数和MIoU三大指标达到了94.82%、84.08%和84.92%,满足医学图像分割需求。
Abstract: Cell segmentation in medical images is an important step in biomedical research and clinical diagnosis. To solve the problems of discontinuous cell boundary segmentation and the low accuracy of current cell segmentation methods, an improved model LSKA-Unet based on Unet is proposed. Firstly, to address the problem that medical images are affected by noise, artifacts, and other interferences that lead to low image quality, a large detachable kernel attention module (LSKA) is introduced to reduce the interference of background noise on the segmentation results, and to emphasize the important features and enhance the expression of the features by adaptively learning the spatial attention weights; second, the Leaky ReLU activation function is used to replace the ReLU activation function, and by allowing a small negative output, the Leaky ReLU can maintain certain feature information, which helps to better capture subtle changes in medical images; finally, BCE Loss (Binary Cross Entropy Loss) is used to replace Dice Loss for such problems as cell segmentation binary classification, which can evaluate the classification of cells and background pixel by pixel, thus improving the accuracy of the model in the boundary region. The experimental results show that the LSKA-Unet network achieves 94.82%, 84.08%, and 84.92% in Precision, Dice coefficient, and MIoU on the cellular dataset Drosophila EM, which meets the requirements of medical image segmentation.
文章引用:孟令然. LSKA-Unet:结合大型可分离核注意力机制与Unet的细胞分割模型[J]. 建模与仿真, 2024, 13(6): 6027-6036. https://doi.org/10.12677/mos.2024.136552

1. 引言

医学图像包含了丰富了解剖结构和病变的形态和边界,在医学领域的研究、提高诊断精度和辅助治疗计划等多个方面具有重要的应用价值。医学图像分割分为人工、半自动和自动分割三种,其中人工分割耗时耗力,依赖于临床经验丰富的医生并且分割结果不够精确,难以满足临床的实际需要。半自动化分割虽然结合了人工和自动化的优点,但也依赖于操作者在初始标记和后期调整过程中的经验,需要专业知识来进行参数配置和调优,消耗了大量的时间。自动化分割能够借助计算机提取出感兴趣区域边缘并且迅速处理大量的医学图像,显著减少分割时间和人力成本。然而,面对医学图像结构复杂且相互重叠、噪声、伪影和病变与周围组织的低对比度等因素的影响,使得医学图像可变性很高。因此,高效率、高精度的医学图像分割方法成为当前的热点研究之一[1]

当前,基于深度学习[2]的语义分割方法已经在计算机视觉和医学图像处理领域取得了显著进展。全卷积网络(Fully Convolutional Networks, FCN) [3]是第一个将卷积网络应用于像素级分类的模型,它通过去掉全连接层,将网络的最后一层替换为卷积层,从而输出与输入图像相同大小的分割图。FCN采用的特征图为局部特征,在处理复杂场景时缺乏全局上下文信息并且难以捕捉细小结构的边界,从而导致分割结果并不是很好。Unet模型[4]是一种针对生物医学图像分割而设计的卷积神经网络(CNN),采用了编码器–解码器结构,能够有效地捕捉图像中的细节信息并提高分割精度。Unet在处理医学图像分割任务时,相较于FCN能够提供更高的分割精度和更好的边界保留效果。Zhang等人[5]提出一种新颖的脑卒中图像分割方法,称为MI-UNet。通过引入LDDMM-image算法以及预先了解GM、WM和LV的位置来利于病变分割,最终获得的精确率(Precision)、召回率(Recall)和Dice系数分别达到了65.45%、59.38%和56.72%。Wang等人[6]提出了一种模拟RNN的递归密集连体解码器结构,与超密集编码器相结合称为HD-RDS-UNet。从淋巴瘤患者全身CT扫描中手机的三维体积进行了与患者无关的五倍交叉验证,得到的Dice系数和灵敏度(Sensitivity)分别为85.85%和95.01%。此外,还进行了模型剪枝操作,在保证分割性能的同时减少推理时间和内存使用。

综上所述,考虑到Unet模型在医学领域的分割性能优秀,在原Unet模型上融入了大型可分离核注意力模块(Large Separable Kernel Attention, LSKA) [7],引入了Leaky ReLU [8]激活函数和BCE Loss (Binary Cross Entropy Loss),称为LSKA-Unet。经多次实验验证,LSKA-Unet模型相较于Unet模型以及其它主流语义分割模型具有更优秀的分割性能。

2. 改进Unet模型

本文提出一种基于Unet的改进模型LSKA-Unet,网络结构如图1所示。首先,在原始Unet的编码器部分的池化层后引入LSKA注意力机制以提高特征的表示能力,增强模型在复杂环境下的鲁棒性。其次,将原ReLU激活函数替换为Leaky ReLU激活函数来缓解梯度消失的问题,加速模型的收敛来提高模型整体的分割效率。最后将原Dice Loss损失函数替换为BCE Loss (Binary Cross Entropy Loss),更适用于细胞分割这类二分类问题,在处理边界模糊或重叠的细胞时,BCE Loss能够更好地优化像素级别的分类,增强细胞与背景的区分。

Figure 1. Network structure of LSKA-Unet

1. LSKA-Unet网络结构

2.1. 引入LSKA注意力机制

LSKA模块通过不同大小的核(kernel size)来构建不同感受野的操作,从而捕捉图像中多尺度的局部空间特征。LSKA使用深度可分离卷积(Depthwise Separable Convolution),通过水平和垂直方向的卷积操作,结合膨胀深度卷积(Dilated Convolution),在保持计算效率的同时,进一步扩大感受野。通过这种方式,LSKA可以有效的捕捉到图像的局部边缘、纹理等重要特征,适用于语义分割任务。具体网络结构如图2所示。

Figure 2. Network structure of LSKA

2. LSKA网络结构

根据图2,对任何输入特征图 F R C ×H×W ,通过将深度卷积和深度扩展卷积的二维权核拆分为两个级联的一维可分离权核。计算公式如下:

(1)

(2)

A C = W 1×1 * Z C (3)

F ¯ C = A C F C (4)

式中:CHW分别表示输入通道数、特征图的高度和宽度; Z ¯ C 表示深度卷积的输出,其核大小为 ( 2d1 )×1 Z C 是通过深度卷积核的大小为k/d一维可分离权核并且向下取整的输出,其中kd分别表示最大感受野和扩张率;*和⊙分别代表卷积和Hadamard积,然后使用1 × 1的卷积进行卷积以获得注意力图 A C ,LSKA的输出 F ¯ C 是注意力图 A C F C 的Hadamard积。

2.2. Leaky ReLU激活函数

激活函数是神经网络中的一个关键组件,起到了将输入的线性变换映射到非线性空间的作用。ReLU激活函数由于计算简单,只需取输入值与零的最大值而被广泛使用。而对于负输入,ReLU输出为零导致一些神经元在训练过程中无法被激活,权重无法更新。计算公式如下:

ReLU( x )=max( 0,x ) (5)

本文采用Leaky ReLU激活函数解决ReLU可能出现神经元死亡的问题,相比于ReLU,Leaky ReLU对于负数部分不在输出零,用很小的斜率α乘以输入。具体计算公式如下:

LeakyReLU( x )={ x, ifx0 α×x, otherwise (6)

Figure 3. Comparison of ReLU and LeakyReLU activation function curves

3. ReLU和LeakyReLU激活函数曲线对比

结合图3可以明显看出,ReLU激活函数在输入小于零时直接输出为零,而Leaky ReLU在输入小于零时输出一个很小的线性值αx (α取0.01),即使在负区域也能保持神经元的激活,避免部分神经元完全失活。针对细胞图像的复杂结构和边界,Leaky ReLU引入的非线性性质不仅有助于网络学习更复杂的特征,还能加快模型的收敛和提升训练的稳定性。这一改进使得模型在不同细胞类型和不同条件下的分割效果更优。

2.3. 改进损失函数

由于医学图像领域细胞分割难度较大,要求在细胞的边缘进行精准分割。在细胞分割中,背景和细胞比例往往不均衡,Dice Loss对小细胞或小区域的分割结果过于敏感。当小细胞的预测结果不准确时,Dice Loss会显著增加,这可能导致模型在训练过程中更加关注这些小目标,而忽视其他更大的区域。具体计算公式如下。

Dice Loss=1 2| XY | | X |+| Y | (7)

式中:X表示真实分割图像的像素标签;Y表示模型预测分割图像的像素类别; | XY | 表示预测图像像素与真实标签图像的像素之间的点乘结果相加; | X | | Y | 分别表示它们各自对应的图像中的像素相加。

二元交叉熵损失(binary cross entropy, BCE)通过引入权重来处理不平衡的数据问题,BCE Loss的平滑特性使得在训练初期,当模型预测不准确时,梯度仍然能够有效传递,这有助于模型更快地收敛。因此选用BCE Loss用于细胞分割这等二分类的问题,在评估二分类模型的输出与实际标签之间的差异有很好的效果,帮助模型学习正确的分割。其计算公式为:

L BCE = xΩ ( y( x )logp(x)( 1y( x ) )log( 1p( x ) ) ) ,Ω Z 2 (8)

式中:p(x)表示在概率图中位置x处正样本的预测概率,而y(x) ∈ {0, 1}则是对应的真实标签。面对分割图像正负样本比例不平衡的问题,用带权重图二值交叉熵改进,计算公式如下:

(9)

w( x )= w c ( x )+ w o exp( ( d 1 ( x )+ d 2 ( x ) ) 2 2 σ 2 ) (10)

式中: w c 表示平衡正负样本的权重图; d 1 ( x ) 表示位置x离最近分割个体边界的距离; d 2 ( x ) 则表示离第二近分割个体边界的距离; w o σ为常数。

3. 实验结果与分析

3.1. 数据集的准备

本文采用公开数据集Drosophila EM,包含30张果蝇中枢神经干细胞如图4所示,来源为ISBI 2012细胞分割大赛。由于数据集过少,本文采用翻转、旋转和裁剪等数据增强技术将数据集增加到128张并按9:1的比例分为训练集和测试集,最终训练集和测试集分别为115和13张。

Figure 4. Training set image

4. 训练集图像

3.2. 实验环境与参数配置

本文实验以Windows11为操作系统,处理器为i7-13700H,GPU为NVIDIA GeForce RTX4060-8G,内存为16G,使用Pytorch为深度学习框架搭建实验平台,本文实验配置能够应对大多数的深度学习任务。编程语言使用Python3.10,选用Adm优化器,初始学习率设置为0.001,epoch设置为50次,batch size设置为4。

3.3. 评价指标

本文采用精确率(Precision, P)、Dice系数(Dice coefficient)和平均交并比(Mean Intersection over Union, MIoU)作为衡量模型分割整体性能的评价指标。

精确率表示预测出的正样本中,实际为正样本所占的比例,即:

P= TP TP+FP (11)

式中:TP表示为准确预测为细胞的像素总数;FP表示为将细胞预测为背景的像素总数。

Dice系数用来计算两个样本的相似度,即:

Dice= 2×TP ( TP+FN )+( TP+FP ) (12)

式中:FN表示错误的将目标分割为背景,其余参数与式(10)一致。

MIoU用来衡量模型的整体性能,为细胞和背景IoU的平均值,即:

MIoU= 1 k i=1 k TP FN+FP+TP (13)

式中:k为类别数本文是二分类任务,所以k = 2,其余参数与式(11)一致。

3.4. 消融实验结果与分析

本文将未改进的Unet作为基础模型,实验配置和实验参数保持一致。为了验证Leaky ReLU激活函数、LSKA模块和BCE Loss损失函数的有效性,以上述改进不同的组合进行消融实验,具体结果如表1所示。

Table 1. Experimental environment and configuration

1. 消融实验结果

Group

ReLU

Leaky ReLU

LSKA

BCE Loss

P (%)

Dice (%)

MioU (%)

1

92.53

82.21

82.43

2

92.65

82.43

82.71

3

94.15

83.89

84.69

4

93.12

82.74

83.06

5

94.82

84.08

84.92

通过表1,Group 1使用ReLU激活函数为原始Unet模型。Group 2是将原激活函数替换为Leaky ReLU激活函数,模型的P、Dice和MIoU分别提高了0.12%、0.22%和0.28%。Group 3在Group 2的基础上引入了LSKA模块,P、Dice和MIoU分别提高了1.62%、1.68%和2.26%。Group 4是在Group 2的基础上将Dice Loss替换为BCE Loss,P、Dice和MIoU分别提高了0.59%、0.53%和0.63%。Group 5是本文改进模型LSKA-Unet,相比于原始的Unet模型,同时各项指标提升的最为明显,P、Dice和MIoU分别提高了2.29%、1.87%和2.49%,这充分说明了改进模型和改进模块的有效性。证实了改进的Unet模型在医学图像领域的分割性能有着显著提升,适用于细胞的分割,改进前后Unet模型分割效果对比如图4所示。

Figure 5. Comparison of segmentation results between Unet and LSKA-Unet models

5. Unet和LSKA-Unet模型的分割结果对比

根据图5可以发现,细胞边界和细胞核的边界非常接近,使得难以检测到该处的细胞边界或是有效的分割出该处的细胞边界。在图5的蓝色箭头处可以发现改进后的Unet模型分割效果更好,保证了细胞边界的连续性。

3.5. 对比实验结果与分析

为了进一步验证本文LSKA-Unet模型的优越性,选取主流的语义分割模型FCN [9]、SegNet [10]、BASNet [11]、PANet [12]、Unet和本文算法LSKA-Unet在Drosophila EM数据集上进行对比实验。各模型的分割效果如表2所示。

表2结果显示,在细胞分割领域,FCN模型的分割效果表现最差,然后依次为BASNet、SegNet和PANet。分割效果最好的是本文模型LSKA-Unet,其次是应用较为广泛的Unet。本文改进的Unet模型,相较于FCN、SegNet、BASNet、PANet和Unet,以MIoU作为细胞分割性能的评价指标,分别提升了14.27%、4.31%、5.6%、3.36%和2.49%。并且在P和Dice系数的评价指标均有提升,证实了本文模型的综合性能。

Table 2. Comparison of experimental results

2. 对比实验结果

Model

P (%)

Dice (%)

MioU (%)

FCN

85.21

63.91

70.65

SegNet

90.13

78.87

80.61

BASNet

89.49

76.48

79.32

PANet

91.68

79.69

81.56

Unet

92.53

82.21

82.43

LSKA-Unet

94.82

84.08

84.92

3.6. 曲线训练结果与分析

在同等的实验环境下,Unet和LSKA-Unet的训练损失值变化曲线如图6所示,两种模型经过50次epoch后趋于平滑,都已充分学习。LSKA-Unet模型损失值下降更快并且更为稳定,训练效果明显优于Unet。再次证实了本文引入的BCE损失函数和改进模型的有效性。

Figure 6. Comparison of training loss of Unet model before and after improvement

6. 改进前后Unet模型训练损失对比

4. 结论

本文提出了一种基于LSKA注意力机制的LSKA-Unet模型,用于解决医学图像的细胞分割不连续的难题。具体的研究内容及结论具体包括以下几个方面:

通过引入大型可分离核注意力机制(LSKA),模型能够更有效地捕捉图像中的关键特征,从而提高了分割的准确性和连续性。

将传统ReLU激活函数替换为Leaky ReLU激活函数使得梯度消失的问题得以缓解,并且在处理噪声和不均匀亮度的医学图像时表现得更加鲁棒,提升了分割结果的稳定性。

采用BCE Loss (Binary CrossEntropy Loss)替换原Dice Loss损失函数,针对细胞分割这类二分类问题,提供了更精确的像素级分类,尤其是在处理边界模糊或重叠的细胞时,显著提高了模型的准确性。

在Drosophila EM数据集上经过大量的实验验证,LSKA-Unet模型的P、Dice系数和MIoU三大指标达到了94.82%、84.08%和84.92%,相比于原始Unet模型分别增长了2.29%、1.87%和2.49%,满足医学图像上细胞分割的精度要求,同时也证实了LSKA-Unet模型在医学图像分割领域的有效性和优越性。在后续研究中,将运用模型剪枝和知识蒸馏技术,以保证模型原有的精度而实现更少的参数量。

参考文献

[1] Ling, Y., Wang, Y., Dai, W., Yu, J., Liang, P. and Kong, D. (2024) MTANet: Multi-Task Attention Network for Automatic Medical Image Segmentation and Classification. IEEE Transactions on Medical Imaging, 43, 674-685.
https://doi.org/10.1109/tmi.2023.3317088
[2] Roy, S., Menapace, W., Oei, S., Luijten, B., Fini, E., Saltori, C., et al. (2020) Deep Learning for Classification and Localization of COVID-19 Markers in Point-Of-Care Lung Ultrasound. IEEE Transactions on Medical Imaging, 39, 2676-2687.
https://doi.org/10.1109/tmi.2020.2994459
[3] Long, J., Shelhamer, E. and Darrell, T. (2015) Fully Convolutional Networks for Semantic Segmentation. 2015 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), Boston, 7-12 June 2015, 3431-3440.
https://doi.org/10.1109/cvpr.2015.7298965
[4] Ronneberger, O., Fischer, P. and Brox, T. (2015) U-Net: Convolutional Networks for Biomedical Image Segmentation.
http://arxiv.org/abs/1505.04597
[5] Zhang, Y., Wu, J., Liu, Y., Chen, Y., Wu, E.X. and Tang, X. (2021) MI-UNet: Multi-Inputs UNet Incorporating Brain Parcellation for Stroke Lesion Segmentation from T1-Weighted Magnetic Resonance Images. IEEE Journal of Biomedical and Health Informatics, 25, 526-535.
https://doi.org/10.1109/jbhi.2020.2996783
[6] Wang, M., Jiang, H., Shi, T. and Yao, Y. (2022) HD-RDS-UNet: Leveraging Spatial-Temporal Correlation between the Decoder Feature Maps for Lymphoma Segmentation. IEEE Journal of Biomedical and Health Informatics, 26, 1116-1127.
https://doi.org/10.1109/jbhi.2021.3102612
[7] Lau, K.W., Po, L.M. and Rehman, Y.A.U. (2023) Large Separable Kernel Attention: Rethinking the Large Kernel Attention Design in CNN.
http://arxiv.org/abs/2309.01439
[8] Xu, B., Wang, N.Y., Chen, T.Q. and Li, M. (2015) Empirical Evaluation of Rectified Activations in Convolutional Network.
http://arxiv.org/abs/1505.00853
[9] Yang, F., Li, X. and Shen, J. (2021) MSB-FCN: Multi-Scale Bidirectional FCN for Object Skeleton Extraction. IEEE Transactions on Image Processing, 30, 2301-2312.
https://doi.org/10.1109/tip.2020.3038483
[10] Chen, X., Huang, W., Haller, M.C. and Pittman, R. (2021) Rain-Contaminated Region Segmentation of X-Band Marine Radar Images with an Ensemble of SegNets. IEEE Journal of Selected Topics in Applied Earth Observations and Remote Sensing, 14, 141-154.
https://doi.org/10.1109/jstars.2020.3043739
[11] Bo, W., Liu, J., Fan, X., Tjahjadi, T., Ye, Q. and Fu, L. (2022) BASNet: Burned Area Segmentation Network for Real-Time Detection of Damage Maps in Remote Sensing Images. IEEE Transactions on Geoscience and Remote Sensing, 60, 1-13.
https://doi.org/10.1109/tgrs.2022.3197647
[12] Piao, Y., Jiang, Y., Zhang, M., Wang, J. and Lu, H. (2023) PANet: Patch-Aware Network for Light Field Salient Object Detection. IEEE Transactions on Cybernetics, 53, 379-391.
https://doi.org/10.1109/tcyb.2021.3095512