UNET 3+: A FULL-SCALE CONNECTED UNET FOR MEDICAL IMAGE SEGMENTATION
unet3+:一种用于医学图像分割的全连通UNET
arXiv.2004.08790
文章地址:https://arxiv.org/abs/2004.08790
代码地址:https://github.com/ZJUGiveLab/UNet-Version
摘要
最近,人们对基于深度学习的语义分割产生了越来越大的兴趣。UNet是具有编码器-解码器结构的深度学习网络之一,被广泛用于医学图像分割。结合多尺度特征是准确分割的重要因素之一。UNet++是通过设计一个具有嵌套和密集跳过连接的架构,作为一个改进的Unet而开发的。然而,它并没有从全尺度中发掘足够的信息,仍有很大的改进空间。在本文中,我们提出了一个新颖的UNet 3+,它利用了全尺度跳过连接和深度监督的优势。全尺度跳过连接将低层次的细节与不同尺度的特征图的高层次语义结合起来;而深度监督则从全尺度聚合的特征图中学习层次化的表示。所提出的方法对出现在不同尺度的器官特别有利。除了准确性的提高,提议的UNet 3+可以减少网络参数以提高计算效率。我们进一步提出了一个混合损失函数,并设计了一个分类指导模块,以增强器官边界,减少非器官图像中的过度分割,产生更准确的分割结果。我们在两个数据集上证明了所提方法的有效性。代码见:https://github.com/ZJUGiveLab/UNet-Version
索引项–分割,全尺寸跳过连接,深度监督,混合损失函数,分类
1导言
医学图像中的自动器官分割是许多临床应用中的一个关键步骤。最近,卷积神经网络(CNN)极大地促进了各种分割模型的发展,如全卷积神经网络(FCN)[1]、UNet[2]、PSPNet[3]和一系列DeepLab版本[4-6]。特别是UNet,它是基于编码器-解码器架构的,被广泛用于医学图像分割中。它使用跳过连接来结合来自解码器的高层次语义特征图和来自编码器的相应低层次详细特征图。
为了减少UNet中普通跳过连接对语义不同的特征的融合,UNet++[7]通过引入嵌套和密集的跳过连接进一步加强了这些连接,目的是减少编码器和解码器之间的语义差距。尽管取得了良好的性能,但这种类型的方法仍然无法从全尺度中探索出足够的信息。 正如许多分割研究中所见证的那样[1-7],不同尺度的特征图可以探索出不同的信息。低级别的详细特征图捕捉丰富的空间信息,突出器官的边界;而高级别的语义特征图体现了位置信息,定位器官的位置。然而,在逐步下采样和上采样时,这些精美的信号可能会逐渐被稀释。为了充分利用多尺度特征,我们提出了一个新的基于U形的架构,命名为UNet 3+,其中我们重新设计了编码器和解码器之间的互连以及解码器之间的内部互连,以便从全尺度上捕捉细粒度的细节和粗粒度的语义。为了进一步从全比例聚集的特征图中学习分层表征,每侧输出都与混合损失函数相连,这有助于准确分割,特别是对于在医学图像体中出现不同比例的器官。除了准确性的提高,我们还表明,所提出的UNet 3+可以减少网络参数以提高计算效率。
为了解决医学图像中更精确的分割需求,我们进一步研究如何有效地减少非器官图像中的假阳性。现有的方法通过引入注意力机制[8]或在推理时进行预先定义的细化方法,如CRF[4]来解决这个问题。与这些方法不同,我们将分类任务扩展到预测输入图像是否有器官,为分割任务提供指导。 综上所述,我们的主要贡献有四个方面。(i) 设计了一个新颖的UNet 3+,通过引入全尺度的跳过连接来充分利用多尺度的特征,它结合了全尺度特征图的低层次细节和高层次语义,但参数较少;(ii) 开发了一个深度监督,从全尺度聚合的特征图中学习层次化的表征,它优化了一个混合损失函数来增强器官边界。(iii) 提出一个分类指导模块,通过与图像级分类联合训练,减少非器官图像的过度分割;(iv) 在肝脏和脾脏数据集上进行广泛的实验,其中UNet 3+比一些基线产生一致的改进。
2.方法
图1给出了UNet、UNet++和提议的UNet 3+的简化概述。与UNet和UNet++相比,UNet 3+通过重新设计跳过连接以及利用全面的深度监督来结合多尺度特征,它提供了较少的参数,但产生了一个更准确的位置感知和边界增强的分割图。
2.1. 全面跳过连接
提出的全尺寸跳过连接转换编码器和解码器之间的互连以及解码器子网络之间的内部连接。具有普通连接的UNet和具有嵌套和密集连接的UNet++都缺乏从全尺度上探索足够信息的能力,无法明确了解器官的位置和边界。为了弥补UNet和UNet++中的缺陷,UNet3+中的每个解码器层都合并了来自编码器的较小和相同比例的特征图以及来自解码器的较大比例的特征地图,从而以全比例捕获细粒度细节和粗粒度语义。
作为示例,图2说明了如何构建
X
D
e
3
X^3_{De}
XDe3。与UNet类似,来自同一尺度编码器层
X
E
n
3
X^3_{En}
XEn3的特征映射直接在解码器中接收。与UNet相反,一组编码器间解码跳过连接从较小尺度编码器层传递低级别详细信息
X
E
n
1
X^1_{En}
XEn1和
X
E
n
2
X^2_{En}
XEn2(通过应用非重叠最大池操作;而解码器内跳过连接链从更大规模的解码器层传输高级语义信息
X
E
n
4
X^4_{En}
XEn4和
X
E
n
5
X^5_{En}
XEn5,通过利用双线性插值。有了五个相同分辨率的特征图,我们需要进一步统一频道的数量,并减少多余的信息。3×3的64个滤波器的卷积可以是一个令人满意的选择。为了将浅层精细信息与深层语义信息无缝融合,我们进一步从五个尺度对级联的特征图执行特征聚合机制,该机制由320个大小为3×3的过滤器、批量规范化和ReLU激活功能组成。形式上,我们将跳过连接公式化为:让
i
i
i沿着编码器对下采样层进行索引,
N
N
N表示编码器的总数。由
X
D
e
i
X^i_{De}
XDei,计算为:
X
D
e
i
=
{
X
E
n
i
,
i
=
N
H
(
[
C
(
D
(
X
E
n
k
)
)
k
=
1
i
−
1
,
C
(
X
E
n
i
)
C
(
U
(
X
D
e
k
)
)
k
=
i
+
1
N
]
)
i
=
1
,
⋯
,
N
−
1
(1)
X^i_{De}=\left\{ \begin{matrix}X^i_{En},i=N\\ \\ H\left( {\left[ {C\left( {D\left( {X_{En}^k} \right)} \right)_{k = 1}^{i - 1},C\left( {X_{En}^{\rm{i}}} \right)C\left( {U\left( {X_{D{\rm{e}}}^k} \right)} \right)_{k = i + 1}^N} \right]} \right){\rm{i}} = 1, \cdots ,N - 1\tag{1} \end{matrix} \right.
XDei=⎩
⎨
⎧XEni,i=NH([C(D(XEnk))k=1i−1,C(XEni)C(U(XDek))k=i+1N])i=1,⋯,N−1(1)
where函数 C ( ⋅ ) C(\cdot ) C(⋅)表示卷积运算, H ( ⋅ ) H(\cdot ) H(⋅)通过卷积实现了特征聚合机制,随后是批量规范化和ReLU激活功能。 D ( ⋅ ) D(\cdot ) D(⋅) 和 U ( ⋅ ) U(\cdot ) U(⋅) 分别指示上采样和下采样操作,以及 [ ⋅ ] [\cdot ] [⋅]表示连接。
值得一提的是,我们建议的UNet 3+更高效,参数更少。在编码器子网中,UNet、UNet++和UNet3+共享相同的结构,其中
X
E
n
i
X^i_{En}
XEni有
32
×
2
i
32×2^i
32×2i通道。对于解码器,UNet中特征图的深度与编码器对称,因此
X
D
e
i
X^i_{De}
XDei也有
32
×
2
i
32×2^i
32×2i通道信道。UNet(
P
U
−
D
e
i
P^i_{U-De}
PU−Dei)的
i
t
h
i^{th}
ith解码器级中的参数数量可计算为:
P
U
−
D
e
i
=
D
F
×
D
F
×
[
d
(
X
D
e
i
+
1
)
×
d
(
X
D
e
i
)
+
d
(
X
D
e
i
)
2
+
d
(
X
E
n
i
+
X
D
e
i
)
×
d
(
X
D
e
i
)
]
(2)
P^i_{U-De}=D_F×D_F×[d(X^{i+1}_{De})×d(X^{i}_{De})+d(X^{i}_{De})^2+d(X^{i}_{En}+X^i_{De})×d(X^i_{De})]\tag{2}
PU−Dei=DF×DF×[d(XDei+1)×d(XDei)+d(XDei)2+d(XEni+XDei)×d(XDei)](2)
其中
D
F
D_F
DF是卷积核大小,
D
(
⋅
)
D(\cdot)
D(⋅) 表示节点的深度。当涉及UNet++时,它利用每个跳过路径上的密集卷积块,其中
P
U
+
+
−
D
e
i
P^i_{U^{++}-De}
PU++−Dei可以计算为:
P
U
+
+
−
D
e
i
=
D
F
×
D
F
×
[
d
(
X
D
e
i
+
1
)
×
d
(
X
D
e
i
)
+
d
(
X
D
e
i
)
2
+
d
(
X
E
n
i
+
∑
k
=
1
N
−
1
−
i
X
M
e
i
,
k
+
X
D
e
i
)
×
d
(
X
D
e
i
)
]
(3)
P^i_{U^{++}-De}=D_F×D_F×[d(X^{i+1}_{De})×d(X^{i}_{De})+d(X^{i}_{De})^2+d(X^{i}_{En}+\sum^{N-1-i}_{k=1}X^{i,k}_{Me}+X^i_{De})×d(X^i_{De})]\tag{3}
PU++−Dei=DF×DF×[d(XDei+1)×d(XDei)+d(XDei)2+d(XEni+k=1∑N−1−iXMei,k+XDei)×d(XDei)](3)
如图所示,
P
U
+
+
−
D
e
i
P^i_{U^{++}-De}
PU++−Dei大于
P
U
−
D
e
i
P^i_{U-De}
PU−Dei。而在UNet 3+中,每个解码器特征图都是从N尺度导出的,产生64×N信道。
P
U
3
+
−
D
e
i
P^i_{U^{3+}-De}
PU3+−Dei可计算为:
P
U
3
+
−
D
e
i
=
D
F
×
D
F
×
[
(
∑
k
=
1
i
d
(
X
E
n
k
)
+
∑
k
=
i
+
1
N
d
(
X
D
e
k
)
)
×
64
+
d
(
X
D
e
i
)
2
]
(4)
P^i_{U^{3+}-De}=D_F×D_F×[(\sum^i_{k=1}d(X^k_{En})+\sum^N_{k=i+1}d(X^k_{De}))×64+d(X^i_{De})^2] \tag{4}
PU3+−Dei=DF×DF×[(k=1∑id(XEnk)+k=i+1∑Nd(XDek))×64+d(XDei)2](4)
为了减少通道,UNet 3+中的参数比UNet和UNet++中的少。
2.2全面深度监管
为了从全尺度聚集的特征图中学习分层表征,UNet 3+中进一步采用了全尺度深度监督。与UNet++中对生成的全分辨率特征图进行的深度监督相比,提议的UNet 3+从每个解码器阶段产生一个侧面输出,该输出由地面真实监督。为了实现深度监督,每个解码器阶段的最后一层被送入一个普通的3×3卷积层,然后是一个双线性上采样和一个sigmoid函数。
为了进一步增强器官的边界,我们提出了一个多尺度结构相似性指数(MS-SSIM)[9]损失函数,为模糊边界分配更高的权重。得益于此,UNet 3+将关注模糊边界,因为区域分布差异越大,MS-SSIM值越高。两个相应的
N
×
N
N×N
N×N大小的斑块从分割结果P和地面真实掩码G中裁剪出来,可以表示为
p
=
p
j
:
j
=
1
,
⋅
⋅
⋅
,
N
2
p={p_j:j=1,···,N^2}
p=pj:j=1,⋅⋅⋅,N2和
g
=
g
j
:
j
=
1
,
⋅
⋅
⋅
,
N
2
g={g_j:j=1,···,N^2}
g=gj:j=1,⋅⋅⋅,N2。
p
p
p和
g
g
g的MS-SSIM损失函数定义为
ℓ
m
s
−
s
s
i
m
=
1
−
∏
m
=
1
M
(
2
μ
p
μ
g
+
C
1
μ
p
2
+
μ
g
2
+
C
1
)
β
m
(
2
σ
p
g
+
C
2
σ
p
2
+
σ
g
2
+
C
2
)
(5)
{\ell _{ms - ssim}} = 1 - {\prod\limits_{m = 1}^M {\left( {{{2{\mu _p}{\mu _g} + {C_1}} \over {\mu _p^2 + \mu _g^2 + {C_1}}}} \right)} ^{{\beta _m}}}\left( {{{2{\sigma _{pg}} + {C_2}} \over {\sigma _p^2 + \sigma _g^2 + {C_2}}}} \right)\tag{5}
ℓms−ssim=1−m=1∏M(μp2+μg2+C12μpμg+C1)βm(σp2+σg2+C22σpg+C2)(5)
其中k是标度的总数, μ p 、 μ g \mu _p、\mu _g μp、μg和 σ p , σ p \sigma _p,\sigma _p σp,σp是 p p p和 g g g的均值和标准差, σ p g \sigma _{pg} σpg表示它们的协方差。 β m \beta_m βm和 γ m \gamma _m γm定义了每个量表中两个分量的相对重要性,根据[9]设置。添加两个小常数 C 1 = 0.0 1 2 C_1=0.01^2 C1=0.012(和 C 2 = 0.0 3 2 C_2=0.03^2 C2=0.032),以避免被零除的不稳定情况。在我们的实验中,我们根据[9]将比例设置为5。
通过组合焦距损失(
ℓ
f
l
ℓ_{fl}
ℓfl )[10],MS-SSIM丢失(
ℓ
m
s
−
s
s
i
m
ℓ_{ms-ssim}
ℓms−ssim) 和IuU损失(
ℓ
i
o
u
ℓ_{iou}
ℓiou)[11],我们开发了一种用于三级层次(像素级、补丁级和地图级)分割的混合损失,它能够捕捉具有清晰边界的大规模和精细结构。混合分割损失(
ℓ
s
e
g
ℓ_{seg}
ℓseg) 定义为:
ℓ
s
e
g
=
ℓ
f
l
+
ℓ
m
s
−
s
s
i
m
+
ℓ
i
o
u
(6)
ℓ_{seg}=ℓ_{fl}+ℓ_{ms-ssim}+ℓ_{iou}\tag{6}
ℓseg=ℓfl+ℓms−ssim+ℓiou(6)
2.3.分类指导模块(CGM)
在大多数医学图像分割中,在非器官图像中出现假阳性是一个不可避免的情况。 它很可能是由残留在较浅层的背景噪声信息引起的,导致了过度分割的现象。为了实现更准确的分割,我们试图通过增加一个额外的分类任务来解决这个问题,该任务旨在预测输入图像是否有器官。
如图3所示,在通过一系列的操作,包括丢弃、卷积、maxpooling和sigmoid之后,从最深层的 X E n 5 X^5_{En} XEn5中产生了一个二维张量,每个张量都代表有/无器官的概率。受益于最丰富的语义信息,分类结果可以在两个步骤中进一步指导每个分割侧的输出。首先,在argmax函数的帮助下,二维张量被转移到一个{0,1}的单一输出,表示有/无器官。随后,我们将单一分类输出与侧面分割输出相乘。由于二元分类任务的简单性,该模块在二元交叉熵损失函数[12]的优化下毫不费力地取得了准确的分类结果,实现了对非器官图像过度分割缺点的补救指导。
3 实验和结果
3.1.数据集和实施该方法在肝脏和脾脏两个器官上进行了验证。
肝脏分割数据集来自ISBI LiTS 2017挑战赛。它包含131个增强3D腹部CT扫描,其中103个和28个体积分别用于训练和测试。来自医院的脾脏数据集通过了伦理认证,包括40和9个用于训练和测试的CT体积。为了加快训练速度,输入图像有三个通道,包括要分割的切片和上下切片,被裁剪为320×320。我们利用随机梯度下降来优化网络,并将其超参数设置为默认值。骰子系数用作每种情况的评估指标。
3.2.与UNet和UNet++的比较
在本节中,我们首先将建议的UNet 3+与UNet和UNet+进行了比较。每种方法中使用的损失函数是焦距损失。
(i) 定量比较:基于Vgg-16和ResNet-101的主干,表1比较了UNet、UNet++和拟议的UNet3+架构在肝脏和脾脏数据集上的参数数量和分割精度。如图所示,没有深度监管的UNet 3+实现了优于UNet和UNet++的性能,在两个数据集上执行的两个主干之间平均提高2.7和1.6点。考虑到肝脏和脾脏在CT切片中以不同的比例出现,UNet 3+结合全面深度监护进一步提高了0.4分。
(ii)定性比较:图2显示了基于ResNet-101的UNet、UNet++和UNet3+的分割结果,并对肝脏数据集进行了全面深度监控。可以观察到,我们提出的方法不仅准确地定位器官,而且即使在小物体环境下也能产生相干边界。
3.3.与现有技术的比较
我们将基于ResNet-101的UNet 3+与几种最新的先进方法进行了定量比较:PSPNet[3]、DeepLabV2[4]、DepeLabV3[5]、DeepLabV3+[6]和Attention UNet[8]。值得一提的是,所有结果都直接来自于单模型测试,而不依赖于任何后处理工具。此外,所有的网络都通过他们自己的论文中提出的损失函数进行了优化。
M e t h o d Method Method | D i c e l i v e r Dice_{liver} Diceliver | D i c e s p l e e n Dice_{spleen} Dicespleen |
---|---|---|
PSPNet [3] | 0.9242 | 0.9240 |
DeepLabV2 [4] | 0.9021 | 0.9097 |
DeepLabV3 [5] | 0.9217 | 0.9217 |
DeepLabV3+ [6] | 0.9186 | 0.9290 |
Attention UNet [8] | 0.9341 | 0.9324 |
UNet 3+ (focal loss) | 0.9601 | 0.9560 |
UNet 3+ (Hybrid loss) | 0.9643 | 0.9588 |
UNet 3+ (Hybrid loss + CGM) | 0.9675 | 0.9620 |
表2总结了定量比较结果。如图所示,所提出的混合损失函数通过考虑像素级、补丁级和地图级优化,大大提高了性能。特别是,贴片级MSSSIM损失函数有助于为模糊边界分配更高的权重,从而产生更增强的边界感知分割图。此外,利用分类指导模块,UNet 3+巧妙地避免了复杂背景下的过度分割。可以看出,与所有其他以前的方法相比,这种方法是突出的。还值得注意的是,所提出的方法在肝脏(0.9675对0.9341)和脾脏(0.9620对0.9324)方面优于次优结果。
4. 结论
在本文中,我们提出了一个全面连接的UNet,名为UNet3+,具有深度监控,以便最大限度地利用全尺度的特征图,以较少的参数实现准确的分割和高效的网络架构。进一步引入了分类引导模块和混合损失函数,以产生更准确的位置感知和边界感知分割图。对肝脏和脾脏数据集的实验结果表明,UNet 3+超越了所有以前的最先进方法,突出了器官并产生了连贯的边界。
5.感谢
确认此项工作部分得到了浙江省实验室重大科研项目(2018DG0ZX01号)、杭州市重点科技创新支持计划(20172011A038号)、日本教育、科学、文化部科学研究资助项目(2017DG0ZZ01号)的支持,文化与体育(MEXT)。18H03267和编号17H00754。
6. REFERENCES
[1] J. Long, E. Shelhamer and T. Darrell, “Fully Convolutional Networks for Semantic Segmentation,” The IEEE Conference on Computer Vision and Pattern Recognition, pp. 3431-3440, 2015.
[2] O. Ronneberger, P. Fischer and T. Brox, “U-Net: Convolutional Networks for Biomedical Image Segmentation,” Medical Image Computing and Computer-Assisted Intervention, pp.234241, 2015.
[3] H.S. Zhao, J.P. Shi, X.J. Qi, X.G. Wang and J.Y. Jia, “Pyramid scene parsing network,” The IEEE Conference on Computer Vision and Pattern Recognition, pp. 2881-2890, 2017.
[4] L-C. Chen, G. Papandreou, I. Kokkinos, K. Murphy and A.L. Yuille, “Semantic image segmentation with deep convolutional nets, atrous convolution, and fully connected crfs,” IEEE transactions on pattern analysis and machine intelligence, vol.40, no.4, pp: 834-848, 2018.
[5] L-C. Chen, G. Papandreou, F. Schroff and H. Adam, “Rethinking atrous convolution for semantic image segmentation”, arXiv preprint arXiv:1706.05587, 2017.
[6] L-C. Chen, Y.K. Zhu, G. Papandreou and H. Adam, “Encoderdecoder with atrous separable convolution for semantic image segmentation,” Proceedings of the European Conference on Computer Vision, 2018.
[7] Z.W. Zhou, M.M.R. Siddiquee, N. Tajbakhsh and J.M. Liang, “UNet++: A Nested U-Net Architecture for Medical Image Segmentation,” Deep Learning in Medical Image Anylysis and Multimodal Learning for Clinical Decision Support, pp: 3-11, 2018.
[8] O.O et al., “Attention u-net: Learning where to look for the pancreas,” Medical Imaging with Deep Learning, 2018.
[9] Z. Wang, E.P. Simoncelli and A.C. Bovik, “Multiscale structural similarity for image quality assessment,” The Thrity-Seventh Asilomar Conference on Signals, Systems & Computers, 2003.
[10] T.-Y. Lin, P. Goyal, R. Girshick, K.M. He and P. Dollar. “Focal loss for dense object detection,” The IEEE international conference on computer vision, pp. 2980-2988, 2017.
[11] G. Mattyus, W.J. Luo, and R. Urtasun, “Deep-roadmapper: Extracting road topology from aerial images”, The IEEE international conference on computer vision, 2017.
[12] P.-T. Boer et al., “A tutorial on the cross-entropy method,” Annals of Operations Research, vol.134, no.1, pp. 19–67, 2005.