ASF-YOLO:一种基于注意力尺度序列融合的细胞实例分割YOLO模型

摘要:

我们提出了一种基于注意力尺度序列融合的ASF-YOLO框架,该框架融合了空间和尺度特征,来实现快速准确的细胞实例分割。建立在YOLO分割框架上面,采用多尺度序列特征融合模块(SSFF)模块来增强网络的多尺度信息提取能力,采用三重特征编码器(TPE)模块来融合不同尺度的特征映射,增加详细信息。我们引入了通道和位置注意力机制(CPAM)来继承SSFF模块和TPE模块,它关注于信息通道和空间位置相关的小物体,以此来提高检测和分割性能。在两个细胞数据集上的实验表现了提出的ASF-YOLO模型的卓越的分割准确度和速度。在2018年的Data Science Bowl 数据集上面实现了0.91的box mAP、0.887的mask mAP和47.3的FPS,性能优于现在最先进的模型。源代码开源:https://github.com/mkang315/ASF-YOLO

关键词:医学图像分割、小目标分割、YOLO、系列特征融合、注意力机制。

1 引言

     随着样品制备技术和显微成像技术的快速发展,细胞图像的定量处理和分析在医学、细胞生物学等领域发挥着重要作用。基于卷积神经网络(CNN),通过神经网络训练可以学习到不同细胞图像的特征信息,具有较强的泛化性能。两阶段R-CNN系列及其单阶段变体是用于实例分割任务的经典基于cnn的框架。在最近的研究中,You Only Look Once (YOLO)系列已经成为实时实例分割中最快、最准确的模型之一。由于采用了一阶段的设计思想和特征提取能力,YOLO实例分割模型比两阶段分割模型具有更高的精度和速度。但是,单元格实例分割的难点在于对象小、密集、重叠,以及单元格边界模糊,导致单元格精度较差。细胞实例分割需要对细胞图像中不同类型的物体进行精确的细节分割。如图1所示,由于细胞形态、制备方法和成像技术的差异,不同类型的细胞图像在颜色、形态形态学、纹理等特征信息上存在较大差异。

图1 不同的细胞(左侧)和他们的特征图(右侧)

  典型的YOLO框架由三个主要部分组成:主干、颈部和头部。YOLO的骨干网络是一个提取不同细粒度图像特征的卷积神经网络。CSPDark-net53就是在YOLOv4的基础上面修改过来的,设计为包括C3模块和BNSiLU的YOLOv5的主干网络。在YOLOv8中,C3模块被C2f所替代。这是和YOLOv5唯一不同的地方。如图2所示,YOLOv5和YOLOv8主干中的1-5级特征提取分支{P1,P2,P3,P4,P5}对应和这些特征映射相关联的YOLO网络输出。YOLOv5 v7 和YOLOv8是第一个基于YOLO结构的框架,除了目标检测和分类之外,还可以处理分割的任务。在v5特征提取阶段,使用多个C3模块堆叠的CSPDarkNet53骨干网,然后将骨干网的三个有效特征分支P3、P4和P5作为特征金字塔(FPN)的输入,在Neck构建多尺度融合结构。在特征层的解码过程中,使用骨干网有效特征分支对应的三个不同大小的头部,对目标进行边界框预测,对P3特征进行上采样后,逐像素解码作为目标的分割掩码预测,完成目标的实例分割。在分割头中,三个尺度的特征输出三个不同的锚框,掩码原型模块负责输出原型掩码,对原型掩码进行处理,得到用于实例分割任务的检测框和分割掩码。

  本文提出一种将注意力尺度序列融入到YOLO框架中的细胞单阶段实例分割模型。在特征提取阶段,首先利用CSPDarknet53骨干网络从细胞图像中提取多维特征信息。在Neck部分设计新的网络结构,本文的工作贡献总结如下:

(1)针对不同类型细胞的多尺度问题和小细胞的目标检测和分割问题,设计了尺度序列融合模块(SSFF)和三重特征编码器(TFE),将路径聚合网络(PANet)结构中提取的主干网络多尺度特征映射进行融合。

(2)设计通道和位置注意机制(CPAM),整合SSFF和TFC模块的特征信息,进一步提高实例分割精度。

(3)我们在训练阶段使用EIoU进行优化,通过最小化边界框和锚框的宽度和高度之间的差异来减少边界框位置损失,在后处理阶段使用最大抑制(Soft-max)来改善密集重叠的细胞问题。

