计算机系统应用  2024, Vol. 33 Issue (2): 159-165   PDF    
基于改进U2Net的岩石薄片图像分割
舒小锋1, 吴晓红1, 卿粼波1, 滕奇志1, 罗彬彬2     
1. 四川大学 电子信息学院, 成都 610065;
2. 成都西图科技有限公司, 成都 610065
摘要:了解岩石的孔隙度、孔径分布、孔隙连通性等特征对于油气的寻找和开采有着重要的意义, 而这些特征的分析和判断需要借助岩石薄片图像分割技术. 岩石薄片图像有大量细小颗粒, 这些颗粒之间的边缘特征十分相似, 无法做出精准的区分, 同时制造切片过程中染色不均会造成薄片孔隙的颜色特征不平衡而导致无法分割. 因此为了改善岩石薄片分割效果, 本文提出基于一种改进的U2Net的分割算法. 主要内容如下: (1)以U2Net网络为骨干进行改进, 结合coordinate attention注意力机制, 用来提高模型对图像特征的表达能力. (2)通过引入多尺度特征提取模块, 增加卷积层的感知区域, 且能够利用特征图的多尺度特征信息. 实验证明, 该方法与传统分割方法和其他分割网络相比在较小颗粒的分割上表现更好, 所提出的算法具有较高的分割准确度和鲁棒性.
关键词: 注意力机制    岩石薄片图像    图像分割    U2Net    多尺度特征提取    
Image Segmentation of Rock Thin Sections Based on Improved U2Net
SHU Xiao-Feng1, WU Xiao-Hong1, QING Lin-Bo1, TENG Qi-Zhi1, LUO Bin-Bin2     
1. College of Electronics and Information Engineering, Sichuan University, Chengdu 610065, China;
2. Chengdu Xitu Technology Co. Ltd., Chengdu 610065, China
Abstract: It is important to understand the characteristics of rock porosity, pore size distribution, and pore connectivity for oil and gas exploration and exploitation, and the analysis and judgment of these characteristics need to rely on the image segmentation technology of rock thin sections. There are a large number of fine particles in the images of rock thin sections, and the edge features among these particles are very similar, which cannot be accurately distinguished. Meanwhile, uneven staining during section manufacturing will cause unbalanced color characteristics of the pores of the thin sections, resulting in the inability to segment. Therefore, to improve the segmentation effect of rock thin sections, this study proposes an improved segmentation algorithm based on U2Net. The main contents are as follows. (1) The U2Net network is adopted as the backbone to improve the model’s ability to express image features, and coordinate attention is combined to enhance the ability to express image features. (2) The introduction of a multi-scale feature extraction module enlarges the receptive field of the convolutional layers and enables the utilization of multi-scale feature information from the feature map. Empirical evaluations demonstrate that the proposed method outperforms conventional segmentation techniques and other state-of-the-art segmentation networks in small particle segmentation. Additionally, the proposed algorithm exhibits superior segmentation accuracy and robustness.
Key words: attention mechanism     image of rock thin section     image segmentation     U2Net     multi-scale feature extraction    

利用岩石薄片图像进行岩石颗粒的粒度分析是地质实验室的一种常用方法, 对岩石薄片的分割是粒度分析和矿物识别的前提[1]. 岩石薄片孔隙是流体可以存在于岩石中的先决条件[2], 储集层的岩石中独特的孔缝结构能够储集和渗流油气[3], 对孔隙结构精准地分析有助于增加石油采收率, 评价储集层和对油气产能的预测.

薄片图像的分割往往是对岩石薄片图像处理的第一步, 使地质研究人员所关注的孔隙、矿物等目标更加清晰地呈现, 但目前岩石薄片图像分割面临多目标相似特征、样本差异性和设备变化等挑战. 并且在传统岩石分割中, 该项工作一般由专业人员根据个人专业知识完成; 但这种方式工作量较大, 且分割结果易受主观因素干扰, 缺乏确定性与稳定性[4].

