论文阅读|小目标分割算法ASF-YOLO
论文题目:ASF-YOLO: A novel YOLO model with attentional scale sequence fusion for cell instance segmentation
论文链接: https://arxiv.org/abs/2312.06458
代码链接:https://github.com/mkang315/ASF-YOLO
摘要(Abstract)
作者提出了一种新颖的基于 “You Only Look Once”(YOLO)的注意力尺度序列融合框架(ASF-YOLO),该框架结合了空间和尺度特征,可实现准确、快速的细胞实例分割。在 YOLO 分割框架的基础上,采用尺度序列特征融合(SSFF)模块来增强网络的多尺度信息提取能力,并采用三重特征编码器(TPE)模块来融合不同尺度的特征图,以增加细节信息。进一步引入了通道和位置注意力机制(CPAM),将 SSFF 和 TPE 模块整合在一起,重点关注信息通道和与空间位置相关的小目标,以提高检测和分割性能。在两个细胞数据集上进行的实验验证表明,所提出的 ASF-YOLO 模型具有出色的分割精度和速度。它在 2018 Data Science Bowl 数据集上实现了 0.91 的边界框 mAP、0.887 的掩膜 mAP 和 47.3 FPS 的推理速度,优于最先进的方法。源代码见 https://github.com/mkang315/ASF-YOLO。
关键词:医学图像分析、小目标分割、You Only Look Once(YOLO)、序列特征融合、注意力机制
1 引言(Introduction)
随着样品制备技术和显微成像技术的快速发展,细胞图像的定量处理和分析在医学和细胞生物学等领域发挥着重要作用。基于卷积神经网络(Convolutional Neural Networks,CNN),可以通过神经网络训练学习不同细胞图像的特征信息,具有很强的泛化性能。两阶段 R-CNN 系列 [1, 2, 3] 及其单阶段变体 [4, 5] 是用于实例分割任务的基于 CNN 的经典框架。在最近的研究中,"You Only Look Once(YOLO)"系列 [6, 7, 8, 9] 已成为实时实例分割中最快、最准确的模型之一。与两阶段分割模型相比,YOLO 实例分割模型由于采用了单阶段设计思想和特征提取功能,因此具有更好的准确性和速度。然而,细胞实例分割的难点在于对象小、密集、重叠,以及细胞边界模糊,从而导致细胞分割的准确性较差。细胞实例分割需要对细胞图像中不同类型的物体进行精确细致的分割。如图 1 所示,由于细胞形态、制备方法和成像技术的不同,不同类型的细胞图像在颜色、形态、纹理和其他特征信息方面存在很大差异。
典型的 YOLO 框架结构由三个主要部分组成:骨干backbone、颈部neck和头部head。YOLO 的骨干网络是一个卷积神经网络,用于提取不同粒度的图像特征。具有 53 个卷积层的 Cross Stage Partial [10] Darknet(CSPDarknet53)[11] 是在 YOLOv4 [12] 的基础上修改的,被作为 YOLOv5 [8] 的骨干网络,包含 C3(CSP bottleneck 包括 3 个卷积层)和 ConvBNSiLU 模块。 YOLOv8[9]的骨干网络与 YOLOv5 唯一的不同之处,C3 模块被 C2f(CSP bottleneck 包括 2 个直连的卷积层)模块替代。如图 2 所示,YOLOv5 和 YOLOv8 主干网中的 1-5 级特征提取分支 P1、P2、P3、P4、P5 分别对应于与每个特征图相关的 YOLO 网络输出。YOLOv5 v7 和 YOLOv8 是首个基于 YOLO 的主流架构,除检测和分类外,还能处理分割任务。在 YOLOv5 的特征提取阶段,采用了由多个 C3 模块堆叠而成的 CSPDarkNet53 主干网络,然后将主干网络的三个有效特征分支 P3、P4 和 P5 作为特征金字塔网(FPN)结构的输入,在颈部建立多尺度融合结构。在特征层的解码过程中,与主干网络有效特征分支相对应的三个不同大小的头被用于物体的边界框预测。在对 P3 特征进行上采样后,逐像素解码作为目标的分割掩码预测,以完成目标的实例分割。在分割头中,三个尺度的特征输出三个不同的锚框,掩码原型模块负责输出原型掩码,经过处理后得到检测框和分割掩码,用于实例分割任务。
本文提出了一种细胞图像的单阶段实例分割模型ASF-YOLO,该模型将注意力尺度序列融合(Attentional Scale Sequence Fusion)融入了 You Only Look Once(YOLO)框架中。在特征提取阶段,首先使用 CSPDarknet53 骨干网络从细胞图像中提取多维特征信息。在颈部部分提出了用于细胞实例分割的新型网络设计。这项工作的贡献总结如下:
1)针对不同类型细胞的多尺度问题以及小细胞的目标检测和分割问题,设计了尺度序列特征融合(Scale Sequence Feature Fusion,SSFF)模块和三重特征编码器(Triple Feature Encoder,TFE)模块,在路径聚合网络(Path Aggregation Network,PANet)[13]结构中融合从主干网提取的多尺度特征图。
2)然后,设计了一个通道和位置注意力机制(Channel and Position Attention Mechanism,CPAM)来合并 SSFF 和 TFC 模块的特征信息,从而进一步提高实例分割的准确性。
3)在训练阶段,利用 EIoU [14]最小化边界框与锚点框的宽度和高度之差,以优化边界框位置损失,在后处理阶段,使用软非最大值抑制(Soft NonMaximum Suppression,Soft-NMS)[15] 来改善密集重叠细胞问题。
4)将提出的 ASF-YOLO 模型应用于高难度的密集重叠和各种细胞类型的实例分割任务。据作者描述,这是第一项利用基于 YOLO 模型进行细胞实例分割的工作。对两个基准细胞数据集的评估表明,与其他最先进的方法相比,该方法的检测精度和速度都更胜一筹。
2 相关工作(Related work)
2.1 细胞实例分割(Cell instance segmentation)
细胞实例分割可进一步帮助完成图像中的细胞计数任务,而细胞图像的语义分割则无法完成。深度学习方法提高了自动细胞核分割的准确性[16]。Johnson 等人[17]、Jung 等人[18]、Fujita 等人[19]和 Bancher 等人[20]提出了基于 Mask R-CNN [2]同时检测和分割细胞的改进方法。Yi 等人[21] 和 Cheng 等人[22] 利用单次多框检测器(Single-Shot multi-box Detector, SSD)[23] 方法检测和分割神经细胞实例。Mahbod 等人[24] 采用基于 U-Net [25] 模型的语义分割算法进行细胞核分割。带有注意力机制的 SSD 和 U-Net 混合模型[19]或 U-Net 和 Mask R-CNN 模型[26]在细胞实例分割数据集上取得了一定的提升。BlendMask [27] 是一种细胞核实例分割框架,带有扩张卷积聚合模块和上下文信息聚合模块。Mask R-CNN 是一种两阶段物体分割框架,速度较慢。SSD、U-Net 和 BlendMask 是统一的端到端(即单阶段)框架,但在分割密集细胞和小细胞时性能较差。
2.2 改进的YOLO用于实例分割(Improved YOLO for instance segmentation)
最近,针对实例分割任务对 YOLO 的改进主要集中在注意力机制、改进的骨干或网络以及损失函数上。挤压和激发(Squeeze-and-Excitation,SENet)[28] 块被集成到改进的 YOLACT[6] 中,用于识别显微图像中的瘤胃原生动物[29]。YOLOMask[30]、PR-YOLO[31]和YOLO-SF[32] 利用卷积块注意模块(CBAM)[34]增强了YOLOv5[8]和YOLOv7-Tiny[33]。改进后的主干网络中加入了有效的特征提取模块,使 YOLO 特征提取过程更加高效[35, 36]。YOLO-CORE [37]通过使用设计的多阶约束(由极距损失和扇形损失组成)进行显式直接轮廓回归,有效增强了实例的掩码。此外,另一种 YOLOMask [38] 和 YUSEG [39] 相结合的混合模型,将优化的 YOLOv4 [12] 和带有语义分割 U-Net 网络的原始 YOLOv5s ,以确保实例分割的准确性。
3 提出的ASF-YOLO(The proposed ASF-YOLO model)
3.1 总体框架(Overall architecture)
图 3 显示了提出的 ASF-YOLO 框架,该框架结合了空间和多尺度特征,用于细胞图像实例分割。作者开发了一种新颖的特征融合网络架构,由可以为小物体分割提供互补信息的两个主要组件网络组成:(1) SSFF 模块结合了来自多尺度图像的全局或高级语义信息;(2) TFE 模块可捕捉小目标物体的局部细节。局部和全局特征信息的融合可以生成更精确的分割图。作者对从骨干网络中提取的 P3、P4 和 P5 的输出特性进行了融合。首先,SSFF 模块旨在有效融合 P3、P4 和 P5 的特征图,以捕捉不同空间尺度、不同大小和形状的不同细胞类型。在 SSFF 中,P3、P4 和 P5 特征图被归一化为相同的大小,经过上采样,然后叠加在一起,作为三维卷积的输入,以结合多尺度特征。其次,作者开发了 TFE 模块,通过在空间维度上拼接大、中、小三种不同尺寸的特征来捕捉小目标的细节信息,从而增强对密集细胞小目标的检测。然后通过 PANet 结构将 TFE 模块的细节信息整合到每个特征分支中,再与 SSFF 模块的多尺度信息结合到 P3 分支中。还在 P3 分支中引入了通道和位置注意力机制(Channel and Position Attention Mechanism,CPAM),以充分利用高级多尺度特征和细节特征。CPAM 中的通道和位置注意力机制可分别捕获信息通道和细化与细胞等小物体相关的空间定位,从而提高其检测和分割精度。
3.2 尺度序列特征融合模块(Scale sequence feature fusion module)
针对细胞图像的多尺度问题,现有文献采用特征金字塔结构进行特征融合,其中只采用了和或并集来融合金字塔特征。然而,各种特征金字塔网络结构无法有效利用所有金字塔特征图之间的相关性。作者提出了一种新的尺度序列特征融合方法,它能更好地结合深度特征图的高维信息和浅层特征图的详细信息。尺度空间是沿着图像的尺度轴构建的,它不仅表示一个尺度,还表示一个物体可能具有的各种尺度的范围[40]。比例意味着图像的细节。模糊的图像可能会丢失细节,但图像的结构特征可以保留。作为 SSFF 输入的缩放图像可以通过以下方法获得
F σ ( w , h ) = G σ ( w , h ) × f ( w , h ) G σ ( w , h ) = 1 2 π σ 2 e − ( w 2 + h 2 ) / 2 σ 2 \begin{aligned}F_\sigma(w,h)&=G_\sigma(w,h)\times f(w,h)\\G_\sigma(w,h)&=\frac1{2\pi\sigma^2}e^{-(w^2+h^2)/2\sigma^2}&\end{aligned} Fσ(w,h)Gσ(w,h)=Gσ(w,h)×f(w,h)=2πσ21e−(w2+h2)/2σ2
其中, f ( w , h ) f (w, h) f(w,h) 表示宽度为 w w w、高度为 h h h 的二维输入图像。 F σ ( w , h ) F_\sigma(w,h) Fσ(w,h)通过使用二维高斯滤波器 G σ ( w , h ) G_σ(w, h) Gσ(w,h)的一系列卷积平滑生成, σ σ σ 是用于卷积的二维高斯滤波器标准偏差的缩放参数。
这些生成的图像具有相同的分辨率,但尺度不同。因此,可以将不同大小的特征图视为尺度空间,将不同分辨率的有效特征图调整为相同分辨率进行拼接。受对多个视频帧进行二维和三维卷积操作[41]的启发,作者将不同尺度的特征图进行水平堆叠,并使用三维卷积提取其尺度序列特征。高分辨率特征图 P3 层包含了对小型目标的检测和分割至关重要的大部分信息,因此 SSFF 模块是基于 P3 层设计的。如图 3 所示,拟议的 SSFF 模块由以下部分组成:
- 使用 1×1 卷积层将 P4 和 P5 特征层的通道数改为 256。
- 使用近邻插值法 [42] 将其大小调整为 P3 层的大小。
- 使用取消挤压法(unsqueeze method)增加每个特征层的维度,将其从三维张量[高度、宽度、通道]变为四维张量[深度、高度、宽度、通道]。
- 然后沿深度维度将 4D 特征图串联起来,形成一个 3D 特征图,用于后续卷积。
- 最后,使用三维卷积、三维批量归一化和 SiLU[43]激活函数完成尺度序列特征提取。
3.3 三重特征编码模块(Triple feature encoding module)
为了识别密集重叠的小物体,可以通过放大图像来参考和比较不同尺度下的形状或外观变化。由于骨干网络的不同特征层具有不同的尺寸,传统的 FPN 融合机制只是对小尺寸的特征图进行上采样,然后将其分割或添加到上一层的特征中,忽略了大尺寸特征层丰富的细节信息。因此,作者提出了 TFE 模块,它可以分割大、中、小尺寸的特征,添加大尺寸的特征图,并进行特征放大,以完善详细的特征信息。
图 4 展示了 TFE 模块的结构。在进行特征编码之前,首先要调整特征通道的数量,使其与主要尺度特征保持一致。卷积模块处理大尺寸特征图(Large)后,将其通道数调整为 1C,然后采用最大池化 + 平均池化的混合结构进行降采样,这有助于保留高分辨率特征和细胞图像的有效性和多样性。对于小尺寸的特征图(Small),同样使用卷积模块来调整通道数,然后使用最近邻插值法进行上采样。这有助于保持低分辨率图像局部特征的丰富性,防止小目标特征信息的丢失。最后,将尺寸相同的大、中、小三个特征图进行一次卷积,然后在通道维度上进行拼接,具体方法如下。
F T F E = C o n c a t ( F l , F m , F s ) F_{TFE}=Concat(F_{l},F_{m},F_{\mathrm{s}})