代码连接
这篇文章主要是提出了多视图以及多视图的损失,提高了模型旋转对比的不变性。同时模型在多个数据集上进行预训练,随后通过自监督学习微调,从而达到了较高的精度。能同时实现2D/3D分割。
摘要
近年来,大规模视觉 transformer 的进步在改善医学图像分割的预训练模型方面取得了重大进展。然而,这些方法在获取大量的预训练数据方面面临着显著的挑战,特别是在医学领域。为了解决这一限制,我们提出了带有Swin transformer的遮罩多视图(SwinMM),这是一种新的多视图管道,用于实现准确和数据高效的自监督医学图像分析。我们的策略通过结合两个主要组成部分来利用多视图信息的潜力。在预训练阶段,我们部署了一个屏蔽多视图编码器,旨在通过一系列不同的代理任务并发训练屏蔽多视图观测值。这些任务包括图像重建、旋转、对比学习和一个采用相互学习范式的新任务。这项新任务利用了来自不同角度的预测之间的一致性,从而能够从3D医疗数据中提取隐藏的多视图信息。在微调阶段,开发了一种跨视点解码器,通过交叉注意块聚合多视点信息。与之前最先进的自监督学习方法Swin UNETR相比,SwinMM在多个医学图像分割任务上表现出显著的优势。它允许多视图信息的平滑集成,显著提高了模型的准确性和数据效率。
引言
医学图像分割是计算机辅助诊断、治疗计划和干预的一项关键任务。虽然大型 transformer 在各种计算机视觉任务中表现出令人印象深刻的性能,例如自然图像的识别、检测和分割[5,16],它们在应用于医学图像分析时面临着重大挑战。主要的挑战是标记医学图像的稀缺性,因为很难收集和标记它们,这需要专业的医学知识,并且耗时[23,25,12]。第二个挑战是识别医学图像中稀疏和模糊模式的能力,包括具有小分割目标的模糊和暗淡图像。因此,为了提高医学图像分析网络在计算机辅助医学诊断中的准确性和可靠性,迫切需要开发一种精确、高效的医学图像分析网络管道。
自监督学习是一种通过设计 pretext task 来构建特征嵌入空间的技术,已成为解决标签不足问题的一种有希望的解决方案。自监督学习的一种代表性方法是掩码自编码器(MAE)[11]。MAEs学习在随机屏蔽某些输入特征后重建输入数据。该方法已成功地应用于各种应用,包括图像去噪、文本补全、异常检测和特征学习。在医学图像分析领域,MAE预训练也被发现是有效的[32]。然而,这些研究存在局限性,因为它们需要大量未标记的数据,并且没有优先考虑提高输出可靠性,这可能会破坏它们在现实世界中的实用性。MAE 这类方法的缺点
在本文中,我们提出了基于Swin的掩膜多视图(SwinMM),这是第一个用于自监督医学图像分割的综合多视图管道。我们从以往的研究中得到启发[31,33,26,28],旨在通过将多视图学习纳入自监督学习管道来提高输出可靠性和数据利用率。在预训练阶段,提出的方法随机屏蔽3D医学图像,并从不同的角度创建各种观察结果。屏蔽多视图编码器同时处理这些观察结果,以完成四个代理任务:图像重建,旋转,对比学习,以及利用相互学习范式最大化不同视图预测之间一致性的新型代理任务。该方法有效地利用了3D医疗数据中隐藏的多视图信息,并允许编码器学习原始图像的丰富高级表示,这有利于下游分割任务。在微调阶段,来自同一图像的不同视图被编码成一系列表征,这些表征将在专门设计的跨视图注意块中相互作用。通过多视图一致性损失,从多个角度生成对齐的输出预测,提高了最终输出的可靠性和精度。在SwinMM中使用的不同视图的互补性导致更高的精度,需要更少的训练数据和注释,这对于推进该领域的最新技术具有重大潜力。综上所述,我们研究领域的贡献如下: 这里的多个视图是多个看的角度吗?还是类似多种模态,如何实现的?
我们提出了SwinMM,一个独特的、数据高效的3D医学图像分析管道,在该领域提供了第一个全面的多视图、自我监督的方法。
我们的设计包括一个屏蔽多视图编码器和一个新的基于相互学习的代理任务,促进有效的自监督预训练。
我们结合了一个跨视图解码器,通过交叉注意块优化多视图信息的利用。
SwinMM在WORD数据集上的平均Dice得分为86.18%,在数据效率和分割性能方面都优于其他领先的分割方法。
方法、
图1提供了SwinMM的概述,包括一个屏蔽多视图编码器和一个跨视图解码器。SwinMM通过随机屏蔽输入图像来创建多个视图,随后将这些被屏蔽的视图输入编码器进行自监督预训练。在微调阶段,我们在解码器中构建了一个跨视图注意模块。这种设计有利于多视图信息的有效利用,能够生成更精确的分割预测。
预训练
遮罩多视图编码器。跟随[11],我们将3D图像分成大小相同的子卷,并随机屏蔽其中的一部分,如图2所示。这些被屏蔽的3D补丁,从不同的角度,然后被屏蔽的多视图编码器用于自监督预训练。如图1所示,编码器由一个补丁划分层、一个补丁嵌入层和四个Swin Transformer层组成[17]。值得注意的是,与典型的变压器编码器不同,我们的遮罩多视图编码器可以处理来自不同视图的不同图像的多个输入这里的多个视图应该就是图2里的旋转和交换次序,但是我觉得交换次序是不对的,这破坏了空间信息,如果不交换次序是不是能达到更好的结果
,使其更适用于广泛的应用。
预训练的策略。为了结合三维体的多个视角,我们从不同的观察角度生成了视图,包括轴向、冠状和矢状。此外,我们应用了与每个视角对齐的旋转操作,包括沿相应方向的0°,90°,180°和270°角度。为了便于自我监督的预训练,我们设计了四个代理任务。重建和旋转任务分别衡量模型在每个输入上的性能,而对比和相互学习任务使模型能够跨多个视图集成信息。
重建任务比较未掩码输入 X 与重建图像
y
r
e
c
y^{rec}
yrec 之间的差异。跟随[11],我们采用均方误差(Mean-SquareError, MSE)来计算重构损失:
L
r
e
c
=
(
X
−
y
r
e
c
)
2
\mathcal{L}_{rec}=(\mathcal{X}-y^{rec})^2
Lrec=(X−yrec)2
旋转任务的目的是检测被遮挡输入沿所选视角轴线的旋转角度,旋转角度可能为0°、90°、180°和270°。模型的性能使用交叉熵损失进行评估,如下式所示,其中
y
r
o
t
y^{rot}
yrot 和
y
r
y_r
yr 分别表示旋转角度和地面真值的预测概率。
L
r
o
t
=
−
∑
r
=
1
R
y
r
l
o
g
(
y
r
o
t
)
\mathcal{L}_{rot}=-\sum_{r=1}^{R}y_rlog(y^{rot})
Lrot=−r=1∑Ryrlog(yrot)对比学习任务旨在通过比较多个视图的高级特征来评估模型在表示输入数据方面的有效性。我们的工作假设是,尽管从不同的角度来看,同一样本的表征在局部层面上可能有所不同,但在全局层面上它们应该是一致的。为了计算对比损失,我们使用余弦相似度
s
i
m
(
−
)
sim(-)
sim(−),其中
y
i
c
o
n
y_i^{con}
yicon 和
y
j
c
o
n
y_j^{con}
yjcon 代表对比对这里的对比对是不同视图的统一局部?
,t是温度常数,
1
1
1 是指示函数。log 外面是不是应该关于 j 求和
L
c
o
n
=
−
l
o
g
e
x
p
(
s
i
m
(
y
i
c
o
n
,
y
j
c
o
n
)
/
t
)
∑
k
2
N
1
k
≠
i
e
x
p
(
s
i
m
(
y
i
c
o
n
,
y
k
c
o
n
)
/
t
)
\mathcal{L}_{con}=-log\frac{exp(sim(y_i^{con},y_j^{con})/t)}{\sum_k^{2N}1_{k\neq i}exp(sim(y_i^{con},y_k^{con})/t)}
Lcon=−log∑k2N1k=iexp(sim(yicon,ykcon)/t)exp(sim(yicon,yjcon)/t)
互学习任务评估来自不同视图的重建结果的一致性,使模型能够从多视图输入中学习对齐的信息。重建结果被转换成统一的视角,并用于计算互损失
L
m
u
l
\mathcal{L}_{mul}
Lmul,该
L
m
u
l
\mathcal{L}_{mul}
Lmul 与重建任务一样,使用 MSE 损失。这里,
y
i
r
e
c
y^{rec}_i
yirec 和
y
j
r
e
c
y^{rec}_j
yjrec 分别表示视图
i
i
i 和
j
j
j 的预测重建。
L
m
u
l
=
(
y
i
r
e
c
−
y
j
r
e
c
)
2
\mathcal{L}_{mul}=(y_i^{rec}-y_j^{rec})^2
Lmul=(yirec−yjrec)2
总的预训练损失如下式所示。权系数
α
1
\alpha_1
α1,在我们的实验中,
α
2
\alpha_2
α2,
α
3
\alpha_3
α3 和
α
4
\alpha_4
α4 设为相等(
α
1
=
α
2
=
α
3
=
α
4
=
1
\alpha_1=\alpha_2=\alpha_3=\alpha_4=1
α1=α2=α3=α4=1)。
L
p
r
e
=
α
1
L
r
e
c
+
α
2
L
r
o
t
+
α
3
L
c
o
n
+
α
4
L
m
u
l
\mathcal{L}_{pre}=\alpha_1\mathcal{L}_{rec}+\alpha_2\mathcal{L}_{rot}+\alpha_3\mathcal{L}_{con}+\alpha_4\mathcal{L}_{mul}
Lpre=α1Lrec+α2Lrot+α3Lcon+α4Lmul
微调
Cross-view解码器。交叉视图解码器的结构如图1所示,包括用于跳过连接的卷积块、用于上采样的上行块和用于视图交互的交叉视图注意块。卷积块在不同的层上操作,通过执行卷积来重塑来自掩码多视图编码器的各个级别的潜在表示,使它们符合相应解码器层的特征大小(
H
2
,
W
2
,
D
2
,
i
=
0
,
1
,
2
,
3
,
4
,
5
\frac{H}{2},\frac{W}{2},\frac{D}{2},i = 0,1,2,3,4,5
2H,2W,2D,i=0,1,2,3,4,5)。在u型结构的底部,交叉视图注意力模块集成了来自两个视图的信息。假定此级别的表示包含类似的语义。横视注意机制的细节如图1和下式所示。式中,
f
i
、
f
j
f_i、f_j
fi、fj为不同视图的表示,
Q
i
、
K
i
、
V
i
Q_i、K_i、V_i
Qi、Ki、Vi分别为
f
i
f_i
fi 的查询矩阵、键矩阵、值矩阵这三个矩阵是?
。
C
r
o
s
s
A
t
t
e
n
t
i
o
n
(
f
i
,
f
j
)
=
[
S
o
f
t
m
a
x
(
Q
i
K
j
⊤
d
)
V
j
,
S
o
f
t
m
a
x
(
Q
j
K
i
⊤
d
)
V
i
]
Cross Attention(f_i,f_j)=[Softmax(\frac{Q_iK_j^{\top}}{\sqrt{d}})V_j,Softmax(\frac{Q_jK_i^{\top}}{\sqrt{d}})V_i]
CrossAttention(fi,fj)=[Softmax(dQiKj⊤)Vj,Softmax(dQjKi⊤)Vi]
多视图一致性损失。我们假设在相同体积的不同视图中应该实现一致的分割结果。为了量化多视图结果的一致性,我们引入了一致性损失
L
M
C
\mathcal{L}_{MC}
LMC,在微调阶段使用KL散度计算,如之前在相互学习方面的工作[29]。KL散度的优点是它不需要类标签,并且在微调阶段显示出更健壮。我们评估了消融研究中不同相互损失函数的有效性(见补充)。KL散度计算如式7所示:
L
M
C
=
D
K
L
(
V
i
∣
∣
V
j
)
=
∑
m
=
1
N
V
i
(
x
m
)
⋅
l
o
g
V
i
(
x
m
)
V
j
(
x
m
)
\mathcal{L}_{MC}=D_{KL}(V_i||V_j)=\sum_{m=1}^NV_i(x_m)\cdot log\frac{V_i(x_m)}{V_j(x_m)}
LMC=DKL(Vi∣∣Vj)=m=1∑NVi(xm)⋅logVj(xm)Vi(xm)
其中
V
i
(
x
m
)
,
V
j
(
x
m
)
V_i(x_m),V_j(x_m)
Vi(xm),Vj(xm) 表示第m个体素的不同视图预测。N表示案例
x
x
x 的体素数,
V
i
(
x
)
V_i(x)
Vi(x) 和
V
j
(
x
)
V_j(x)
Vj(x) 表示案例
x
x
x 的不同视图预测。我们根据[24]使用结合Dice Loss和Cross Entropy Loss的
L
D
i
c
e
C
E
\mathcal{L}_{DiceCE}
LDiceCE来衡量分割性能。
L
D
i
c
e
C
E
=
1
−
∑
m
=
1
N
(
2
∣
p
m
∩
y
m
∣
N
(
∣
p
m
∣
+
∣
y
m
∣
)
+
y
m
l
o
g
(
p
m
)
N
)
\mathcal{L}_{DiceCE}=1-\sum_{m=1}^N(\frac{2|p_m\cap y_m|}{N(|p_m|+|y_m|)}+\frac{y_mlog(p_m)}{N})
LDiceCE=1−m=1∑N(N(∣pm∣+∣ym∣)2∣pm∩ym∣+Nymlog(pm))
其中
p
m
,
y
i
p_m, y_i
pm,yi分别表示第
m
m
m 个体素的预测真值标签和ground真值标签,N为体素总数。我们在微调阶段使用
L
f
i
n
\mathcal{L}_{fin}
Lfin,如下式所示,并为不同的损失函数添加权重系数
β
D
i
c
e
C
E
β_{DiceCE}
βDiceCE 和
β
m
c
\beta_{mc}
βmc,均设置为默认值1。
L
f
i
n
=
β
D
i
c
e
C
E
L
D
i
c
e
C
E
+
β
M
C
L
M
C
\mathcal{L}_{fin}=\beta_{DiceCE}\mathcal{L}_{DiceCE}+\beta_{MC}\mathcal{L}_{MC}
Lfin=βDiceCELDiceCE+βMCLMC
SwiMM 半监督学习
如前所述,SwinMM 的多视图特性可以大大提高其最终输出的可靠性和准确性,同时最大限度地减少对大型,高质量标记医疗数据集的需求,使其成为半监督学习的有希望的候选者。在本研究中,我们提出了一个简单的SwinMM变体来处理半监督。如图3所示,我们利用来自不同视图的不同预测来处理未标记的数据,并为训练过程生成聚合的伪标签。与单视图模型相比,SwinMM的多视图方案可以通过整合来自不同视图的更全面的信息来减轻预测的不确定性,而集成操作可以减轻个体偏差这种多视图的设置使得网络不需要标签,因此适用于半监督
。
实验
数据集和评估。我们的预训练数据集包括来自8个公共数据集的5833卷:
A
b
d
o
m
e
n
t
C
T
−
1
K
[
19
]
,
B
T
C
V
[
13
]
,
M
S
D
[
1
]
,
T
C
I
A
C
o
v
i
d
19
[
9
]
,
W
O
R
D
[
18
]
,
T
C
I
A
−
C
o
l
o
n
[
14
]
,
L
i
D
C
[
2
]
和
H
N
S
C
C
[
8
]
AbdomentCT-1K[19], BTCV[13], MSD[1], TCIACovid19[9], WORD[18],TCIA-Colon[14], LiDC[2]和HNSCC[8]
AbdomentCT−1K[19],BTCV[13],MSD[1],TCIACovid19[9],WORD[18],TCIA−Colon[14],LiDC[2]和HNSCC[8]。我们选择两个流行的数据集
W
O
R
D
(
T
h
e
W
h
o
l
e
a
b
d
o
m
i
n
a
l
O
R
g
a
n
D
a
t
a
s
e
t
)
WORD (The Whole abdominal ORgan Dataset)
WORD(TheWholeabdominalORganDataset) 和
A
C
D
C
[
3
]
(
A
u
t
o
m
a
t
e
d
C
a
r
d
i
a
c
D
i
a
g
n
o
s
i
s
C
h
a
l
l
e
n
g
e
)
ACDC [3] (Automated Cardiac Diagnosis Challenge)
ACDC[3](AutomatedCardiacDiagnosisChallenge) 来测试下游分割性能。我们的分割结果的准确性是用两个常用的指标来评估的:Dice系数和Hausdorff距离(HD)Hausdorff距离衡量两个集合之间的距离,计算一个集合到另一个集合最近点距离的最大值
。
结果
单视图vs多视图。为了评估我们提出的多视图自监督预训练管道的有效性,我们将其与WORD[18]数据集上最先进的自监督学习方法SwinUNETR[24]进行了比较。具体来说,比较了两种基于swinunet的方法:使用固定的单视图(轴向,矢状和冠状)和使用多视图的集成预测(表示为swinunet - fuse)。我们的结果如表3所示,表明我们的SwinMM超过了包括SwinUNETRFuse在内的所有其他方法,突出了我们独特的多视图设计的优势。此外,通过结合多视图集成操作,SwinMM可以有效地减少硬标签中的异常值,并产生更精确的输出,特别是在处理更难的情况(如较小的器官)时。补充材料提供了2D/3D分割结果的定性比较。
消融实验
为了公平地评估多视图设计的好处,分别研究了在预训练阶段、微调阶段以及两个阶段的影响。此外,还分析了每个预训练损失函数的作用。
预训练损失函数。通过代理任务实现多视图预训练。每个任务的作用可以通过去掉其他损失函数来揭示。为了更便宜的计算,在这些实验中,我们只在来自5个数据集(Abdomentct - 1k、BTCV、MSD、tcia - covid - 19和WORD)的2639个体素上预训练我们的模型,并且我们在测试期间应用了50%的重叠窗口比。如表4所示,我们提出的相互损失比原始的SwinUNETR设置带来了Dice的显著改进(大约1%)。当组合所有代理任务时,我们的SwinMM达到了最佳性能。
数据的效率。在各种半监督设置下对数据效率进行了评估。最初,使用来自WORD数据集的一定比例的监督数据从头开始训练基本模型,训练时间为100个epoch。然后是基础模型用无监督数据完成剩余的训练过程。监督数据的比例(用标签率表示)在10%到100%之间变化。表5显示,SwinMM始终比SwinUNETR获得更高的Dice(%),并且在监督数据较少的情况下,其优越性更加显著。
结论
本文介绍了用于医学图像分析的自监督多视图管道swimm。SwinMM在预训练阶段集成了一个掩模多视图编码器,在微调阶段集成了一个跨视图解码器,实现了多视图信息的无缝集成,从而提高了模型精度和数据效率。值得注意的是,它引入了一个新的代理任务,采用相互学习范式,从3D医疗数据中提取隐藏的多视图信息。该方法实现了与现有方法相比具有竞争力的分割性能和更高的数据效率,并强调了自监督学习领域中多视图学习的潜力和有效性。