点击下方卡片,关注 「AI视界引擎」 公众号
( 添加时备注:方向+学校/公司+昵称/姓名 )
核分割在组织病理学图像分析中尽管扮演着基础性角色,但仍然是一项具有挑战性的工作。这项任务的主要挑战在于存在重叠区域,这使得分离独立的核更加复杂。
在本文中,作者提出了一种新的双分支架构,通过结合Unet和TransUnet网络来进行核分割任务。
所提出的架构名为Trans2Unet,输入图像首先被送入Unet分支,该分支去除了最后的卷积层。这个分支使网络能够结合输入图像不同空间区域的特征,并更精确地定位感兴趣区域。输入图像同时也被送入第二个分支。在第二个分支,即TransUnet分支中,输入图像将被划分为图像块。
借助架构中的视觉 Transformer (VIT),TransUnet可以作为一个强大的编码器用于医学图像分割任务,并通过恢复局部空间信息来增强图像细节。为了提升Trans2Unet的效率和性能,作者提出将TransUnet与一种计算效率更高的变体相结合,这种变体被称为带跳跃连接的“瀑布”式带孔空间池化(WASP-KC)模块,其灵感来源于“瀑布”式带孔空间池化(WASP)模块。
在2018年数据科学碗基准测试上的实验结果表明,与之前的分割模型相比,所提出架构的有效性和性能。
I Introduction
细胞核分割一直是因其在癌症诊断中的实际应用而受到关注的关键问题。一般来说,这项任务与自然图像分割相似,涉及从细胞核图像(2D或3D图像)中提取所需目标的过程,可以通过手动、半自动或全自动的方式完成[5][6][7]。最近,许多高精度的深度学习模型被用于细胞核分割[8]。2015年,Unet网络因其编码器-解码器架构结合跳跃连接以保留重要特征,在分割任务中,特别是医学图像上表现出色。
尽管Unet和其他CNN网络架构强大,但它们通常仍限于复制由卷积操作的内在局部性所导致的直接的远程相互关系。与基于CNN的网络不同,基于Transformer的模型具有全局计算特征。在[2]中,TransUnet被提出用以解决这一问题,通过采用混合CNN-Transformer方法来增强CNN特征图中的精细高分辨率空间信息以及由Transformer编码的全局上下文。尽管Transformer因其全局特征在计算机视觉中受到欢迎,但缺乏低层细节使得局部特征信息提取不足[10]。
为了充分利用Unet和TransUnet的优势,在本研究中,作者提出将这两种架构结合在一起,以获得更强大的架构。所提出的架构命名为Trans2Unet,包括两个主要分支。一个分支将输入图像通过Unet网络,另一个分支将输入图像通过TransUnet网络。最后,这些分支的输出被连接起来,以重新创建输入图像的特征图,从而提高模型的预测能力。此外,作者没有使用原始的TransUnet架构,而是加入了WASP-KC模块,以利用级联方法中更大视野(FOV)块的逐步提取。
作者的主要贡献可以简述如下:
- 引入了一种新的、更健壮、高效的架构,使用Unet和TransUnet网络。
- 在TransUnet模型的CNN块后加入了WASP-KC块。
通过在2018年数据科学碗挑战数据集上的实际实验,结果表明,与同一数据集上的其他SOTA架构相比,所提出的网络取得了相当好的准确性。具体来说,作者获得了2个参数DSC和IoU,其值分别为0.9225和0.8613。
本文的组织结构如下:首先,在第二节描述了相关工作。第三节介绍了作者提出的模型。在第四节中获取了2018年数据科学碗挑战数据集上的实验结果。最后,在第五节中进行了总结、局限性及进一步工作的描述。
II Related Work
(Unet)
Unet最初在2015年被提出,被认为是一种有效的用于生物医学图像分割的卷积网络。Unet架构包含两条路径:编码器和解码器。编码器路径是下采样部分,每个模块包含一个整流线性单元(ReLU)和一个步长为2的2x2最大池化操作[9]。解码器路径是用于重建图像的高分辨率特征图的上采样部分。特别是,Unet使用跳跃连接来保留空间信息,因为在编码器路径的下采样过程中,输入图像的空间信息会丢失,导致架构准确度下降。
(ViT)
在自然语言处理任务(NLP)方面,人们已经知道Transformer架构是一项关键标准。然而,在计算机视觉任务中,这个模型仍然存在许多限制[3]。视觉Transformer(ViT)是一个开创性的模型,它通过将输入图像嵌入为一串视觉标记,并使用一组Transformer块来建模这些序列之间的全局依赖关系,将Transformer模型适配到计算机视觉(CV)任务。ViT简单地将输入图像视为1D序列,这导致了在建模局部视觉结构时缺乏归纳偏置[11]。最近,ViT在多种应用中取得了极具竞争力的准确度基准:图像分类、目标检测和语义图像分割。受到将输入图像作为序列处理的启发,视觉Transformer是Transformer架构部分和多层感知机(MLP)块的组合[1]。ViT的Transformer编码器包括多头自注意力层(MHSA)、多层感知机(MLP)层和层归一化(LN)[12]。MHSA是Transformer块的关键组成部分。它是经过n次重复单头自注意力(SHSA)后实现的,其中n是头的数量。MHSA旨在从图像中重现长距离的结构数据[13]。
(TransUnet)
TransUnet可以被看作是Unet的升级版本。TransUnet是首个将 Transformer (transformers)应用于计算机视觉任务的网络架构,并且它通过在图像任务上成功应用 Transformer ,开辟了新的研究方向。TransUnet与Unet之间的主要区别在于编码器路径(Encoder Path)。图2中对TransUnet编码器路径架构有相当详细的描述。它包括卷积神经网络块(在研究[2]中,作者使用ResNet50作为 Backbone 网络)和视觉 Transformer (ViT)。应用了ViT中的 Transformer 的编码器包括连续的多头自注意力(MHSA)层和MPL块。与使用批量归一化(BatchNorm,BN)不同, Transformer 块在每个块之前使用层归一化(LayerNorm,LN),并在每个块之后添加残差连接[16][17]。
DeepLabv3+
在研究[14]中,提出了将Atrous Spatial Pyramid Pooling模块(ASPP)与编码器-解码器结构相结合的方法,该研究显示对输入图像中分割目标的边界有更好的改进。ASPP的特殊结构将具有不同 Level 的四个并行分支中的扩张卷积组合在一起。最终,通过快速双线性插值与额外的八倍因子相结合,所得到的特征图被恢复到原始分辨率[3]。DeepLabv3+在准确性方面显著优于前一个版本。
Waterfall Atrous Spatial Pooling (WASP)
WASP是一种高效率的语义分割架构。它利用级联结构中的逐步过滤,同时与空间金字塔配置相比,保持多尺度视场(FOV)。根据[3]的研究,当WASP与Resnet Backbone 网络结合使用时,将为分割问题提供一个健壮的架构并取得潜在的结果。此外,这种变体在计算上有效,是DeepLabv3+架构中的一种带孔空间池化(ASP)类变体。[15]展示了WASP模块在训练过程中的计算时间大幅提升,并且相比于原始的ASPP模块,参数数量有所减少。
III Methodology
Waterfall Atrous Spatial Pooling with Skip Connection (WASP-KC) Module
WASP-KC模块如图1所示,灵感来源于WASP模块。WASP-KC包含四个大型视场(FOV)单元,它们合并在一起形成瀑布状以产生输出。除了多尺度方法[26][23]外,该模块还受到级联配置[3][14]的启发,同时也受到了ASPP[24]和Res2Net模块[25]的并行结构的启发。WASP模块有助于减少所需的参数和内存,从而降低了计算的昂贵程度,这是Atrous卷积[3][15]的主要局限。根据作者在[3]中进行的实验,WASP模块成功减少了20.69%的参数,并且在使用基于此模块构建的WASPnet网络上,与Res2Net-Seg或ASPP模块相比,模型的性能提高了2%(mIoU)。在这项研究中,作者通过密集连接替换了WASP块,这些连接受到了DenseNet模型的启发。在这种技术中,每一层都将所有前一层输出作为输入,其特征图将被传递到更深的层,这意味着每一层都接收前面所有层的完整信息。这将确保特征的可重用性,因为之前层的特征图被保留并全部相加,这有助于输入图像数据在没有任何损失的情况下得到良好保持。这是一个重要的修改,使得WASP可以更健壮地工作。WASP-KC块直接添加在CNN模块(使用ResNet-50 Backbone 网)之后,以提高所提模型的性能和效率。
Model architecture
针对开发用于细胞核图像分割的新深度学习架构,本研究提出了Trans2Unet,该架构结合了Unet和TransUnet分支。首先,为了提高TransUnet分支的效率,作者采用了如图2所示的附加WASP-KC块。WASP-KC块由四个卷积单元组成。每个单元包含三个块,第一个块使用3x3卷积,随后两个块应用1x1卷积。3x3卷积块在水平方向上共享信息,通过这种方式,信息将被用于模块的所有单元中。此外,每个单元还使用了跳跃连接,以利用前一层特征。这一调整显著提高了相较于WASP模块的性能。模块的输出是这四个单元和全局平均池化块的输出之和,并且也将作为ViT网络的输入。
图3展示了所提出Trans2Unet的总体结构,包括Unet分支和 Proposal 的TransUnet+WASP-KC分支。输入图像经过这两个分支后,两个分支的输出将被连接在一起。最终,在上面的两个分支聚合输出之后,作者继续通过一个卷积块进行处理,然后产生预测输出。这是一种相当新颖且简单的组合,但其性能提升远胜于仅使用常规的Unet或TransUnet。
Loss function
损失函数,也称为成本函数,是一个表示q(模型的预测结果)与p(实际值)之间关系的方程。作者的任务是使该方程的值最小化。损失函数用于优化模型,这也是评估模型质量的一个参数。与图像分割相关的任务应用了许多损失函数,例如二元交叉熵(BCE)、Dice损失等。
二元交叉熵(BCE)损失函数 计算两个概率分布之间的差异,它们是实际概率分布p和预测概率分布q。它通常用于目标分类任务,在图像分割任务中也是如此,因为它是对像素的分类。这应该用于平衡的数据集。BCE损失由以下方程表示:
其中p代表 GT 标签,q代表Trans2Unet模型的预测值。(1)式的值反映了实际值与模型预测值之间的差异。
Dice损失 是在涉及弧形图像分割或医学图像分割任务中广泛使用的损失函数。该损失函数的值衡量 GT 与预测值之间的差异。Dice损失由以下方程表示:
数学符号的含义与二元交叉熵部分相似。
Evaluation Metrics
当前,Dice相似性系数(DSC)和Jaccard指数或交并比(IoU)是评估医学图像分割模型中最受欢迎的指标[18][19][20]。在本研究中,作者也采用这两个参数,以便在2018年数据科学碗挑战赛数据集上与其他模型进行公平的比较。
DSC和IoU的定义如下数学表达式[21]:
其中:TP、FP、FN、TN分别代表真正例、假正例、假反例和预测负例的数量。此外,
在当前研究中,Dice相似性系数(DSC)和Jaccard指数或交并比(IoU)是医学图像分割模型评估中广泛采用的指标[18][19][20]。在本项研究中,作者也采用这两个指标,以便在2018年数据科学碗挑战赛的数据集上与其他模型进行公正的比较。
DSC和IoU的计算公式如下[21]:
图1展示了本研究中使用的瀑布式扩张空间池化(WASP-KC)模块[22]。对于图像分割任务,还存在其他评估指标,例如精确度、准确率、体积相似性等。
其中,TP、FP、FN、TN分别代表真阳性、假阳性、假阴性和真阴性的数量。另外,
IV Experimental Results
作者在此部分展示了实验的结果。通过对数据集的分析,作者评估了所提出方法的有效性和性能。所有实验均按照第III节中描述的实验设计进行。以下是详细的结果和讨论。
Dataset
为了准确评估Trans2Unet模型的性能,作者使用了公共的生物医学图像数据集——2018年数据科学碗挑战数据集和GlaS数据集。2018年数据科学碗挑战数据集包含了原始图像及其 Mask (或称 GT 值)。总共有670幅图像,作者将这个数据集按照80% - 10% - 10%的比例分为训练集、验证集和测试集。一些在2018年数据科学碗上测试的最先进模型,如SSFormer-L、MSRFNet、DoubleUnet、Unet++等,已经取得了显著成果。遵循此数据集的相同分割比例,通过尝试和错误,作者确信670幅图像对于所提出的模型来说是足够的,能够稳健地执行。GlaS数据集包含了165张显微图像及其相应的目标 Mask 标注。在本研究中,作者将GlaS数据集分为85张训练图像和80张测试图像。
Implementation detail
作者使用Pytorch框架实现了整个 Proposal 的架构,并在NVIDIA K80 GPU上进行了实验。作者弃用了Adam优化函数,初始学习率(LR)设为0.0003,并采用了丢弃法正则化,其概率为p=0.2。在连续三个周期未见改进后,新的学习率通过将当前学习率与一个足够小的因子相乘来计算,以降低当前学习率,同时仍能到达全局最小值。2018年数据科学碗挑战赛和GlaS数据集中的所有图像将被调整至256 x 256的分辨率。使用的批处理大小为10,训练作者模型的周期数为300。
Evaluation
在本研究中,作者参考了2018年数据科学碗挑战赛和GlaS数据集中取得显著成果的一些模型,以客观评估作者模型的性能。如表格1所示,Trans2Unet模型的参数数量达到了110M,这是未来研究需要改进的一个缺点,而SSFormer-L模型的参数数量为66.2M。作者提出的网络规模之所以巨大的解释在于,TransUnet分支中使用了ViT模型。如文献[1]所述,ViT有3种变体,包括ViT-Base(86M参数)、ViT-Large(307M参数)和ViT-Huge(632M参数)。考虑到这些规模,作者决定在网络中使用ViT-Base模型。
尽管与当前在该数据集上的其他SOTA架构相比,作者的结果仍然有限,但作者相信,通过这种方法,架构将在未来得到改进。
为了更清楚地展示整合了WASP-KC模块的Trans2Unet模型的改进,该模型的IoU和Dice指标与原始TransUnet以及整合了原始WASP的Trans2Unet模型进行了比较,所有实验均在同一设备上进行测试。表2中的结果显示,作者提出的模型的IoU和Dice指标均无与伦比,具体而言,该模型的IoU和Dice指标分别为86.13%和92.25%。
从表3中可以看出,结果表明,提出的Trans2Unet网络在GlaS数据集上也取得了优异的性能,Dice系数为89.94%,平均IoU为82.54%。
Results
图5:Trans2Unet在2018年数据科学碗挑战数据集中的核图像分割的一些代表性结果。
为了展示新架构在2018年数据科学碗挑战数据集上的性能,作者在图4中展示了学习曲线。如图所示,经过100个周期后,模型损失和包括Dice(DSC)以及IoU在内的分数趋于收敛并保持稳定。为了进行定性评估,作者还展示了该数据集测试集中的一些代表性分割结果,如图5所示。在图5中可以明显看出,所提出方法的预测与 GT 值相吻
合。
V 结论
在本研究中,作者介绍了一种新的架构,它是Unet和TransUnet这两种深度学习网络的结合,用于核图像分割。此外,为了利用级联方法中逐渐提取更大视野(FOV)的方法,作者将带 Short-Cut 的WASP-KC(WASP模块与 Short-Cut )模块集成到TransUnet架构中。
通过对2018年数据科学碗挑战数据集的实验,作者展示了作者提出的模型通过DSC或IoU分数取得了相当好的结果。
通过将Unet与TransUnet架构结合,模型可以保持CNN的局部特征,并利用Transformer的全局特征进行更稳健的分割。
作者相信,这种结构不仅可以提高核细胞模型的效率,也可以为一般的图像分割任务提供一个良好的方法。
参考
[1].Trans2Unet: Neural fusion for Nuclei Semantic Segmentation.
点击上方卡片,关注 「AI视界引擎」 公众号