由于传统图像分割算法中图像大多都缺乏全局的语义信息, 因此泛化性较差, 无法很好地应对复杂多变的场景. 伴随着深度学习图像技术的发展, 传统分割方法在分割速度和精度方面已经远不如深度学习技术. 其中编解码基础模型、全卷积网络[5]、空洞空间特征金字塔[6]等技术的出现使得分割后输出的图像分辨率更高. 常见的编解码基础模型如, Ronneberger等提出的UNet [7]网络结构, 2020年, 阿尔伯塔大学在UNet基础上提出了U2Net[8]网络结构, 该网络结构为RSU结构, 每个RSU就是一个缩版的UNet, 最后通过FPN的跳跃连接构建完整模型. 许多学者利用这些网络在医学、遥感等领域取得了一定的成绩, 例如贺军义等[9]使用改进的UNet网络对混凝土CT图像孔隙裂隙进行分割, 有效提升了图像中裂隙与孔隙的分割精度, 梁燕等[10]利用UNet3+编解码网络提出遥感影像建筑变化检测. 而在岩石领域, 也有不少学者将卷积神经网络应用到岩石图像的处理中, 例如刘勇等[11]提出基于改进的SKnet和Bi-GRU对岩石薄片图像矿物识别, 董领等[12]提出了基于改进的SLIC的岩心颗粒图像边缘分割算法.

但目前的图像分割方法在分割颗粒时容易受到复杂背景的干扰, 并且颗粒之间存在相互粘连, 使得分割并不准确. 因此本文提出一种基于U2Net网络结构结合coordinate attention注意力机制[13]和receptive field block-small模块[14]进行改进的图像分割网络, CA来突出重要的空间位置, 从而增强网络对空间上局部细节的感知能力, 改善对细节和边缘的重建. RFB-S提取更丰富的特征并增强网络的表征能力, 增强网络鲁棒性和对细小颗粒的分割效果.

1 U2Net网络

大多数的显著目标检测网络有一个共性, 就是注重利用现有的主干网络提取深层特征, 比如ResNet[15], ResNeXt[16], DenseNet[17]等, 这些主干最终都是为图像分类任务而设计的, 他们提取的特征代表语义, 而不是局部细节和全局对比信息. U2Net很好地解决了多层次与多尺度特征提取的问题, 它的可以将图片的前景和背景分割出来, 相当于语义分割中的二分类任务.

U2Net的体系结构是一个两层嵌套的U结构, 主要包括残差U型块 (residual U-blocks, RSU)和连接RSU的外层U型结构[8], 这样能够捕捉到更多的上下文信息, 并且整个架构的深度有所增加没有增加计算成本, 因为在这些RSU块中使用了池化操作. 该算法在设计过程中无需使用图像分类的预训练主干网络, 在不会显著增加内存和低计算成本的前提下可以提取更高分辨率的特征图[18], 因此, 本文选取U2Net网络来分割图像. 其中, RSU由3个部分构成: 输入卷积层、U型结构以及残差连接层连接输入层和中间层, RSU结构图1所示.

图 1 RSU结构

图1中, CinCout分别为输入层、输出层的通道数; HW分别为输入的长和宽; M代表内层通道数; L代表RSU层级数量, 如图1为7层, 简称RSU-7, 对应的是整体网络结构中的Encoder1和Decoder1, RSU的网络深度越深, 池操作越多, 提取到的局部特征和全局特征越丰富. 图1中绿色方框代表卷积层, 每个卷积层中都包含3个操作: 卷积运算(Conv)、批规范化(batch normalization, BN)、激活函数非线性化(ReLU). RSU结构左侧是编码结构, 蓝色方框代表卷积层加降采样操作, 通过多层3×3卷积操作获得多尺度特征, 降采样操作增大感受野; 结构右侧是解码结构, 黄色方块代表卷积层加上采样操作, 上采样是为了将卷积层输出的特征编码为高分辨率的特征图.

U2Net是一个两层嵌套的U型结构, 主要由3部分组成6 阶段的编码器、5 阶段的解码器构成以及连接最后阶段编码器和解码器的特征融合输出模块, 每一个阶段由一个配置良好的RSU填充. U2Net整体结构如图2所示.