(4)提出ASF-YOLO模型来应对密集挑战和各种各样细胞类型的实例分割任务。就我们所知道,这是第一个基于YOLO模型用于细胞实例分割的工作。与其他最先进的方法进行比较,在两个细胞数据集上的评估都表现出更好的检测准确性和速度。

2 相关工工作

2.1 细胞实例分割

  细胞实例分割可以更进一步帮助在图像中的细胞计数工作但是语义分割不可以。深度学习提高了自动核分割的准确性。Johnson et al. , Jung et al. , Fujita et al.和Bancher等人基于Mask R-CNN提出了改进的细胞同步检测和分割方法。。Yi et al.[21] 和Cheng et al. 利用 Single-Shot multi-box Detector (SSD)方法检测和分割神经细胞实例。Mahbod等人采用基于U-Net模型的语义分割算法进行细胞核分割。混合模型SSD和U-Net具有注意机制[19]或U-Net和Mask R-CNN在单元实例分割数据集上取得了一定的提升。BlendMask[27]是一个包含展开卷积聚合模块和上下文信息聚合模块的核实例分割框架。Mask R-CNN是一种速度较慢的两阶段目标分割框架。SSD、U-Net和BlendMask是统一的端到端(即单阶段)框架,但在密集和小单元的分割方面性能较差。

2.2 改进实例分割的YOLO

  最近YOLO在实例分割任务中的改进主要集中在注意力机制改进骨干网络和损失函数上面。将SENet块集成到改进的YOLACT[6]中,用于识别显微图像中的瘤胃原生动物。YOLOMask, PR-YOLO和YOLO-SF用卷积块注意模块CBAM增强YOLOv5和YOLOv7-Tiny。改进后的骨干网中加入了有效的特征提取模块,使得YOLO特征提取过程更加高效。YOLO-CORE利用设计的由极距离损失和扇区损失组成的多阶约束,通过显式和直接的轮廓回归,有效地增强了实例的掩码。此外,混合模型另一个YOLOMask和YUSEG结合了优化的YOLOv4和原始的YOLOv5s与语义分割U-Net网络,以确保实例分割的准确性。

3 提出的ASF-YOLO框架

3.1 整体框架

  图3展示了所提出的ASF-YOLO框架的概述,该框架结合了空间和多尺度特征,用于细胞图像实例分割。我们开发了一种新的特征融合网络框架,该架构由两个可以提供互补信息的主要组件网络构成

对于小目标分割:(1)SSFF模块,该模块结合了图像的多个尺度的全局或高级语义信息;(2)TFE模块,该模块可以捕获小目标对象的局部精细细节。局部和全局特征信息的融合可以生成更精确的分割图。我们从骨干网络中提取的P3、P4和P5的输出特征先进行融合。首先,设计了SSFF模块,有效融合P3、P4和P5的特征图,这些特征图捕获了覆盖不同细胞类型的各种大小和形状的不同空间尺度。在SSFF中P3、P4和P5特征映射被归一化到相同尺寸,上采样。然后堆叠在一起作为3D卷积的输入,来组合多尺度特征。其次,开发了TFE模块,通过在空间维度上拼接大、中、小三种不同尺寸的特征来捕获小目标的详细信息,增强对密集细胞的小目标检测。详细特征然后通过PANet结构将TFE模块的信息整合到各个特征分支中,再与SSFF模块的多尺度信息结合到P3分支中。我们进一步在P3分支中引入信道和位置注意机制(CPAM),以利用高级多尺度特征和详细特征。CPAM中的通道注意机制和位置注意机制可以分别捕获信息通道和细化与细胞等小物体相关的空间定位,从而提高其检测和分割精度。

3.2 尺度序列特征融合模块

    对于细胞图像的多尺度问题,现有文献采用特征金字塔结构进行特征融合,仅采用求和或拼接方法融合金字塔特征。然而,各种特征金字塔网络的结构并不能有效地利用所有金字塔特征映射之间
的相关性。我们提出了一种新的尺度序列特征融合,可以更好地将深度特征图的高维信息与浅层特征图的详细信息结合起来,在图像的下采样过程中,图像的大小会发生变化,但尺度不变的特征不会发生变化。比例尺空间沿着图像的比例尺轴线构建,它不仅代表一个比例尺,还代表一个物体可以具有的各种比例尺的范围。比例尺表示图像的细节。模糊的图像可能会失去细节,但图像的结构特征可以保留下来。作为SSFF输入的缩放后的图像可以通过

