【图像去雨】MPRNet:性能炸裂!MPRNet:多阶段渐进式图像恢复(图像去雨、去模糊、去噪)
原文:https://arxiv.org/pdf/2102.02808.pdf
源码:http://github.com/swz30/MPRNet
在进行图像去雨的过程中需要恢复图像在空间细节和上下文信息之间达成的平滑,MPRNet进行了新的设计,通过一个多阶段的体系结构,该体系结构逐步学习降级输入的恢复特征,从而将整个恢复过程分解为更易于管理的步骤。具体来说,模型首先使用编码器-解码器体系结构学习上下文相关的特征,然后将它们与保留局部信息的高分辨率分支相结合。
网络结构
渐进式图像恢复的多阶段体系结构。Stage 1和Stage 2阶段使用编码器/解码器来提取多尺度上下文相关的特征,而Stage 3则以原始图像分辨率运行以生成空间精确的输出。该体系架构中包含模块有:
(1)Conv:卷积模块,用于图像特征的提取;
(2)Channel Attention Block(CAB):用于通道注意力的增强
(3)U-Net:编码器-解码器结构,用于深层特征提取并恢复图像分辨率;
(4)Supervised Attention Block(SAM):利用真实监督信号和局部监督预测来生成注意力图,抑制当前阶段不太重要的特征;
(5)Original Resolution(OPB):在保持输入图像原始分辨率的情况下,生成空间上丰富的高分辨率特征
此外,结构中的绿色实线表示通道维度的特征图融合,用于融合不同层的特征信息;虚线的粉红色箭头表示跨阶段特征融合机制。
核心模块如下:
编码器-解码器结构
(1)编码器
编码器负责对输入图像进行特征提取,将图像从高分辨率逐渐下采样到低分辨率,从而捕获图像的高级语义信息。编码器由多个卷积层和池化层组成。卷积层用于提取图像的特征,池化层用于减小特征图的尺寸,增加感受野。随着网络层次的加深,特征图的分辨率逐渐降低,而通道数逐渐增加,从而学习到更抽象的特征表示。在编码器阶段中使用通道注意块(CABs)提取每个尺度上的特征。
(2)解码器
解码器的任务是将编码器学习到的低分辨率特征图恢复到原始图像的分辨率,从而实现对图像的重建或分割。通过多个上采样层和卷积层增加特征图的分辨率并进一步提取和整合特征。在解码器的最后阶段使用一个双线性上采样层提高解码器中特征的空间分辨率。
(3)跳跃链接
跳跃链接用于连接编码器和解码器的对应层,将编码器中提取的低级特征直接传递到解码器的相应层,从而帮助解码器更好地恢复图像的细节。跳跃链接可以有效地解决在深度神经网络中由于信息丢失导致的图像细节恢复困难的问题。通过将低级特征与高级特征相结合,网络可以更好地学习到图像的全局和局部信息,提高图像恢复的准确性。
Original Resolution Block(OPB)
ORB 的主要作用是在保持输入图像原始分辨率的情况下,生成空间上丰富的高分辨率特征。每个 ORB 包含多个通道注意力块(CAB),CAB 用于提取特征并对特征进行重新校准,以突出重要特征并抑制不重要的特征。具体来说,ORB 的工作原理如下:
(1)接收输入特征:ORB 接收来自前一层的特征作为输入。
(2)CAB 处理:输入特征经过多个 CAB 的处理,CAB 通过对特征进行通道维度上的注意力计算,重新调整特征的重要性分布。
(3)生成高分辨率特征:经过 CAB 处理后的特征被组合和处理,以生成具有丰富空间信息的高分辨率特征。
(4)输出特征:ORB 将生成的高分辨率特征输出给下一层或后续的处理步骤。
通过使用多个 ORB能够有效地保留图像的精细细节,从而在图像恢复等任务中发挥重要作用。这种在原始分辨率上操作的设计有助于提高网络对空间细节的捕捉能力,与编码器 -解码器子网相结合,可以更好地实现图像的恢复和重建。
跨阶段特征融合(Cross-stage Feature Fusion,CSFF)
引入了两个编码器-解码器之间的CSFF模块,图4是结构中CSFF模块,即图1中的粉色虚线模块。左图是Stage 1和Stage 2之间的跨阶段特征融合;右图是Stage 2和Stage 3之间的跨阶段特征融合。来自一个阶段的特性首先通过1×1卷积进行细化,然后将它们传播到下一个阶段进行聚合。CSFF有几个优点。首先,由于在编码解码器中重复使用上采样和下采样操作,它使网络不那么容易出现信息损失。第二,一个阶段的多尺度特征有助于丰富下一阶段的特征。第三,网络优化过程变得更加稳定,简化了信息流,从而允许在整个架构中添加几个阶段。
Supervised Attention Module(SAM)
多阶段网络中一个阶段的预测结果直接输入到下一阶段,而SAM在两个阶段之间引入了一个有监督的注意模块,能够显著提升模型性能。如图5所示,Fin表示输入特征图,将特征图进行卷积操作,获取残差特征Rsx,与原始图像进行加权求和,并计算相应损失,这样充分利用了标签的监督,最后再结合输入特征信息,获取输出图像Fout。SAM为图像恢复的每个阶段提供了地面真实监督信号,这对于渐进式的图像恢复非常有用;其次通过生成注意力掩码,SAM 能够抑制当前阶段不太重要的特征,只允许有用的特征传播到下一个阶段,从而提高了信息的有效性和准确性。
部分测试效果