图 2 U2Net结构

图2中, 在En_1和De_1两个阶段的编解码中, 使用了RSU-7模块填充, En_2和De_2则采用的RSU-6, 即RSU的层级数量为6层, En_3和De_3、En_4和De_4分别使用的RSU-5和RSU-4模块, 以此获得不同尺度的空间信息; 而En_5、De_5、En_6阶段的编解码器使用的是在RSU-4基础上改进的模块称为RSU-4F, 这一模块主要是将卷积上下采样替换成膨胀卷积, 使得中间层的特征图的大小与输入层一致. 最后部分是将显著特征图进行融合, 网络首先通过3×3卷积和Sigmoid 激活函数将不同阶段解码器En_6, De_5, De_4, De_3, De_2和De_1生成的S6、S5、S4、S3、S2、S1显著性概率图用线性插值进行上采样到输入图像的大小; 通过一个级联互相融合, 然后通过1×1卷积和一个Sigmoid函数将概率图融合, 最终输出显著. 由于U2Net网络结构复杂且深, 通过 RSU 和跳跃连接能够提取图片不同尺寸的级内和级间信息, 但在连接中容易出现非颗粒区域等无效特征[19]的保留和缺陷边缘等信息的丢失[20], 因此要从提高有效特征的提取能力和降低信息的丢失去改进网络.

2 改进的U2Net网络 2.1 整体网络结构

由于薄片图像背景的复杂性, 并且包含细小颗粒, 且存在颗粒边界难以辨认等情况, 使用U2Net网络难以精准提取目标, 时常导致图像欠分割或者过分割. 因此, 本文把CA注意力机制和RFB-S模块与U2Net网络结合, 改进出一种新的目标分割模型. 本文在U2Net的左侧编码器底层下采样部分加入CA注意力机制和RFB-S模块, 以加强对关注对象细节信息的提取能力, 更好地捕捉到图片中那些细微特征, 增强目标边缘定位能力提高分割结果的准确性. 改进后的网络结构如图3所示.

由于编码器负责逐层提取图像的语义信息和边缘特征, 通过引入RFB-S模块可以增强网络的分割精度和增强感受野, 减少网络的计算量, 引入CA注意力机制, 可以使网络更加关注重要的空间位置, 提高特征的表达能力和感知能力. 因此在编码网络部分, 输入特征图经过RSU-7模块后, 在下采样部分引入RFB-S模块和CA模块. 在RFB-S模块中通过引入不同尺度的卷积核并且该卷积核拥有更小的尺寸和参数量, 通过空洞卷积和多尺度卷积核后将输出的图像输入至CA模块. 而在CA模块中, 通过对特征图进行编码, 可以获得一对具有方向感知和位置敏感性的特征图, 其中一个特征图能够捕捉远程依赖关系, 而另一个特征图则保留了精确的位置信息. 这两种特征图相互补充, 并应用于输入特征图, 从而提升对所关注对象的表征能力. 最后输出的特征图用于下一阶段的RSU模块进行特征提取和后续进行双线性插值法进行上采样. 在编码阶段的其余4个编码器后引入同样的结构构成改进后的网络.

图 3 改进后的网络结构

2.2 Coordinate Attention注意力机制

大部分注意力机制用于深度神经网络能够带来很好的功能提高, 例如SE[21]、BAM[22]、CBAM[23]等. 但由于普通卷积难以建模信道关系, SE很容易忽略空间定位信息, 但在视觉中, 目标的空间结构却扮演着至关重要的角色. 而BAM和CBAM对通道进行全局池化只能引入短距离的位置信息, 长距离的相关信息无法捕获. 基于此提出一种新型的坐标注意力机制, 命名为coordinate attention. 该注意力在有效保持精确的位置信息的同时可以获取范围更广的区域信息. CA注意力机制模块结构如图4所示.