其中,f(w,h)表示宽度为w,高度为h的二维输入图像经过二维高斯滤波器G(w,h)一系列卷积之后进行平滑后生成的,其中,\sigma是用于卷积的二维高斯滤波器的标准差缩放参数。

这些生成图像具有相同的分辨率,但是,尺度不同。因此,可以将不同尺寸的特征图视为尺度空间,将不同分辨率的有效特征图调整到相同的分辨率进行拼接。受多视频帧上二维和三维卷积操作
的启发,我们将不同尺度的特征图水平叠加,并利用三维卷积提取其尺度序列特征。高分辨率级别P3包含了对小目标至关重要的大部分信息,基于P3级别设计了SSFF模块。如图3所示,所提出的SSFF模块由以下几个部分构成:

(1)使用1×1卷积将P4和P5特征图的通道调整为256。

(2)使用最近邻插值方式将他们的大小调整为P3级别的大小。

(3)unsqueeze方法用于增加每个特征层的维度,将其从3D张量[高度、宽度、通道]变为4D张量[深度、高度、宽度、通道]。

(4)然后沿着深度维度将4D特征图连接起来,形成一个3D特征图,用于后续的卷积。

(5)最后利用3D卷积、3D批归一化和SiLU[43]激活函数完成尺度序
列特征提取。

3.3 三重特征编码器

   为了识别密集重叠的小物体,人们可以通过放大图像来参考和比较不同尺度下的形状或外观变化。由于骨干网的不同特征层具有不同的尺寸,传统的FPN融合机制只是对小尺寸的特征图进行上采样,然后将其拆分或添加到前一层的特征中,而忽略了大尺寸特征层丰富的详细信息。因此,我们提出了TFE模块,将大、中、小特征拆分,增加大尺寸特征图,并进行特征放大,以提高详细的特征信息。

图4展示了TFE模块的结构,在特征编码之前,首先对特征的通道数量进行了调整,使其与主尺度特征保持相一致。大尺度特征图(Large)经过卷积模块处理之后,将其通道数调整至1C,然后使用最大池化+平均池化的混合结构进行下采样,有助于保留高分辨率特征和细胞图像的有效性和多样性。对于小尺寸特征图(Small)使用卷积来调整其通道数,然后使用最近邻插值法进行上采样。这有助于保持低分辨率图像局部特征的丰富性,防止小目标特征信息的丢失。最后,将具有相同维度的大、中、小三个尺寸的特征图进行一次卷积,然后在通道维度上进行拼接,如下图所示。

 可以看到,输出通道是原来的3倍。

3.4 频道和位置关注机制

为了提取不同通道中包含的代表性特征信息,我们提出了将详细特征信息和多尺度特征信息相结合的CPAM方法。CPAM的体系结构如图5所示。它包括从TFE接收输入的通道注意网络和从通道注意网络和SSFF输出叠加接收输入的位置注意网络。

通道注意网络的输入1是PANet之后的特征图,其中包含了TFE的详细特征。SENet频道注意块首先对每个频道采用全局平均池化

 并使用两个完全连接的层和一个非线性Sigmoid函数来生成信道权值。这两个完全连接的层旨在捕获非线性跨通道相互作用,这涉及到降维来控制模型的复杂性,但降维会给通道注意力预测带来副作用,并且捕获所有通道之间的依赖关系是低效和不必要的。我们引入了一种没有降维的注意机制,以有效的方式捕获跨通道的相互作用。在不降低维数的通道全局平均池化之后,通过考虑每个通道及其k个最近邻居来捕获局部跨通道相互作用,即使用大小为k的1D卷积实现,其中核大小k表示局部跨通道相互作用的覆盖范围,即:即有多少邻居参与到一个通道的注意力预测中。为了获得最优覆盖,人们可能会在不同的网络结构和不同数量的卷积模块中手动调整k,然而这是繁琐的。由于卷积核大小k与通道维数C成正比,可见k与C之间存在映射关系,通道维数一般为2的指数,其映射关系如下。

 为了在通道数较大的层中实现更多的跨通道相互作用,可以通过一个函数来调整一维卷积的卷积核大小。卷积核大小k可以计算为:

      其中|·奇数|表示最近邻居的奇数。将γ的值设为2,b设为1。根据上述非线性映射关系,高值通道的交换时间较长,低值通道的交换时间较短。因此,通道注意机制可以对多通道特征进行更深层次
