基于对抗性混频的鲁棒图像去噪
项目链接:https://github.com/dhryougit/AFM
Abstract
基于深度神经网络的图像去噪方法经常与训练数据中存在的特定噪声分布的过拟合作对抗。这一挑战存在于现有的现实世界去噪网络中,这些网络使用有限的真实噪声分布谱进行训练,因此对分布外的真实噪声类型表现出较差的鲁棒性。为了缓解这个问题,我们开发了一种新的训练框架,称为对抗性混频(AFM)。AFM利用频域的混合产生具有独特和具有挑战性的噪声特征的噪声图像,同时保留真实世界噪声的属性。随后,将这些噪声图像纳入训练pipeline中,增强去噪网络对噪声分布变化的鲁棒性。在广泛的真实噪声基准上进行的大量实验和分析表明,使用我们提出的框架训练的去噪网络在对看不见的噪声分布的鲁棒性方面表现出显着的改进。
1. Introduction
基于深度神经网络的图像去噪[5,6,22,26 - 28,30,31]由于问题表述的简单性和新数据集的构建,取得了前所未有的成功。传统上,去噪网络使用合成噪声模型进行训练,例如高斯或泊松噪声,这些模型被人为地添加到图像中以用于训练和评估目的。然而,现实世界的噪声受图像信号处理(ISP)pipeline中的各种因素(如去马赛克和伽马校正)的影响,表现出明显的信号依赖性,并且通常遵循与合成对应物不同的分布。当将去噪模型应用于真实的噪声图像时,合成噪声和真实噪声分布之间的差异提出了重要的泛化问题。
虽然人们一直致力于在现实世界中创建具有干净和噪声图像对的数据集,但收集这些数据集是一项重大挑战。因此,各种自监督方法[10,12,18]已经成为图像去噪的有前途的解决方案,旨在减少对成对的无噪声图像数据集的依赖。然而,与监督学习方法相比,这些方法在性能上往往不足。因此,在实际场景中,对真实噪声数据集的监督仍然是有利的。
然而,在真实噪声数据集上进行监督训练的图像去噪器有其自身的缺陷。最值得注意的是,我们观察到,即使是最先进的去噪器也难以泛化到噪声分布的变化,这些变化是由不同的相机传感器类型、拍摄环境和ISP过程等因素引起的。我们在图1中说明了这种现象,并在图2中提供了频率分析。虽然缺乏真实的噪声泛化对去噪模型的广泛使用产生了不利影响,但在图像去噪研究领域,这是一个相对未被探索的问题。
在这项工作中,我们提出了对抗混频(AFM),这是一种模型不可知的训练框架,可以提高去噪网络对实际噪声分布变化的泛化能力。为此,AFM构建了具有独特且难以去噪的噪声的图像,同时保持了真实世界噪声的特性。然后,通过将这些噪声图像整合到训练pipeline中,去噪网络对实际噪声分布的变化具有鲁棒性。更具体地说,AFM的工作原理是根据频域的混合掩码将有噪声和去噪的图像混合在一起。这个混合掩码是由一个单独的轻量级神经网络生成的,它产生了我们方法的两个变体:AFM-E和AFM-B。AFM-E以元素方式生成混频掩码,而AFM-B通过为每个频带分配一个混频值来生成掩码。此外,AFM使用对抗性损失进行训练,并可以生成对抗性混合掩码。
总体而言,我们的贡献总结如下:
- 我们提出了AFM,这是一种模型不可知的训练框架,通过在频域中的对抗性混合生成新的噪声图像,提高了去噪网络对现实世界分布外噪声图像的泛化和鲁棒性。
- 我们设计了两个对抗性掩码生成网络,AFM-E和AFM-B。AFM-E以元素方式生成混合掩码,而AFM-B为单个波段生成混合掩码。
- 我们在多个真实世界的分布外图像去噪基准上证明了AFM的有效性。我们的方法在各种去噪架构甚至最先进的去噪模型上始终优于普通训练。
2. Related Works
本节回顾了现有的基于有监督和无监督的深度神经网络的图像去噪模型,并讨论了最近在鲁棒图像去噪方面的努力。
2.1. 有监督图像去噪
近年来,有监督图像去噪领域取得了重大进展,其中有噪声和干净的图像配对可用于训练。最初的突破主要是由基于CNN的模型推动的,其中去噪卷积神经网络(Denoising Convolutional Neural Network, DnCNN)[30]为该领域的进一步创新奠定了基础。在此基础上,基于U-Net的架构[4,5,26,27]已经成为突出的模型,利用跳过连接来有效地结合局部和全局上下文信息。此外,基于Transformer的模型的引入[22,28]标志着去噪策略的范式转变。配备了注意力机制,基于Transformer的模型擅长识别复杂的依赖关系,有效地减少了过程中的噪声失真。这些模型在降低噪声方面显示出卓越的技能,大大减轻了它们使用成对数据集训练的特定噪声分布的影响。然而,他们在有效管理与他们所训练的噪声分布不同的噪声分布方面遇到了挑战,这对他们在现实场景中的应用构成了障碍。
2.2. 自监督图像去噪
从现实世界中收集干净和有噪声的对数据是成本很高的。为了缓解这些问题,人们积极研究探索无监督和自监督学习技术,如N2N[12]、N2S[2]、N2V[10]和R2R[18]。这些方法证明了使用仅包含噪声对或单个噪声图像的数据集训练去噪网络的可行性,而不需要干净的噪声对数据集。此外,DIP[21]、Self2Self[19]和Neighbor2Neighbor[7]引入了新的方法来训练去噪网络,即使在没有训练数据的情况下,也可以使用单个带噪声的图像来产生干净的对应图像。然而,由于网络需要对每个不同的图像进行重新训练,因此这些方法的效率低下且耗费大量时间。重要的是,与监督方法相比,它们的性能往往不足。
2.3. 去噪泛化
现有的去噪模型存在泛化问题,对训练噪声分布存在过拟合问题。只有少数作品解决了这一挑战[3,15]。Mohan等[15]观察到噪声过拟合是由偏置项引起的,并从网络中去除所有的偏置。然而,这种方法的目的是增强在不同程度的噪声中的鲁棒性,而不是噪声类型。Chen等人[3]在训练过程中引入输入像素和注意特征掩码,专注于重建干净的图像本身,而不是去噪,但其在真实噪声上的表现仍然不足。
3. Methods
我们的目标是训练一个去噪网络,在各种看不见的真实噪声的例子中表现出强大的泛化性能。为了实现这一目标,我们提出了对抗混频(AFM),这是一种基于具有真实噪声分布的频域增强图像的新型训练框架。
3.1. 背景
去噪网络的目的是生成干净的图像,而不考虑输入图像上的噪声模式。这种步骤表述为:
D θ ( x + n ) = x , (1) \mathcal{D}_\theta(\boldsymbol{x}+\boldsymbol{n})=\boldsymbol{x}, \tag{1} Dθ(x+n)=x,(1)
式中, D θ \mathcal{D}_θ Dθ表示用 θ θ θ参数化的去噪网络, ( x , n ) (x, n) (x,n)表示干净图像及其噪声的一对。
对于监督学习,一种流行的方法是使用真实的噪声数据集(如SIDD[1])训练去噪模型,该数据集由干净图像 ( x ) (x) (x)和噪声图像 ( x + n ) (x + n) (x+n)对组成。问题的目标是通过最小化以下损失来优化模型参数 θ θ θ:
L = ∥ D θ ( x + n ) − x ∥ . (2) \mathcal{L}=\|\mathcal{D}_\theta(\boldsymbol{x}+\boldsymbol{n})-\boldsymbol{x}\|. \tag{2} L=∥Dθ(x+n)−x∥.(2)
然而,传统的方法容易学习干净图像和噪声图像之间的固定映射,这往往导致较差的泛化。在实际场景中,由于相机传感器和图像信号处理(ISP)pipeline的不同特性,存在很大范围的噪声变化。因此,真实的噪声去噪模型必须能够推广到这种噪声分布的变化。
3.2. 对抗性混频
图3(a)说明了所提出的AFM方法的概述,该方法旨在提高对看不见的真实噪声图像的泛化。给定一个带噪声的输入图像 I ∈ R C × H × W I∈\mathbb{R}^{C×H×W} I∈RC×H×W,我们首先使用去噪网络 D θ \mathcal{D}_θ Dθ预测去噪图像 I I I,该网络由公式(3)给出:
I ^ = D θ ( I ) . (3) \hat{I}=\mathcal{D}_\theta(I). \tag{3} I^=Dθ(I).(3)
然后,我们将原始输入 I I I与频域预测 I ^ \hat{I} I^混合如下:
I h a r d = F − 1 ( F ( I ) ⊙ m + F ( I ^ ) ⊙ ( 1 − m ) ) , (4) I_{\mathrm{hard}}=\mathcal{F}^{-1}\left(\mathcal{F}(I)\odot\boldsymbol{m}+\mathcal{F}(\hat{I})\odot(\mathbf{1}-\boldsymbol{m})\right), \tag{4} Ihard=F−1(F(I)⊙m+F(I^)⊙(1−m)),(4)
其中 F \mathcal{F} F是快速傅里叶变换(FFT), ⊙ \odot ⊙是元素乘法算子, m ∈ [ 0 , 1 ] 1 × H × W m∈[0,1]^{1×H×W} m∈[0,1]1×H×W是频域中的任意掩码, I h a r d I_{hard} Ihard是得到的混合图像。由于掩码 m m m以[0,1]为界,公式(4)的右侧表示频域中 I I I和 I ^ \hat{I} I^之间的元素插值。给定 I I I和 I ^ \hat{I} I^,这里的目标是生成带有一些噪声的新图像,其特征是与原始噪声图像的分布不同,同时也类似于现实场景中遇到的噪声。我们在公式(4)中的设计选择—特别是频率混频的使用—直接受到这个目标的启发。
I I I和 I ^ \hat{I} I^上的傅里叶变换将每个图像映射到频域,与空间域相比,在频域操作图像的噪声分布相对容易。图2显示了支持这一主张的直观可视化,其中噪声和干净信号的区域在频域中可以清晰区分。此外,频率分量的变化对图像起全局作用,从而最大限度地降低了对底层图像内容进行不必要操作的风险。从本质上讲,结果 I h a r d I_{hard} Ihard停留在真实噪声图像的流形中,因为它是 I I