针对通道注意力存在的问题, CA通过精确的位置信息可以同时对通道相关性与远程依赖性建模, 也就是分为CIE (coordinate information embedding)与CAG (coordinate attention generation)两个步骤, 分别称作Coordinate信息嵌入和Coordinate 注意生成. 其中CIE通过一对一维特征编码操作实现了全局池化的效果使模块能够获取精确的位置信息, 再通过CAG充分利用CIE所产生的表征就使得网络获取更大范围信息的同时还可以有效保持精确的位置信息.

图 4 CA注意力模块结构

2.3 RFB-S模块

Receptive field block (RFB)类似于人类视觉的RFs结构, 考虑感受野的离心率和大小, 即便采用轻量网络结构的情况下, 仍然有效提取到高度区分性的特征, 其可以看成是Inception的基础上加入空洞卷积. 而receptive field block-small (RFB-S)模块相比于RFB模块具有更小的卷积核尺寸和参数量, 可以在一定程度上减少网络的参数数量和计算复杂度. RFB-S结构如图5所示.

图 5 RFB-S结构

与RFB相比, RFB-S将原来的5×5卷积用3×3卷积替换, 既减少模块参数数量, 还增加了其非线性的关系能力, 并且还用1×3和3×1卷积层代替原先的3×3卷积层. 将特征图输入模块, 先采用1×1的卷积降维通道数量, 再经过卷积层1×3、3×1和3×3得到多尺度特征, 最后为了扩大模块的感受野, 对模块进行空洞卷积(1、3、5扩张率), 再将合并输出4个卷积核实现特征融合.

3 实验结果与分析

为了验证改进后的U2Net网络有效性, 本文将该算法与部分分割网络SegNet、R2UNet、U2Net、Deeplabv3+以及传统阈值分割进行对比, 使用的数据集是岩石薄片图像. 此外, 为了分析本文算法中提出的改进方法对图像分割结果的影响, 还设计了消融实验来评估算法的有效性.

3.1 数据集

本次实验所用的岩石薄片数据集来源于GitHub中的公开项目, 此数据集的许可与GitHub仓库中的许可一致, 根据MIT许可证提供. 薄片图像通过显微镜观察拍摄获取得到, 岩石薄片图和标注图如图6所示, 考虑到GPU承受能力并且为了在训练中提高效率并且降低训练的耗能, 将采集得到的原始图像进行裁剪, 总共得到1400张分辨率为256×256大小的图像, 使用图形图像注释工具Labelme对图像进行标注, 得到JSON文件并转换成png格式的图像.

图 6 岩石薄片图和标注图

3.2 参数设置与训练环境

本文的分割算法实验基于深度学习框架PyTorch 1.8.0并结合编程语言Python在工作站上进行环境搭建, 工作站配置如下: 操作系统为Ubuntu 18.04 LTS系统, 处理器为Intel(R)Core(TM)i7-9700, GPU为NVDIA GeForce RTX 2080Ti, 软件环境为CUDA 10.2.

本文训练时使用Adam优化器优化网络, 初始学率$ lr=0.001,\ betas=(0.9,0.999),\ eps=1\times10^{-8} $, 权重衰减为0; 经过多次对比实验后, 将训练次数设定为200, 批量大小设置为2, 每100次保存一次模型参数文件, 最后在网络内使用最优的参数文件对测试集进行预测.

3.3 实验评价指标

本文将采用准确率(Accuracy)、交并比(IoU)、均像素精度(MPA)、均交并比(MIoU)和频权交并比(FWIoU)、F1分数来对分割效果和网络模型进行评估. 本文中选取分割出的颗粒作为正例, 背景部分(孔隙)作为反例.

在二分类问题中, 混淆矩阵的4个基本元素如下.

TP: 模型预测为正例, 实际是正例.

FP: 模型预测为正例, 实际是反例.

FN: 模型预测为反例, 实际是正例.

TN: 模型预测为反例, 实际是反例.

$ Accuracy=\frac{TP+TN}{TP+TN+FP+FN} $ (1)
$ IoU=\frac{TP}{TP+FP+FN} $ (2)
$ MPA = \frac{1}{2}\left( {CP{A_1} + CP{A_2}} \right) $ (3)
$ MIoU=\frac{1}{2}(IoU_1+IoU_2) $ (4)
$ FWIoU = \frac{{TP + FN}}{{TP + FP + TN + FN}} \times \frac{{TP}}{{TP + FP + FN}} $ (5)
$ {F_1} = \frac{{2 {\textit{Precision}} \times R}}{{{\textit{Precision}} + R}} $ (6)

