今天读了一篇SOTA的分割论文:PraNet:Parallel Reverse Attention Network for Polyp Segmentation(中文可以翻译为:PraNet:用于肠息肉分割的并行反向注意力网络),效果表现炸裂,在各个数据集上平均超过前SOTA 7个以上百分点,并且提供了代码。
摘要
结肠镜检查是一种检测大肠息肉的有效技术,息肉与大肠癌高度相关。在临床实践中,从结肠镜检查图像中分离息肉非常重要,因为它为诊断和手术提供了有价值的信息。但是,准确的息肉分割术是一项艰巨的任务,主要有两个原因:
(i)相同类型的息肉具有不同的大小,颜色和质地;
(ii)息肉及其周围粘膜之间的边界不清晰。
为了解决这些挑战,我们提出了一个并行反向注意网络(PraNet),用于在结肠镜检查图像中进行精确的息肉分割。具体来说,我们首先使用并行部分解码器(PPD)对高层级网络的特征进行聚合,基于聚合特征,我们生成一个全局map,并作为后面组件的初始指导区域(guidance area)。此外,我们使用反向注意力(RA)模块挖掘边界线索(boundary cues),该模块能够建立区域与边界线索之间的关系。由于区域和边界之间的循环合作(recurrent cooperation)机制,我们的PraNet能够校准一些未对齐的预测,从而提高了分割精度。在六个指标上对五个具有挑战性的数据集进行定量和定性评估表明,我们的PraNet显著提高了分割精度,并在泛化性能和实时分割效率(〜50fps)方面显示了许多优势。
1. Introduction
在本文中,我们提出了一种新颖的深度神经网络,称为并行反向注意力网络(PraNet),用于息肉分割任务。
动机:
在进行息肉标注时,临床医生首先粗略定位息肉,然后再根据局部特征准确提取其轮廓。因此,我们认为面积和边界是区分正常组织和息肉的两个关键特征。
做法:
我们首先预测粗糙区域,然后通过反向注意力隐式地对边界建模。
该策略具有三个优点,包括更好的学习能力,改进的泛化能力和更高的训练效率。
贡献:
(1)通过使用并行部分解码器(PPD)聚合高层网络层中的特征,生成全局特征图(global map),并将其作为后续步骤的初始指导区域(guidance area)。
为了进一步挖掘边界线索,利用一组循环反向注意力(RA)模块(recurrent reverse attention mudules)来建立区域与边界线索之间的关系。由于区域和边界之间的这种循环合作机制,我们的模型能够校准一些未对齐的预测。
(2)我们针对息肉分割引入了几种新颖的评估指标(evaluation metrics),并为公开可用的现有SOTA模型提供了一个综合基准。
(3)广泛的实验表明,在五个具有挑战性的数据集上,所提出的PraNet优于大多数最先进的模型,并且大大提高了SOTA的实时推断和训练时间。
2. 算法
图1 展示了PraNet的整体架构,它使用并行部分解码器生成高级语义全局图和一组反向注意力模块,用于从结肠镜检查图像中进行精确的息肉分割。 每个组件将详细说明如下。
2.1 通过并行部分解码器进行特征聚合
当前流行的医学图像分割网络通常依赖于UNet 或类似U-Net的网络(例如,U-Net++,ResUNet等)。这些模型实质上是编码器-解码器框架,通常会汇总从CNN中提取的所有层级的特征。Wu的文章表明,与高层级的特征相比,低层级的特征(由于其较大的空间分辨率)需要更多的计算资源,但对性能的贡献较小。受此观察启发,我们提出并行的部分解码器组件来聚合高层级特征。更具体地,对于尺寸为的输入息肉图像,可以通过基于Res2Net的骨干网络,生成五层特征, 每层分辨率为。 然后,将分为低层级特征和高层级特征。我们引入了部分解码器(一个新的SOTA解码器组件),对高层级特征进行聚合,通过计算部分解码器特征,最终获得全局图。
2.2 反向注意力模块
在临床环境中,医生首先大致定位息肉区域,然后仔细检查局部组织以准确标记息肉。 如2.1节所述,我们的全局图来自最深的CNN层,该层只能捕获息肉组织的相对粗糙的位置,而没有结构细节(参见图1)。 为了解决这个问题,我们提出了一种原则性策略,通过逐步擦除前景物体的方式挖掘更具判别性的息肉区域[27,4]。 与其像[4,13,36,33]中那样聚合所有层级的特征,我们提出自适应地学习三个并行高层级特征中的反向注意力。 换句话说,我们的架构可以通过擦除从高层级输出特征中评估出的现有息肉区域,来依次挖掘互补的息肉区域和细节,而现有的息肉区域评估是从更深层特征进行上采样得到的。
具体来说,我们通过将高层级侧输出特征乘以(逐元素$$)反向注意力权重来获得输出反向注意力特征,如下所示:
反向注意权重实际上是计算机视觉界中用于显着对象检测的事实[4,34],其公式可表示为:
其中,表示上采样操作,表示Sigmoid函数,表示从元素全为1的矩阵中减去输入的操作。图1的RA给出了这个过程的细节。值得注意的是,由反向注意力驱动的擦除策略最终可以将不精确和粗略的估计精炼成准确而完整的预测图。
2.3 学习过程和实施细节
损失函数。 损失函数定义为,其中和分别表示在全局限制和局部限制(像素级)下的加权IoU损失和二元交叉熵(BCE)损失。和标准IoU不同,加权IoU损失增加了难像素的权重,BCE损失也给难像素更大权重。这两个损失函数的定义同[21,26],已经在物体显著性监测领域得到了有效性验证。这里,我们对3个边侧输出()和全局图都进行监督,每一个图都上采样到和真实值一样的大小。这样总损失可以表示为:
实施细节。基于Pytorch,GPU选用一块NVIDIA TITAN RTX,所有输入图片归一化到并采用多尺度训练策略, 不使用数据扩增。优化算法采用Adam,初始学习率,整个网络采用端到端的训练方式,批大小为16,20个epochs,32分钟收敛。最终的预测图通过对进行Sigmoid操作后得出。