UNET 3+

论文: https://arxiv.org/abs/2004.08790

github: https://github.com/ZJUGiveLab/UNet-Version

UNet是医学影像分割领域应用最广泛的的网络,其性能和网络中多尺度特征的融合密切相关。此后的UNet++通过嵌套结构和密集的跳过连接原始网络进行了改进。本文提出的UNet3+通过全尺度的连接和深度监督来融合深层和浅层特征的同时对各个尺度的特征进行监督。提出的UNet3+网络可以在减少网络参数的同时提高计算效率,在两个数据集上验证了方法有效性。相关代码已经开源。

现有的分割网络如UNET、PSPNET和DeepLab等网络都通常会通过多尺度的方式提取图像的信息。低层次的细节特征图中具有更丰富的例如边界这样空间信息,高层特征图中包含更多的例如物体位置这样的高级语义特征。然而,随着网络的下采样和上采样,这些高低层的信息并没有被充分地利用。因此,文章提出的UNet3+对网络的编码器与解码器连接以及解码器内部之间的连接进行了改进。此外,文章通过提出的混合损失函数对各层进行深度监督和分类分支指导分割的方式,进一步提高了分割的精度。总结来说,文章主要有以下四点贡献:

  1. 设计了一种新的网络结构UNet3+,通过引入全尺度的跳过连接,在全尺度特征映射中融合了低层细节和高层语义,充分利用了多尺度特征的同时具有更少的参数;
  2. 通过深度监督让网络从全尺度特征中学习分割表示,提出了更优的混合损失函数以增强器官的边界;
  3. 提出分类指导模块,通过与图像分类分支联合训练的方式,减少了网络在非器官图像的过度分割(over-segmentation);
  4. 在肝脏和脾脏数据集上进行了广泛的实验,证明了UNet 3+的有效性。
  5. 从图中可以看到,UNet3+与UNet主体上非常相似,不同之处在于从编码器到解码器的跳过连接以及不同层级的编码器之间的连接。以图中的节点 XDe3 为例,它的信息来自于两方面,一是比其更浅(包括同一层级)的编码器,二十比其更深的解码器。不同层级的特征通过maxpooling和双线性上采样的方式进行尺寸统一。解码层的卷积分两步,第一步是对来及各个节点的信息进行各自的卷积,第二步是对堆叠的特征通过卷积来进行信息的融合和提取。值得注意的一个细节是,进行第一个卷积层时,来自各层的数据被卷积到相同的特征图数(在这里是n/5,n为所在层的特征图数)。

从图中可以看到,UNet3+与UNet主体上非常相似,不同之处在于从编码器到解码器的跳过连接以及不同层级的编码器之间的连接。以图中的节点 XDe3 为例,它的信息来自于两方面,一是比其更浅(包括同一层级)的编码器,二是比其更深的解码器。不同层级的特征通过maxpooling和双线性上采样的方式进行尺寸统一。解码层的卷积分两步,第一步是对来及各个节点的信息进行各自的卷积,第二步是对堆叠的特征通过卷积来进行信息的融合和提取。值得注意的一个细节是,进行第一个卷积层时,来自各层的数据被卷积到相同的特征图数(在这里是n/5,n为所在层的特征图数)。

2.全尺度的深度监督

为了进一步优化网络对图像边界的分割,文章借鉴了图像质量评估中常用的多尺度SSIM(MS-SSIM)提出了MS-SSIM loss。

本文最终采用了混合损失函数(focal loss,ms-ssim loss和iou loss)来对各层进行监督。

ℓseg=ℓfl+ℓms−ssim+ℓiou

3. 分类指导模块(CGM)

在大多数医学图像分割中,非器官图像中出现假阳性不可避免。这通常是保留在较浅层中背景噪声信息导致的过分割现象。为了实现更精确的分割,文章尝试通过添加一个额外的分类任务来解决这个问题,该分类任务被设计用于预测输入图像是否有器官。简单来说,当预测到图像包含待分割器官的概率较小时,对输出图像乘以0使得输出全黑。

文章采用了LITS的肝脏数据集和自己采集的脾脏数据集通过两组实验来进行验证。

第一组对UNet、UNet++、UNet3+(带深度监督和不带)以Vgg和ResNet101作为backbone进行了对比。可以在以Vgg为backbone时,UNet3+比其UNet在两个数据集上分别有2.8%和4.1%的提升。网络相比于UNet++也有较大的提升。另外,可以看到UNet3+使用了更少的参数得到了更好的结果。可视化的结果表明即使在器官较小的情况下网络也能得到更加精细连贯的分割。

文章进一步以ResNet作为backbone,将网络与当前比较先进的分割网络进行对比。在这里,在验证网络有效性的同时,文章对提出的损失函数和分类分支进行了消融实验。

文章对之前在UNet解码器只接收的来自同一层编码器和深一层解码器的连接方式进行了改进,使得解码器都能获得来自每一个更浅的编码器和更深的解码器的信息,使得网络能够更好地提取和融合多尺度的信息。网络的结构设计简洁优雅,是一篇非常不错的UNet改进文章。另外文章提出的MS-SSIM损失和分类指导模块也挺有意思。当然我对文章也有一些思考。第一,网络结构设计中,对于来自不同层级的特征,进行融合时可以考虑通过PSP或者Deeplab的方式(JPU是一种很好选择),也可以考虑通过SE的方式来进行通道的选择。第二,文章提出的MS-SSIM能够更好地分割图像的边界,那么选取豪斯多夫距离这样的指标可以更好地证明方法的有效性。第三,文章通过CGM来对输出进行限制,但是对于器官的顶端和底端这样本身有比较多歧义性图像,容易造成无法分割的情况,可以考虑进行soft的指导。

发表评论

您的电子邮箱地址不会被公开。 必填项已用*标注