ECCV 2022 | 南大&东南大学新作RAM-DSIR:通过随机振幅混合和特定领域的图像重建实现可推广的医学图像分割
-
论文标题:Generalizable Medical Image Segmentation via Random Amplitude Mixup and Domain-Specific Image Restoration
论文地址:https://doi.org/10.48550/arXiv.2208.03901
论文代码:https://github.com/zzzqzhou/RAM-DSIR
摘要
- 现存问题: 对于医学图像分析,由于不同数据采集策略之间的差异,在一个或多个域上训练的分割模型缺乏对未知域的泛化能力。
- 解决方法: 作者认为分割性能的下降主要归因于对源域的过度拟合和域移位。为此,提出了一种新的可推广的医学图像分割方法。
- 实验结果: 在医学图像的两个公共通用分割基准上证明了方法的性能,这验证了算法可以达到最先进的性能。
算法
将K个源域的集合表示为Ds={(xik,yik)i=1Nk}k=1K,其中xik是第k个源域的第i幅图像;yik是分割标签;Nk是指第k个源域中的样本数量。
目标是在Ds上学习一个可推广的医学图像分割模型Fθ。该模型Fθ有望在未见过的目标域Dt={xi}i=1Nt上表现出令人满意的泛化性能,其中xi代表目标域中的第i幅图像,Nt是目标域中的图像样本数。
提出的方法包含一个编码器-解码器分割模型和一个辅助的特定领域图像恢复(DSIR)解码器。如上图b。
在训练pipeline前面,引入了一个名为随机振幅混合(RAM)的数据增强和损坏模块,如上图a。
工作流程包含以下步骤。
- 首先,在RAM模块中,对两个共享不同域标签的源域图像进行傅里叶变换,以获得它们的频率空间信号;
- 然后,纳入它们的低频信号,并利用反傅里叶变换来生成新的图像。
- 其次,在DSIR模块中,分割模型的编码器获得由RAM生成的图像的低层次特征。训练一个具有特定领域批量归一化的解码器,以便从低级特征中恢复特定源域的原始图像。
- 最后,编码器-解码器的分割模型是通过源域图像和增强图像的分割损失来训练的;同时,源域图像和增强图像的输出之间采用了一致性损失,以帮助分割模型更好地抵御域的转移。接下来我们将详细讨论所有这些组成部分
Random Amplitude Mixup
为了解决源域和目标域之间领域差异的限制,一个合理的想法是在源域上应用数据增量,使源域数据多样化。
在这种情况下,可以对模型进行正则化处理,减轻对源域的过度拟合。在大量的数据增强方法中,Mixup已被广泛用于图像识别任务中。图像水平混合(IM)将两个不同的图像从训练数据集中纳入。然而,IM也会干扰图像的语义信息,这可能对语义分割任务产生负面影响。
因此论文建议利用频率空间中源域的固有信息,并将分布信息(即风格)纳入不同图像的振幅谱中。
具体来说,从源域k中随机抽取一个样本图像xik∈RH×W×C(C代表图像通道数;H和W是图像的高度和宽度),然后进行傅里叶变换F,得到图像xik的频率空间信号,可以写成:
经过傅里叶变换,可以将频率信号F(xik)分解为振幅谱Aik∈RH×W×C和相位图Pik∈RH×W×C。
其中,振幅频谱包含低层次的统计数据(如风格),而相位图像包括原始图像的高层次(如对象)语义。 将来自多个源域的不同图像的振幅谱纳入其中。为此,从源域n随机选择另一个样本图像xjn(n ̸= k),并对其进行傅里叶变换。这样,就得到图像xjn的另一个振幅Ajn。为了将低频分量纳入振幅Aki和Anj中,引入了一个二进制掩码M,它可以控制要纳入的振幅频谱中的低频分量的尺度。之后,我们将图像xki和图像xnj的振幅信息纳入。
最后,可以通过反傅里叶变换F−1将合并后的振幅Ai,λn→k转换为新的风格化图像,如下所示
其中,生成的图像xi,λn→k包含xik的语义信息,其低层次信息(例如,风格)是xki和xnj的低层次信息的混合物。
这进一步表明RAM可以增加源域的多样性并缩小域的差异。图2中展示了Fundus数据集中图像特征的t-SNE可视化。图2(a)显示了Fundus数据集中不同领域的原始分布信息。从可视化中,可以观察到来自不同领域的图像特征是明显分开的。这就导致了这样一个问题:在原始源域上训练模型,使得模型很容易对特定的源域过度拟合,这可能会降低目标域的泛化性能。然而,在图2(b)中,发现,通过在原始源域上应用RAM,可以大大缩小域的差距,显示出域的不变性表示。不同领域的分布更加紧凑和多样化。
为了对目标域的图像进行分割,一种直接的方法是通过直接将多源域图像送入模型,以统一的方式训练一个分割模型。把训练这样一个分割模型命名为 “DeepAll”。尽管 "DeepAll "方法在多源域上可能有很好的泛化性能,但它在目标域图像上可能无法保持令人满意的分割性能。在多源域上训练一个分割模型,并没有引入监督来对抗域的偏移。此外,原始的多源域图像在特征分布上缺乏足够的多样性,这可能会导致对特定源域的过度拟合。
设计了一个语义一致性训练策略来解决 "DeepAll "方法的问题。具体来说,引入一个编码器-解码器结构作为分割模型。编码器E将从图像中提取低层次的语义特征,而分割解码器Dseg则用于预测分割掩码。将源域图像xki上的分割模型的前向传播表述为
其中,ˆyn→ki,λ代表预测值。然后,利用统一的交叉熵(CE)损失和骰子损失作为我们的分割损失来优化模型。 原始源域k上的CE和Dice损失被表述为
所以,xik和xi,λn→k的分割损失可以写成
为了对抗领域转移,提出了一种新的语义一致性损失。具体来说,将生成的图像xi,λn→k视为xik的风格增强。打算强迫分割模型从xik和xi,λn→k中预测一致的分割结果。这样,分割模型就可以对领域的转移不那么敏感。
设计了一个损失项来最小化软预测ˆyki和ˆyn→ki,λ之间的KullbackLeibler(KL)分歧。语义一致性损失如下
Domain-Specific Image Restoration
为了进一步规范分割模型,减少源域的过拟合,提出了一个自监督的辅助任务,以帮助训练一个更稳健的分割模型。具体来说,引入了一个带有特定领域批量规范化(DSBN)层的图像重建解码器。图像重建解码器被用来从分割编码器E从RAM图像xi,λn→k中提取的低层次特征中恢复图像。
为了更好地恢复不同源域的图像,在图像重建解码器中加入了DSBN。让图像修复解码器表示为Drec = {D1rec,D2rec,–,DKrec},其中K代表源域的数量,Dkrec用于恢复由第k个源域图像生成的RAM图像的低级特征。Drec中的所有解码器共享相同的模型参数,但有不同的批处理规范化层。由于多源域的分布信息是相当不同的,在不同的域中使用不同的批归一化层可以更好地保留域的固有特征,以进行图像恢复。图像重建模块在源域k上的前向传播如下。
为了训练图像重建模块,采用L2距离作为恢复损失,以优化Drec和E,在第k个源域的恢复损失为:
总的来说,可以把整个框架表述为一个多任务学习范式。总的训练损失如下。
实验
Comparison with Other DG methods
Analysis of Our Method