的挖掘。
    将通道注意机制的输出与SSFF(输入2)的特征相结合,作为位置注意网络的输入,为从每个细胞中提取关键的位置信息提供了补充信息。与通道注意机制相比,位置注意机制首先根据其宽度
和高度将输入特征图分成两部分,然后分别在(pw和ph)轴上进行特征编码处理,最后合并生成出。
      更准确地说,输入特征图在水平轴(pw)和垂直轴(ph化,以保留特征图的空间结构信息,其计算方法如下:

 其中w和h分别为输入特征图的宽度和高度。E(w, j)和E(i, h)是输入特征图中位置(i, j)的值。

 在生成位置注意坐标时,对水平轴和垂直轴应用连接和卷积操作:

 其中P(aw,ah)表示位置注意坐标的输出,Conv表示1 ×1卷积,Concat表示拼接。

在拆分注意力特征时,生成和位置相关特征图对,如下图所示:

 

 其中sw和sh分别是拆分输出的宽度和高度。

CPAM的最终输出定义为:

其中E表示通道和位置关注的权重。

3.5 锚框优化

 通过对损失函数和非最大抑制(Non-Maximum Suppression, NMS)的优化,对三种检测头的锚盒进行了改进,以准确完成不同大小的细胞图像的实例分割任务。

通常使用IoU作为锚盒损失函数,通过计算标记的边界框与预测框之间的重叠程度来确定收敛性。然而,经典的IoU损失不能反映目标盒与锚盒之间的距离和重叠。为了解决这些问题,已经提出了GIoU、DIoU和CIoU。CIoU引入了一个基于DIoU Loss的影响因子,YOLOv5和YOLOv8采用了该影响因子。在考虑重叠面积和中心点距离对损失函数的影响的同时,也考虑了标注框和预测框的宽高比(即长宽比)对损失函数的影响。然而,它只反映了纵横比的差异,而不是标记框的宽度和高度与预测框之间的真实关系。EIoU[14]使目标盒与锚盒之间的宽度和高度差最小化,可以提高小物体的定位效果。EIoU损失可分为IoU损失函数LIoU、距离损失函数Ldis和方面损失函数Lasp3部分,其公式如下:

 式中ρ(·)=∥b−bgt 表示欧几里得距离,b和bgt分别表示b和Bgt=∥2点;bgt、wgt、hgt为中心点

b点的宽度和地面的高度是真理;wc和hc表示覆盖两个箱子的最小围框的宽度和高度。与CIoU相比,EIoU不仅加快了预测框架的收敛速度,而且提高了回归精度。因此,我们选择EIoU代替头部的CIoU。

为了消除重复的锚盒,检测模型会同时输出多个检测边界,特别是当真实物体周围有很多高置信度的检测边界时。经典NMS算法的原理是获取局部最大值。如果当前边界框与得分最高的检测帧之间的差值大于阈值,则直接将边界框的得分设置为零。为了克服经典NMS算法带来的误差,我们采用了Soft-NMS算法,该算法采用高斯函数作为权函数,降低预测边界的分数来代替原始分数,而不是直接将其设置为零,从而修改了误差边界框的规则。

4 实验

4.1 数据集

 我们在两个细胞图像数据集:DSB2018和BCC数据集上评估了所提出的ASF-YOLO模型的性能。2018年数据科学碗(DSB2018)数据集包含670个带有分割掩模的细胞核图像,旨在评估算法在细胞类型、放大倍率和成像方式(亮场与荧光)变化中的通化性。每个掩模包含一个细胞核,掩模之间没有重叠(没有像素属于两个掩模)。数据集按照8:2的比例随机分为训练集和测试集。训练集和测试集的样本量分别为536张和134张图片。

 乳腺癌细胞(BCC)数据集[48]来自加州大学圣巴巴拉分校(UCSB CBI)生物图像信息学中心。它包括160张用于乳腺癌细胞检测的苏木精和伊红染色的组织病理学图像,以及相关的地面真值数据。该数据集被随机划分为128张图像(80%)作为训练集,32张图像(20%)作为测试集。

4.2 实现细节

 实验是在 NVIDIA GeForce 3090 (24G) GPU 和 Pytorch 1.10、Python 3.7和CUDA 11.3依赖项上实现的。我们采用了预训练的COCO数据集的初始权值。输入图像大小为640×640。训练数据量的批处理大小为16。训练过程持续100个epoch。我们使用随机梯度下降(SGD)作为优化函数来训练模型。SDG的超参数设置为动量的0.9,初始学习率的0.001,权衰减的0.0005。

4.3 定量结果

 表1显示了在DSB2018数据集上,所提出的ASF-YOLO与其他经典和最先进的方法(包括Mask R-CNN、Cascade Mask R-CNN、SOLO、SOLOv2、YOLACT、Mask R-CNN与SwinTransformer 主 干 (Mask RCNN Swin T)、YOLOv5l-seg v7.0和YOLOv8l-seg)的性能比较。

 我们的模型在46.18M参数下获得了最佳精度,Box mAP50为0.91,Mask mAP50为0.887,推理速度达到47.3帧/秒(FPS),是最佳性能。由于图像输入尺寸为800×1200,使用Swin Transformer骨干的MaskR-CNN的精度和速度不高。我们的模型也超越了经典的单阶段算法SOLO和YOLACT。

 我们提出的模型在BCC数据集上也取得了最佳的实例分割性能,如表2所示。实验验证了ASF-YOLO对不同细胞类型数据集的泛化能力。

4.4 定性结果

  图6提供了不同方法对DSB2018数据集样本图像的细胞分割的可视化比较。通过使用TFC模块提高小目标检测性能,ASF-YOLO在单通道内对具有密集小目标的单元图像具有良好的召回值。通过使用SSFF模块增强多尺度特征提取性能,ASF-YOLO对于复杂背景下的大尺寸细胞图像也有很好的分割精度。这表明我们的方法对不同的细胞类型具有良好的泛化性。从图6(a)和(b)来看,每个模型都有很好的结果,因为细胞图像相对简单。从图6(c)和(d)可以看出,由于采用了两阶段算法的设计原则,Mask R-CNN具有较高的误检率。SOLO有很多漏检,YOLOv5l-seg无法分割边界模糊的细胞。

 4.5 消融实验

我们对提出的ASF-YOLO模型进行了一系列广泛的消融研究。

 4.5.1 所提方法的有效性

表3显示了每个提出的模块在提高分割性能方面的贡献。在YOLOv5l-seg中使用Soft-NMS可以克服在检测密集小物体的细胞时由于相互遮挡而导致的错误抑制问题,并提供性能提升。EIoU损失函数提高了小对象包围盒的效果,将mAP50:95提高了1.8%。SSFF、TFC和CPAM模块通过解决单元图像的小目标实例分割,有效地提高了模型的性能。

4.5.2 注意力机制的影响

与通道注意SENet、通道和空间注意CBAM以及空间注意CA相比,本文提出的CPAM注意机制在计算量和参数上略有增加,但具有更好的性能。

图7为ASF-YOLO模型中使用不同注意模块的分割结果可视化。所提出的CPAM具有更好的通道和位置特征信息,从原始图像中挖掘出更丰富的特征。

 4.5.3. 骨干网中卷积模块的效果

从表5可以看出,当我们提出的模型的主干中,将YOLOv5的C3模块替换为YOLOv8的C2f模块时,主干中C2f模块在两个数据集上的性能都有所下降。

 5 结论

本文提出了一种用于细胞图像分析的精确、快速的实例分割模型ASF-YOLO,该模型融合了空间和尺度特征,用于细胞图像的检测和分割。我们在YOLO框架中引入了几个新的模块。SSFF和TFE模块增强了多尺度和小对象实例分割性能。通道和位置注意机制进一步挖掘了两个模块的特征信息。大量的实验结果表明,我们提出的模型能够处理各种细胞图像的实例分割任务,并且由于小而密集的目标,大大提高了原始YOLO模型在细胞分割方面的准确性。我们的方法在细胞实例分割的准确性和推理速度方面都大大优于最先进的方法。由于本文中数据集的规模较小,模型的泛化性能还需要进一步提高。此外,在烧蚀研究中讨论了ASF-YOLO各模块的有效性,为进一步改进提供了研究依据。

 

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值