太强了!在图像去雨、去模糊、去噪等任务上拿下多项SOTA!性能优于MSPFN、PreNet等网络,代码于5个小时前开源!
看看MPRNet在各大数据集刷榜的存在!
作者单位:IIAI, MBZUAI, 谷歌等
代码:https://github.com/swz30/MPRNet
论文:https://arxiv.org/pdf/2102.02808.pdf
图像恢复任务要求在恢复图像时在空间细节和高级上下文信息之间达成复杂的平衡。
在本文中,我们提出了一种新颖的协同设计,可以最佳地平衡这些相互竞争的目标。我们的主要建议是一个多阶段体系结构,该体系结构逐步学习降级输入的恢复特征,从而将整个恢复过程分解为更易于管理的步骤。
具体来说,我们的模型首先使用编码器-解码器体系结构学习上下文相关的特征,然后将它们与保留局部信息的高分辨率分支相结合。
图3:(a)编码器-解码器子网。(b)图示我们的ORSNet子网中的原始分辨率块(ORB),每个ORB包含多个通道注意块,GAP代表全球平均池[47]。(c)阶段1和阶段2之间的跨阶段特征融合。(d)阶段2和最后阶段之间的CSFF
在每个阶段,我们都会介绍一种新颖的每像素自适应设计,该设计利用in-situ监督的注意力来重新加权局部特征。这种多阶段体系结构的关键要素是不同阶段之间的信息交换。
图4:监督注意模型
为此,我们提出了一种两方面的方法,即不仅从早期到后期顺序交换信息,而且在特征处理模块之间也存在横向连接,以避免信息丢失。由此产生的紧密互连的多级体系结构称为MPRNet,可在包括图像去雨,去模糊和去噪在内的一系列任务的十个数据集上获得强劲的性能提升。
图2:提议的渐进式图像恢复的多阶段体系结构。较早的阶段使用编码器/解码器来提取多尺度上下文相关的特征,而最后阶段则以原始图像分辨率运行以生成空间精确的输出。在每两个阶段之间添加一个监督注意力模块,该模块学习在将一个阶段的特征传递到下一阶段之前对其进行细化。虚线的粉红色箭头表示跨阶段特征融合机制。
主要贡献
这项工作的主要贡献是:
- 一种新颖的多阶段方法,能够生成内容丰富且空间精确的输出。由于其多阶段的性质,我们的框架将具有挑战性的图像还原任务分解为子任务,以逐步还原退化的图像。
- 一个有效的监督注意力模块,该模块在进一步传播传入特征之前,会在每个阶段充分利用还原的图像来完善传入特征。
- 跨阶段聚合多尺度特征的策略。
- 我们通过在十个合成的和真实的数据集上设置最先进的最新数据集来证明MPRNet的有效性,这些数据集可用于各种恢复任务,包括图像去油,去模糊和去噪,同时保持较低的复杂度(见图1)。此外,我们提供了详细的烧蚀,定性结果和泛化测试。
Quick Run
要在您自己的图像上测试去模糊、去油、去噪的预训练模型,请运行
python demo.py --task Task_Name --input_dir path_to_images --result_dir save_images_here
这是执行去模糊的示例:
python demo.py --task Deblurring --input_dir ./samples/input/ --result_dir ./samples/output/
实验结果
例如,与最新技术相比,在Rain100L,GoPro和DND数据集上,我们分别获得4 dB,0.81 dB和0.21 dB的PSNR增益。
去雨的性能结果
去模糊的性能结果
表3:去模糊结果。我们的方法仅在GoPro数据集[51]上训练,并直接应用于HIDE数据集[67]。
表4:在两种不同的设置下对RealBlur数据集进行去模糊比较:1) 将我们经过GoPro训练的模型直接应用于RealBlur集(以评估对真实图像的泛化),2) 对RealBlur数据的培训和测试,其中方法用符号‡表示。其他评估方法的PSNR / SSIM分数取自RealBlur基准测试[62]。
去噪的性能结果
表5:不在SIDD [1]和DND [58]数据集上进行去噪比较。∗表示使用其他训练数据的方法。而我们的MPRNet仅在SIDD图像上进行训练,而在DND上直接进行测试。
图7:图像降噪比较。第一个示例来自DND [58],其他示例来自SIDD [1]。提出的MPRNet可以更好地保留去噪图像中的精细纹理和结构图案.