ASF-YOLO开源 | SSFF融合+TPE编码+CPAM注意力,精度提升!
目录
2.1 Cell instance segmentation
2.2 Improved YOLO for instance segmentation
3.2 Scale sequence feature fusion module
3.3 Triple feature encoding module
3.4 Channel and position attention mechanism
4.5.1 Effect of the proposed methods
4.5.2 Effect of attention mechanisms
4.5.3 Effect of convolution module in the backbone
摘要
作者提出了一种新颖的注意力尺度序列融合基于YOLO框架(ASF-YOLO),该框架结合了空间和尺度特征,以实现精确快速的细胞实例分割。在YOLO分割框架的基础上,作者采用了尺度序列特征融合(SSFF)模块来增强网络的多尺度信息提取能力,并采用三特征编码器(TPE)模块将不同尺度的特征图进行融合,以增加详细信息。作者进一步引入了一种通道和位置注意力机制,以将SSFF和TPE模块集成起来,专注于具有信息量大、位置相关的较小目标,以提高检测和分割性能。
在两个细胞数据集上的实验验证表明,所提出的ASF-YOLO模型的分割准确性和速度显著提高。该模型在2018年数据科学竞赛数据集上实现了box mAP=0.91,mask mAP=0.887,推理速度达到47.3 FPS,超过了最先进的方法。
代码:https://github.com/mkang315/ASF-YOLO
论文:ASF-YOLO: A novel YOLO model with attentionalscale sequence fusion for cell instance segmentation
1 Introduction
随着样本制备技术和显微成像技术的快速发展,细胞图像的定量处理和分析在医学和细胞生物学等领域中发挥着重要作用。基于卷积神经网络(CNN),通过神经网络训练可以学习不同细胞图像的特征信息,具有较强的泛化性能。两阶段R-CNN系列及其一阶段变体是经典的基于CNN的实例分割任务框架。
在近年来的工作中,YOLO系列已成为实时实例分割的最快、最准确模型之一。由于一阶段设计理念和特征提取能力,YOLO实例分割模型比两阶段分割模型具有更好的准确性和速度。然而,细胞实例分割的困难之处在于细胞图像中的小、密、重叠物体以及细胞边界模糊,这导致细胞精度较低。细胞实例分割需要对细胞图像中不同类型物体的精确详细分割。如图1所示,由于细胞形态、制备方法、成像技术等因素的差异,不同类型的细胞图像在颜色、形态、纹理和其他特征信息方面存在较大差异。
YOLO框架结构通常包括三个主要部分:backbone、neck和head。YOLO的backbone网络是一个卷积神经网络,在不同的粒度下提取图像特征。Cross Stage Partial Darknet与CSPDarknet53是从YOLOv4修改而来,并被设计为YOLOv5的backbone网络,其中包含C3(包括3个卷积层)和ConvBNSiLU模块。在YOLOv5和YOLOv8的backbone中,第1-5级特征提取分支P1、P2、P3、P4和P5对应着与这些特征图相关的YOLO网络输出。YOLOv5 v7和YOLOv8是第一个主流基于YOLO的架构,可以处理分割任务(除了检测和分类)。
在YOLOv5的特征提取阶段,使用堆叠的多级CSPDarkNet53 backbone网络,然后将backbone网络的三个有效特征分支P3、P4和P5作为FPN结构的多 Scale 融合结构输入。在特征层的解码过程中,backbone网络的三个不同大小的头用于预测目标的边界框。在将P3特征上采样后,像素逐像素解码作为目标分割Mask预测,以完成目标实例分割。在分割头中,三个尺度的特征输出三个不同的 Anchor 框,并且mask proto模块负责输出原型Mask,这些被处理以获取实例分割任务的检测框和分割Mask。
在本文中,作者提出了一种基于单阶段实例分割的细胞图像模型,该模型将注意力 尺度序列融合(Attentional Scale Sequence Fusion)集成到YOLO框架(ASF-YOLO)中。首先使用CSPDarknet53 Backbone 网络在特征提取阶段从细胞图像中提取多维特征信息。作者提出了细胞实例分割部分的新颖网络设计。本工作的贡献如下:
- 针对不同类型细胞的多 Scale 问题和小型细胞的目标检测与分割问题,作者设计了一个 Scale 序列特征融合(Scale Sequence Feature Fusion,SSFF)模块和一个三特征编码器(Triple Feature Encoder,TFE)模块,将Path Aggregation Network(PANet)结构中 Backbone 网络提取的多 Scale 特征图进行融合。
- 接下来,作者设计了一个通道和位置注意力机制(Channel and Position Attention Mechanism,CPAM),将来自SSFF和TFC模块的特征信息进行融合,进一步提高实例分割的准确性。
- 在训练阶段,作者利用EIoU(Elastic Intersection over Union)[14]来优化边界框位置损失,通过最小化边界框和 Anchor 框的宽度和高度差异来实现。在后处理阶段,作者使用软非极大值抑制(Soft Non-Maximum Suppression,Soft-NMS)来改善密集重叠细胞问题。
- 作者将提出的ASF-YOLO模型应用于密集重叠和各种细胞类型的实例分割任务。据作者所知,这是首次利用基于YOLO的模型进行细胞实例分割。对两个基准细胞数据集的评估表明,与其他最先进的相比,检测准确率和速度具有优势。
2 Related work
2.1 Cell instance segmentation
细胞实例分割可以进一步帮助在图像上完成细胞计数任务,而细胞图像的语义分割则不能。深度学习方法提高了自动核分割的准确性。Johnson等人,Jung等人,Fujita等人和Bancher等人提出了基于Mask R-CNN的细胞同时检测和分割的改进方法。Yi等人和Cheng等人利用SSD方法检测和分割神经细胞实例。Mahbod等人采用语义分割算法U-Net模型进行细胞核分割。
带有注意机制的SSD和U-Net与Mask R-CNN的混合模型或U-Net与Mask R-CNN在细胞实例分割数据集上取得了一定的提升。BlendMask是一个核实例分割框架,具有膨胀卷积聚合模块和上下文信息聚合模块。Mask R-CNN是一个两阶段目标分割框架,速度较慢。SSD、U-Net和BlendMask是端到端(即单阶段)框架,但在分割密集和小型细胞方面性能较差。
2.2 Improved YOLO for instance segmentation
最近,YOLO在实例分割任务上的改进主要集中在注意力机制,改进的 Backbone 网络或网络,以及损失函数。SENet块被集成到改进的YOLACT中,用于在显微镜图像中识别反刍颗粒。YOLOMask,PR-YOLO和YOLO-SF增强了YOLOv5和YOLOv7-Tiny,使用了卷积块注意力模块(CBAM)。改进的 Backbone 网络网络中添加了有效的特征提取模块,使YOLO特征提取过程更加高效。YOLO-CORE通过使用设计的多级约束(包括极距离损失和扇区损失)的显式和直接轮廓回归有效地增强了实例分割的Mask。此外,混合模型另一个YOLOMask和YUSEG结合优化了YOLOv4和原始YOLOv5s,并使用语义分割U-Net网络确保了实例分割的准确性。
3 The proposed ASF-YOLO model
3.1 Overall architecture
如图3所示,作者提出了一个结合空间和多尺度特征的ASF-YOLO框架,用于细胞图像实例分割。作者开发了一种新颖的特征融合网络架构,由两个主要组件网络组成,可以提供小目标分割的互补信息:
- SSSF模块,它将来自多个尺度图像的全局或高级语义信息组合在一起;
- TFE模块,它可以捕捉小目标目标的局部精细细节。将局部和全局特征信息相结合可以产生更准确的分割图。
作者首先对 Backbone 网络中提取的P3、P4和P5的输出特征进行融合。首先,设计SSSF模块,有效地将P3、P4和P5捕获的不同空间尺度覆盖各种大小和形状的细胞类型的特征图融合在一起。在SSSF中,将P3、P4和P5特征图归一化到相同大小,上采样,然后堆叠在一起作为输入到3D卷积,以组合多尺度特征。其次,开发TFE模块,以增强密集细胞中小目标检测,通过在空间维度中拼接不同大小的特征(大、中、小)来捕捉小目标的详细信息。将详细特征图归一化到相同大小,上采样,然后堆叠在一起作为输入到3D卷积,以组合多尺度特征。最后,开发TFE模块,以增强密集细胞中小目标检测,通过在空间维度中拼接不同大小的特征(大、中、小)来捕捉小目标的详细信息。
TFE模块的信息随后通过PANet结构集成到每个特征分支中,然后与SSSF模块的多 尺度信息结合到P3分支中。作者进一步在P3分支中引入了通道和位置注意力机制(CPAM),以利用高阶多尺度特征和详细特征。CPAM中的通道和位置注意力机制可以分别捕获有用的通道并改进与细胞等小目标相关的空间局部定位,从而提高其检测和分割精度。
3.2 Scale sequence feature fusion module
对于细胞图像的多尺度问题,现有文献中使用特征金字塔结构进行特征融合,其中只使用求和或 ConCat 来融合金字塔特征。然而,各种特征金字塔网络无法有效地利用所有金字塔特征图之间的相关性。作者提出了一种新颖的尺度序列特征融合,可以更好地将深度特征图的高维信息与浅层特征图的详细信息相结合,其中图像大小在降采样过程中发生变化,但尺度不变的特征不会发生变化。 尺度空间沿图像的尺度轴构建,不仅表示一个尺度,还表示目标可以具有的各种不同尺度范围。 尺度意味着图像的细节。一个模糊的图像可能会丢失细节,但图像的结构特征可以得到保留。作为输入到SSFF的缩放图像可以由以下方法获得:
其中表示具有宽度w和高度h的2D输入图像。通过使用一系列卷积在2D高斯滤波器下进行平滑生成。是用于卷积的2D高斯滤波器的标准偏差参数。
这些生成的图像具有相同的分辨率,但具有不同的尺度。因此,不同大小的特征图可以被视为尺度空间,并且可以调整不同分辨率的有效特征图到相同分辨率以便拼接。受到在多个视频帧上的2D和3D卷积操作的启发,作者将不同尺度的特征图水平堆叠,并使用3D卷积提取它们的尺度序列特征。由于高分辨率特征图级别P3包含对小目标检测和分割至关重要的信息,因此,作者基于P3级别设计SSFF模块。如图3所示,所提出的SSFF模块包括以下组件:
- 使用1*1卷积将P4和P5特征级别的通道数更改为256。
- 最近邻插值方法用于调整它们的大小以适应P3级别的尺寸。
- 为了增加每个特征层的维度,使用unsqueeze方法,将每个特征层从3D张量(高度、宽度、通道)转换为4D张量(深度、高度、宽度、通道)。
- 然后,将4D特征图沿着深度维度进行拼接,形成一个3D特征图,以便后续的卷积。
- 最后,使用3D卷积、3D批标准化和SiLU激活函数(SiLU)来完成尺度序列特征提取。
3.3 Triple feature encoding module
为了识别密集重叠的小目标,可以通过放大图像来参考和比较不同尺度下的形状或外观变化。由于 Backbone 网络的不同特征层具有不同的尺寸,传统的FPN融合机制只将小尺寸的特征图上采样,然后将或添加到前一层特征中,从而忽略了较大尺寸特征层丰富的详细信息。因此,作者提出了TFE模块,将大、中、小特征进行分割,添加大尺寸特征图,并进行特征放大以提高详细特征信息。
图4说明了TFE模块的结构。在特征编码之前,首先调整特征通道数,使其与主要尺度特性一致。然后,将大尺寸特征图(Large)处理由卷积模块,其通道数调整为1,然后使用混合结构(最大池化 + 平均池化)进行下采样,有助于保留高分辨率特征和细胞图像的有效性和多样性。对于小尺寸特征图(Small),卷积模块也用于调整通道数,然后使用最近邻插值方法进行上采样,有助于保持低分辨率图像的丰富局部特征,并防止小目标特征信息的损失。最后,将大、中、小尺寸具有相同尺寸的三个特征图卷积一次,然后按通道进行拼接。
其中表示TFE模块的输出特征图。,和分别表示大尺寸、中尺寸和小尺寸的特征图。由,和的拼接得到。的分辨率与相同,但通道数是的三倍。
3.4 Channel and position attention mechanism
为了提取不同通道中包含的代表性特征信息,作者提出了CPAM,以集成详细特征信息和多尺度特征信息。CPAM的结构如图5所示。它包括一个从TFE(输入1)接收输入的通道注意网络,以及一个从通道注意网络和SSFF(输入2)的输出叠加接收输入的位置注意网络。
输入1为通道注意网络,是经过PANet处理后的特征图,包含TFE的详细特征。SENet通道注意力块首先对每个通道进行独立的全局平均池化,并使用两个全连接层以及一个非线性Sigmoid函数来生成通道权重。这两个全连接层旨在捕捉非线性跨通道交互,这涉及降低维数以控制模型复杂度,但维数减少会导致通道注意力预测出现副作用,捕捉所有通道之间的依赖关系是低效和不必要的。
作者引入了一种不需要降维的注意力机制来有效地捕捉跨通道交互。在不对通道维数进行降低的情况下,通过考虑每个通道及其k个最近邻,即使用大小为k的1D卷积来实现局部跨通道交互的捕捉,其中核大小k表示局部跨通道交互的覆盖范围,即参与一个通道注意力预测的邻居数量。为了获得最佳覆盖范围,可能需要对不同网络结构和不同数量的卷积模块进行k的手动调整,但这非常费时。由于卷积核大小k与通道维数C成正比,因此可以看出k和C之间存在映射关系。通道维数通常为2的指数,映射关系如下。
为了使具有较大通道数的层中能够实现更多的跨通道交互,可以通过一个函数调整一维卷积的卷积核大小。卷积核大小可以通过以下公式计算:
其中表示最近邻的奇偶性。参数设置为2,b设置为1。根据上述非线性映射关系,高值通道的交换时间更长,而低值通道的交换时间更短。因此,通道注意力机制可以对多个通道特征进行更深层次的挖掘。
将通道注意力机制的输出与SSFF(输入2)的特征图作为位置注意力网络的输入,可以为每个细胞提取关键的位置信息。与通道注意力机制不同,位置注意力机制首先将输入特征图按宽度和高度分为两部分,然后分别对这两部分在轴和上进行特征编码,最后将它们合并以生成输出。
更精确地说,输入特征图在水平和垂直轴上进行池化,以保留特征图的空间结构信息,这可以计算如下:
3.5 Anchor box optimization
通过优化损失函数和非极大值抑制(NMS),三个检测Head中的 Anchor 框得到了改进,从而准确地完成了不同大小的细胞图像的实例分割任务。
互交面积(IoU)通常被用作 Anchor 框损失函数,通过计算标注框和预测框之间的重叠程度来确定收敛。然而,经典的IoU损失无法反映目标框和 Anchor 框之间的距离和重叠。为了解决这些问题,GIoU,DIoU和CIoU被提出。CIoU引入了基于DIoU损失的影响因子,YOLOv5和YOLOv8使用该影响因子。虽然考虑了重叠区域和中心点距离对损失函数的影响,但同时也考虑了标注框和预测框的宽高比(即,长宽比)对损失函数的影响。然而,它仅反映了标注框和预测框之间的长宽比差异,而没有反映标注框和预测框的实际长宽关系。EIoU最小化了目标框和 Anchor 框之间的宽高差异,可以提高小目标的定位效果。EIoU损失可以分为三个部分:IoU损失函数,距离损失函数和长宽比损失函数,其公式如下所示。
其中表示欧几里得距离,和分别表示框B和的中心点;, ,和分别表示GT框的中心点、宽和高;和分别表示覆盖两个框的最小外接框的宽和高。与CIoU相比,EIoU不仅加快了预测帧的收敛速度,还提高了回归精度。因此,作者在 Head 部分选择EIoU替换CIoU。
为了消除重复的 Anchor 框,检测模型在同一时间输出多个检测边界,特别是在周围有很多高置信度检测边界的情况下。经典的NMS的原则是获取局部最大值。如果当前框与最高得分检测帧之间的差异大于阈值,则将框的分数直接设置为零。为了克服由经典NMS引起的错误,作者采用Soft-NMS,它使用高斯函数作为权重函数将预测边框的分数减少以替换原始分数,而不是直接将其设置为零,从而修改了框的错误规则。
4 Experiments
4.1 Datasets
作者在两个细胞图像数据集上评估了提出的ASF-YOLO模型的性能:DSB2o18和BCC数据集。DSB2o18数据集包含670个细胞核图像带有分割Mask,旨在评估算法在细胞类型、放大倍数和成像模式(平面光 vs. 荧光)变化下的泛化性。每个Mask包含一个细胞核,不同Mask之间没有重叠(没有像素属于两个Mask)。该数据集按8:2的比例将训练集和测试集划分。训练集和测试集的样本量分别为536和134张图像。
乳腺癌细胞(BCC)数据集来自加州大学圣塔芭芭拉分校(UCSB CBI)的生物图像信息中心。该数据集包含160张使用伊红-复染色法(hematoxylin and eosin,H&E)染色的病理图像,这些图像用于乳腺癌细胞检测,并附有相关 GT 数据。该数据集按80%为训练集,20%为测试集的比例随机划分。
4.2 Implementation details
实验在NVIDIA GeForce 3090(24G)GPU上实现,使用Pytorch 1.10,Python 3.7和CUDA 11.3依赖项。作者采用了预训练的COCO数据集的初始权重。输入图像大小为640*640。训练数据的批量大小为16。训练过程持续100个周期。作者使用随机梯度下降(SGD)作为优化函数训练模型。SGD的超参数设置为动量0.9,初始学习率0.001和权重衰减0.0005。
4.3 Quantitative results
表1展示了在DSB2018数据集上,提出的ASF-YOLO与其他经典和最先进的方法(包括Mask R-CNN,Cascade Mask R-CNN,SOLO,SOLOv2,YOLACT,Mask R-CNN with Swin Transformer backbone(Mask RCNN Swin T),YOLOv5l-seg v7.o和YOLOv8l-seg)之间的性能比较。
作者的模型在46.18M参数下达到了最佳性能,Box 为0.91,Mask 为0.887,推理速度达到了47.3帧/秒,这是最佳性能。由于图像输入尺寸为8001200,使用Swin Transformer背心的Mask R-CNN的准确性和速度并不高。作者的模型也超过了经典的单阶段算法SOLO和YOLACT。
作者提出的模型在BCC数据集上也实现了最佳的实例分割性能,如表2所示。实验验证了ASF-YOLO在不同数据集和不同细胞类型下的泛化能力。
4.4 Qualitative results
图6提供了不同方法在DSB2o18数据集样本图像上细胞分割的视觉比较。通过使用TFC模块提高小目标检测性能,ASF-YOLO在单个通道下对密集且小目标的细胞图像具有很好的召回值。通过使用SSSF模块增强多尺度特征提取性能,ASF-YOLO在复杂背景下的大尺寸细胞图像上也具有很好的分割精度。这表明作者的方法对不同细胞类型具有很好的泛化能力。
从图6(a)和(b)可以看出,每个模型都有很好的结果,因为细胞图像相对简单。从图6(c)和(d)可以看出,Mask R-CNN由于两阶段算法的设计原理,具有较高的假阳性检测率。SOLO有许多漏检,YOLOvgl-seg无法分割模糊边界的细胞。
4.5 Ablation study
作者对提出的ASF-YOLO模型进行了系列的广泛消融研究。
4.5.1 Effect of the proposed methods
表3显示了每个建议模块在提高分割性能方面的贡献。YOLOv5l-seg中使用Soft-NMS可以克服密集小目标检测时的错误抑制问题,并提供性能改进。EIoU损失函数改善了小物体边界框的效果,提高了 by 1.8%。SsFF,TFC和CPAM模块有效改善了模型性能,解决了细胞图像中小物体实例分割的问题。
4.5.2 Effect of attention mechanisms
与channel attention SENet,channel and spatial attention CBAM和spatial attention CA相比,提出的CPAM注意力机制尽管计算量和参数略有增加,但提供了更好的性能。
图7展示了使用ASF-YOLO模型中不同注意力模块进行分割的结果可视化。提出的CPAM具有更好的通道和位置特征信息,并从原始图像中挖掘出更丰富的特征。
4.5.3 Effect of convolution module in the backbone
表5显示,在提出的模型中,将YOLOv5的C3模块替换为YOLOv8的C2f模块,导致模型在两个数据集上的C2f模块性能降低。
5 Conclusion
作者开发了一种准确快速的细胞图像实例分割模型ASF-YOLO,用于细胞图像分析,该模型将空间和尺度特征融合用于细胞图像的检测和分割。作者在YOLO框架中引入了几个新颖的模块。SSFF和TFE模块增强了多尺度和小物体实例分割性能。通道和位置注意力机制进一步挖掘了两个模块的特征信息。
大量实验结果表明,作者提出的模型能够处理各种细胞图像的实例分割任务,并显著提高了原始YOLO模型在细胞分割方面的准确性,由于小和密集物体的存在。作者的方法在细胞实例分割的准确性和推理速度方面都显著优于最先进的方法。由于本文中的数据集较小,模型的泛化性能需要进一步提高。此外,在ASF-YOLO中每个模块的有效性在消融研究中进行了讨论,这为进一步改进提供了研究依据。