其中, CPA1CPA2分别为正例和反例的精准率, 正例准确率${\textit{Precision}} = TP/(TP + FP)$, 反例准确率, ${\textit{Precision}} = TN/(TN + FN)$召回率$R = TP/(TP + FP)$, F1中为正例精准率. IoU1IoU2分别为正例和反例的交并比.

3.4 U2Net有效性验证

本文利用注意力机制的优点, 提出了改进U2Net的图像分割方法, 为定性分析该方法的有效性和合理性, 将本文方法与主流深度学习分割网络PSPNet、R2UNet、Deeplabv3+、U2Net以及传统阈值分割相比较. 将本文所改进的网络与上述分割网络用同一数据集进行训练, 并将所有训练完成的网络以及传统阈值分割法使用同一测试集进行测试, 将分割后的结果使用相同的方法进行评价, 结果如表1所示, 图7为各个算法的可视化结果.

表 1 不同网络结果对比(%)

3.5 不同模块的消融实验

为了分析本文算法中提出的各改进方法对岩石薄片分割结果的影响, 设计了在岩石薄片数据集上进行消融实验来评估. 消融实验共分为4组, 第1组实验使用U2Net网络, 第2组实验在U2Net网络的基础上添加CA注意力机制, 第3组实验在U2Net网络的基础上添加RFB-S模块, 第4组则是在U2Net网络的基础上添加CA注意力机制和RFB-S模块, 实验所使用的训练参数一致. 消融实验具体结果如表2所示.

图 7 对比结果可视化

表 2 不同模块的消融实验 (%)

表2中可知, 两种模块都使得网络的效果有一定的提升, CA使得模型对不同位置特征之间的建模能力有所提高, 提高了分割的准确率. RFB-S模块通过构建不同空洞率的卷积层, 融合特征图, 增强了模型对不同尺度颗粒的检测能力, 改善了颗粒的沾粘情况. 再融合两个模块后增强网络的提取精度和鲁棒性, 证明了网络具有较好的分割效果.

4 结语

本文旨在解决岩石薄片图像颗粒在传统分割方法和深度学习图像分割方法上存在精确度低, 由于背景颜色的复杂无法有效地提取岩石薄片图像中的细节信息等问题, 提出了基于改进的U2Net的岩石薄片图像分割算法, 通过引入CA注意力机制来增强网络对颗粒的特征表达能力, 构建RFB-S来提升网络分割精度. 通过在实验中对比和验证, 证明了通过改进后的网络有效地提升了对薄片颗粒和孔隙的分割效果, 对图像中的细节信息、边界检测和区域分割都有了有效提升.

然而, 本文的工作还存在一些局限性和可以改进的方面. 首先, 虽然我们在岩石薄片图像上进行了广泛的实验, 但是我们的算法在其他类型的图像数据集上的适用性还需要进一步验证. 其次, 我们可以进一步提升算法的计算效率, 尤其是在大规模图像数据集上的应用场景下, 以提高算法的实际应用性. 未来将进一步改进和优化岩石薄片图像分割的相关算法, 探索更多的技术手段来提升岩石薄片图像分割的性能.

参考文献
[1]
蒋宜勤, 周基贤, 冯俊羲, 等. DBSCAN算法与数学形态学在岩石薄片图像分割中的应用. 微型机与应用, 2016, 35(17): 39-41, 48. DOI:10.19358/j.issn.1674-7720.2016.17.012
[2]
李霏, 陈铭谦, 赵御庭, 等. 岩石微观孔隙结构研究方法综述. 地下水, 2019, 41(6): 112-114.
[3]
张中亚. 砂岩薄片图像分割与识别研究[硕士学位论文]. 合肥: 中国科学技术大学, 2020. ]
[4]
吕锦涛. 基于超像素与半监督的岩石图像分割与识别. 工程科学与技术, 2023, 55(2): 171-183. DOI:10.15961/j.jsuese.202200459
[5]
Long J, Shelhamer E, Darrell T. Fully convolutional networks for semantic segmentation. Proceedings of the 2015 IEEE Conference on Computer Vision and Pattern Recognition. Boston: IEEE, 2015. 3431–3440.
[6]
Chen LC, Papandreou G, Schroff F, et al. Rethinking atrous convolution for semantic image segmentation. arXiv:1706.05587, 2017.
[7]
Ronneberger O, Fischer P, Brox T. U-Net: Convolutional networks for biomedical image segmentation. Proceedings of the 18th International Conference on Medical Image Computing and Computer-assisted Intervention. Munich: Springer, 2015. 234–241.
[8]
Qin XB, Zhang ZC, Huang CY, et al. U2Net: Going deeper with nested U-structure for salient object detection. Pattern Recognition, 2020, 106: 107404. DOI:10.1016/j.patcog.2020.107404
[9]
贺军义, 冯嘉莘, 焦华喆, 等. 基于改进UNet的混凝土CT孔隙裂隙分割方法. 中国矿业大学学报, 2023, 52(3): 615-624.
[10]
梁燕, 易春霞, 王光宇. 基于编解码网络UNet3+的遥感影像建筑变化检测. 计算机学报, 2023, 46(8): 1720-1733.
[11]
刘勇, 吴晓红, 滕奇志, 等. 基于改进的SKnet和Bi-GRU的岩石薄片图像矿物识别. 智能计算机与应用, 2023, 13(1): 104-111.
[12]
董领, 卿粼波, 何小海, 等. 基于改进的SLIC的岩心颗粒图像边缘分割算法. 智能计算机与应用, 2021, 11(9): 54-58.
[13]
Hou QB, Zhou DQ, Feng JS. Coordinate attention for efficient mobile network design. Proceedings of the 2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition. Nashville: IEEE, 2021. 13708–13717.
[14]
Liu ST, Huang D, Wang YH. Receptive field block net for accurate and fast object detection. Proceedings of the 15th European Conference on Computer Vision. Munich: Springer, 2018. 404–419.
[15]
He KM, Zhang XY, Ren SQ, et al. Deep residual learning for image recognition. Proceedings of the 2016 IEEE Conference on Computer Vision and Pattern Recognition. Las Vegas: IEEE, 2016. 770–778.
[16]
Xie SN, Girshick R, Dollár P, et al. Aggregated residual transformations for deep neural networks. Proceedings of the 2017 IEEE Conference on Computer Vision and Pattern Recognition. Honolulu: IEEE, 2017. 5987–5995.
[17]
Huang G, Liu Z, van der Maaten L, et al. Densely connected convolutional networks. Proceedings of the 2017 IEEE Conference on Computer Vision and Pattern Recognition. Honolulu: IEEE, 2017. 2261–2269.
[18]
严伟健. 基于U2-Net的多模态睡眠分期算法研究与应用[硕士学位论文]. 广州: 广东工业大学, 2022. ]
[19]
徐凯成. 基于U-Net的胃部肿瘤分割技术[硕士学位论文]. 上海: 上海工程技术大学, 2021. ]
[20]
Seo H, Huang C, Bassenne M, et al. Modified U-Net (mU-Net) with incorporation of object-dependent high level features for improved liver and liver-tumor segmentation in CT images. IEEE Transactions on Medical Imaging, 2020, 39(5): 1316-1325. DOI:10.1109/TMI.2019.2948320
[21]
Hu J, Shen L, Sun G. Squeeze-and-excitation networks. Proceedings of the 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition. Salt Lake City: IEEE, 2017. 7132–7141.
[22]
Park J, Woo S, Lee JY, et al. BAM: Bottleneck attention module. Proceedings of the 2018 British Machine Vision Conference. Newcastle: BMVA Press, 2018.
[23]
Woo S, Park J, Lee JY, et al. CBAM: Convolutional block attention module. Proceedings of the 15th European Conference on Computer Vision. Munich: Springer, 2018. 3–19.