相关链接
论文链接
代码链接
摘要
RAW数据的独特属性在低光照图像增强方面展现出巨大潜力。然而,由于现有架构在单阶段和多阶段方法中固有的局限性,性能受到了限制。在单阶段方法中,在两个不同领域之间的混合映射,即噪声到清晰和RAW到sRGB,由于领域模糊性,误导了单阶段方法。多阶段方法仅通过每个阶段产生的图像传播信息,忽略了在有损图像级数据流中丰富的特征。本文探索了这些瓶颈的通用解决方案,并提出了一个名为DNF(Decouple and Feedback,解耦和反馈)的框架。为了缓解领域模糊性,作者解耦了领域特定的子任务,并充分利用了RAW和sRGB领域的独特属性。通过反馈机制跨阶段进行特征传播,避免了由于图像级数据流造成信息丢失的问题。
引言
在低光照场景下的成像引起了越来越多的关注,尤其是随着智能手机夜景模式和监控系统的普及。然而,由于极低的信噪比,低光照图像增强(LLIE)是一个具有挑战性的任务。与sRGB数据相比,RAW数据(RAW data with the unprocessed signal)具有三个独特的优点,使其在LLIE中特别有用:1)信号与RAW域中的光子计数呈线性相关;2)RAW图像上的噪声分布在图像信号处理(ISP)流程之前是可处理的;3)RAW格式的更高位深记录了更多可区分的低强度信号。
为了充分利用RAW图像在LLIE中的潜力,需要一个通用的pipeline。具体来说,神经网络应该利用不同领域的优势,而不是被领域模糊性所困惑。根据RAW和sRGB领域的独特属性,将增强分解为领域特定的子任务是至关重要的。在探索了RAW域中的线性和可处理的噪声之后,可以执行从RAW域到sRGB域的颜色空间转换,而没有噪声干扰。此外,pipeline不能阻碍跨阶段的通信,而不是仅允许一小部分有损信息通过的图像级数据流。
基于这些原则,作者提出针对基于RAW的低光照图像增强(LLIE)量身定制的解耦和反馈(DNF)框架。增强过程被解耦为两个领域特定的子任务:RAW域中的去噪和色彩恢复到sRGB域,如图1(c)所示。
创新点
1、领域特定任务解耦扩展了在RAW和sRGB域中独特属性的利用,避免了领域模糊
2、通过去噪先验反馈实现的特征级数据流减少了误差累积,并在各个阶段聚合了互补的特征
3、与以前的SOTA方法相比,所提出的方法在只有19%参数和63%浮点运算次数(FLOPs)的情况下,取得了显著的性能提升,例如,在SID数据集的Sony子集上PSNR提高了0.97dB,在Fuji子集上PSNR提高了1.30dB
方法论
整体框架
如图2所示,作者提出的DNF框架由两个阶段组成:RAW去噪和颜色渲染,以逐步增强低光照RAW图像。给定输入图像x_raw ,在乘以预定义的放大比例后,放大后的图像X首先通过RAW域中的编码器E _raw和解码器D_raw进行去噪。然后,作者没有使用不准确的y_raw进行颜色渲染,而是将去噪特征F_dn从D_raw反馈回E_raw,进一步区分信号和噪声,并在RAW域中组合丰富的特征。最后,sRGB解码器D_rgb采用RAW域中的多尺度特征来渲染在sRGB域的最终输出Y_rgb。
具体来说,一个共享的编码器E_raw和两个解码器(D_raw和D_rgb)是专门为通过领域特定任务解耦(Domain-Specific Task Decoupling)与任务特定模块而解耦的子任务设计。引入了通道独立去噪(CID)模块,以学习RAW域中不同颜色通道中的可处理且独立的噪声分布。根据颜色空间的定义,矩阵颜色校正(MCC)模块使用全局矩阵变换来完成剩余的增强以进入sRGB域。此外,作者引入了一个去噪先验反馈机制(Denoising Prior Feedback mechanism),以避免跨阶段的误差累积。通过从RAW解码器提取的去噪特征F_dn,RAW编码器用高频信息丰富了浅层特征。此外,作者提出了一个带门控机制的门控融合模块(GFM),用于自适应地探索噪声中的细节。
Domain-Specific Task Decoupling
作者提出领域特定任务解耦(Domain-Specific Task Decoupling)来处理由噪声到清晰和RAW到sRGB的混合映射引起的领域模糊。噪声的RAW和清晰的sRGB领域之间的鸿沟使网络难以学习横跨两个领域的直接映射。因此,作者提出在噪声的RAW领域中引入中间监督,以缓解直接学习混合映射的难度。通过在干净的RAW领域的中间监督,可以:1)将增强分解为RAW去噪和颜色恢复,2)充分利用这一特性,即在RAW图像上的噪声分布是可处理的,用于去噪,3)减少在颜色恢复期间的噪声干扰,从而减少颜色偏移。
Denoising in RAW Domain
如图2所示,作者通过堆叠多个通道独立去噪(CID)块来实现RAW编码器E_raw和RAW解码器D_raw。CID块的设计基于以下两个先验知识:1)RAW格式的低光照图像受到与信号无关的噪声影响,该噪声遵循零均值分布,2)不同通道的信号本质上相关性较小,因此在RAW域中噪声分布在各个通道上趋于独立。因此,需要对几乎相同的信号(邻近像素)进行burst observations,以消除零均值噪声的干扰。同时,在去噪过程中防止通道间信息交换对于处理通道独立的噪声分布至关重要。根据上述讨论,在CID块中引入了大核深度卷积进行去噪。CID块的详细结构如图3(a)所示。具体来说,对于输入特征F_in,经过通道独立去噪块后的输出特征F_out可以表示为:
CID模块的代码如下,
from torch import nn
# CI
class DConv7(nn.Module):
def __init__(self, f_number, padding_mode='reflect') -> None:
super().__init__()
self.dconv = nn.Conv2d(f_number, f_number, kernel_size=7, padding=3, groups=f_number, padding_mode=padding_mode)
def forward(self, x):
return self.dconv(x)
# Post-CI
class MLP(nn.Module):
def __init__(self, f_number, excitation_factor=2) -> None:
super().__init__()
self.act = nn.GELU()
self.pwconv1 = nn.Conv2d(f_number, excitation_factor * f_number, kernel_size=1)
self.pwconv2 = nn.Conv2d(f_number * excitation_factor, f_number, kernel_size=1)
def forward(self, x):
x = self.pwconv1(x)
x = self.act(x)
x = self.pwconv2(x)
return x
class CID(nn.Module):
def __init__(self, f_number, padding_mode) -> None:
super().__init__()
self.channel_independent = DConv7(f_number, padding_mode)
self.channel_dependent = MLP(f_number, excitation_factor=2)
def forward(self, x):
return self.channel_dependent(self.channel_independent(x))
使用大核深度卷积的优点分析:
1、处理信号独立噪声:由于噪声与信号无关,且在整个图像中均匀分布,使用深度卷积可以更有效地处理每个通道内的噪声,而不会受到其他通道噪声的干扰。大核心的深度卷积能够覆盖更广的像素区域,有助于更好地捕捉并去除噪声
2、增强空间上的连续性:大核心的深度卷积能够在进行噪声去除的同时,保持图像中的空间连续性。这对于处理低光照图像中的噪声尤为重要,因为这些噪声往往是在图像的局部区域内更为显著
3、防止通道间信息交换:在去噪过程中,阻止通道间的信息交换是处理信号独立噪声分布的关键。深度卷积通过独立处理每个通道,确保了去噪过程中不会引入其他通道的噪声或信息,从而保持了处理的纯净性和效率
4、利用连续观察:大核心的深度卷积可以覆盖更多的相邻像素,这有助于利用连续观察来去除零均值噪声的干扰。这种方法在处理由于噪声而导致的图像质量下降时,可以有效地恢复图像的细节和清晰度
Color Correction for RAW-to-sRGB
矩阵变换通常在规范的ISP流程中使用。由于全局共享设置,如环境照明和颜色空间规范,图像的颜色主要通过通道矩阵变换来增强或转换到另一种颜色空间。遵循这一原则,作者引入了一个矩阵颜色校正(MCC,Matrixed Color Correction)块来执行全局颜色增强和局部细化,如图3(b)所示。对于sRGB解码器D_rgb,作者堆叠了多个MCC块进行颜色校正。这个块的设计受益于最近的转置自注意力(transposed self-attention)的进展。全局感受野和通道操作非常适合规范ISP中的颜色校正。给定输入源特征F_source,查询Q、键值K、值V依次通过1×1卷积、3×3深度卷积和展平操作生成。然后,通过矩阵乘法获得变换矩阵M。这个过程可以表示为:
MCC模块代码如下,
from einops import rearrange
import torch
from torch import nn
from ..utils import LayerNorm
class MCC(nn.Module):
def __init__(self, f_number, num_heads, padding_mode, bias=False) -> None:
super().__init__()
self.norm = LayerNorm(f_number, eps=1e-6, data_format='channels_first')
self.num_heads = num_heads
self.temperature = nn.Parameter(torch.ones(num_heads, 1, 1))
self.pwconv = nn.Conv2d(f_number, f_number * 3, kernel_size=1, bias=bias)
self.dwconv = nn.Conv2d(f_number * 3, f_number * 3, 3, 1, 1, bias=bias, padding_mode=padding_mode, groups=f_number * 3)
self.project_out = nn.Conv2d(f_number, f_number, kernel_size=1, bias=bias)
self.feedforward = nn.Sequential(
nn.Conv2d(f_number, f_number, 1, 1, 0, bias=bias),
nn.GELU(),
nn.Conv2d(f_number, f_number, 3, 1, 1, bias=bias, groups=f_number, padding_mode=padding_mode),
nn.GELU()
)
def forward(self, x):
attn = self.norm(x)
_, _, h, w = attn.shape
qkv = self.dwconv(self.pwconv(attn))
q, k, v = qkv.chunk(3, dim=1)
q = rearrange(q, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
k = rearrange(k, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
v = rearrange(v, 'b (head c) h w -> b head c (h w)', head=self.num_heads)
q = torch.nn.functional.normalize(q, dim=-1)
k = torch.nn.functional.normalize(k, dim=-1)
attn = (q @ k.transpose(-2, -1)) * self.temperature
attn = attn.softmax(dim=-1)
out = (attn @ v)
out = rearrange(out, 'b head c (h w) -> b (head c) h w', head=self.num_heads, h=h, w=w)
out = self.project_out(out)
out = self.feedforward(out + x)
return out
Denoising Prior Feedback
在以前的基于RAW的增强方法中,一部分高频内容被错误地识别为噪声,严重恶化了最终结果,导致细节丢失和数据流的损失。为了避免现有多阶段方法的有损图像级数据流,作者提出了具有特征级信息传播的去噪先验反馈机制(Denoising Prior Feedback mechanism)。记F_dn={F_dn_1, F_dn_2, ..., F_dn_L}为从RAW解码器D_raw提取的一组去噪特征,其中L表示阶段数。F_dn的每个元素主要包含RAW域中不同尺度的最终噪声估计信息。具体来说,这些特征使噪声更加容易区分,并作为进一步去噪的指导。通过将去噪特征集合F_dn链接到具有多个反馈连接的RAW编码器的相应阶段,编码器逐渐生成更好的去噪特征,以进行进一步增强。因此,sRGB解码器D_rgb可以更专注于颜色校正。反馈的pipeline如图2所示,可以表示为:
Gated Fusion Modules
GFM旨在通过门控机制自适应地将反馈噪声估计与初始去噪特征融合。在特征门控期间,作者希望有用信息被自适应地选择并沿着空间和通道维度合并。为了提高效率,使用逐点卷积和深度卷积来聚合通道和局部内容信息。然后,作者将混合特征沿通道维度分成两个块,即F_lgate和F_lcon。在通过GELU非线性函数激活后,F_lgate通过逐点乘法对F_lcon进行门控。通过这种门控机制实现了空间和通道的适应性。GFM的详细结构如图3(c)所示。第l阶段(l∈{1, 2, ... L,})的操作可以表示为:
GFM模块代码如下,
import torch
from torch import nn
from torch.nn import functional as F
class GFM(nn.Module):
def __init__(self, in_channels, feature_num=2, bias=True, padding_mode='reflect', **kwargs) -> None:
super().__init__()
self.feature_num = feature_num
hidden_features = in_channels * feature_num
self.pwconv = nn.Conv2d(hidden_features, hidden_features * 2, 1, 1, 0, bias=bias)
self.dwconv = nn.Conv2d(hidden_features * 2, hidden_features * 2, 3, 1, 1, bias=bias, padding_mode=padding_mode, groups=hidden_features * 2)
self.project_out = nn.Conv2d(hidden_features, in_channels, kernel_size=1, bias=bias)
self.mlp = nn.Conv2d(in_channels, in_channels, 1, 1, 0, bias=True)
def forward(self, *inp_feats):
assert len(inp_feats) == self.feature_num
shortcut = inp_feats[0]
x = torch.cat(inp_feats, dim=1)
x = self.pwconv(x)
x1, x2 = self.dwconv(x).chunk(2, dim=1)
x = F.gelu(x1) * x2
x = self.project_out(x)
return self.mlp(x + shortcut)
Residual Switch Mechanism
作者只在RAW域的去噪阶段保留global shortcut,以实现更好的去噪效果,而在颜色恢复阶段则去除它,以避免在噪声的RAW域和干净的sRGB域之间产生模糊的连接,如图2所示。因此,编码器在去噪时需要执行噪声估计,而在颜色恢复期间则需要重建信号。为了实现单一编码器中这两种相互矛盾的功能,作者提出了一个简单但有效的残差开关机制(Residual Switch Mechanism, RSM),如图3(a)所示,它赋予共享RAW编码器中的CID块产生两种相互矛盾的特征:噪声和信号。在带有全局残差连接的去噪阶段,局部残差shortcuts被关闭以估计噪声。相反,在渲染阶段,局部残差被激活,通过在shortcut上的原始特征抵消噪声,最终重建信号。如图4所示,带有RSM的共享RAW编码器的CID块能够在不同阶段产生两种不同的特征。然而,如果没有RSM,权重共享的CID块在颜色恢复阶段无法区分噪声和信号,导致特征模糊。剩余的噪声限制了颜色校正过程,并再次引入了领域模糊。
RSM模块的代码如下,
from torch import nn
class ResidualSwitchBlock(nn.Module):
def __init__(self, block) -> None:
super().__init__()
self.block = block
def forward(self, x, residual_switch):
return self.block(x) + residual_switch * x
Training Objectives
为了顺序完成通过领域特定任务解耦所解耦的RAW去噪和颜色恢复子任务,作者在不同领域引入了两种不同的监督,即干净的RAW和干净的sRGB。真实情况是清晰的RAW图像Y_raw。记去噪解码器的输出RAW图像为Y_^_raw。网络损失函数如下:
实验和分析
数据集和评估指标
作者在两个不同的基于RAW的低光照图像增强数据集上对提出的DNF进行了基准测试,即See-In-the-Dark(SID)数据集和Mono-Colored Raw Paired(MCR)数据集。SID数据集包含5094张极低光照RAW图像,以及由两台相机拍摄的正常光照参考图像:Sony A7S2带有Bayer传感器,分辨率为4240×2832,以及Fuji X-T2带有X-Trans传感器,分辨率为6000×4000。低光照图像的曝光时间从0.1秒到0.033秒不等,参考图像的曝光时间是低光照图像的100到300倍。请注意,Sony子集的测试集中长曝光和短曝光的图像对存在不对齐问题,因此作者在测试阶段遵循以前的方法,丢弃了这些图像。为了公平比较,所有比较的方法都在相同的设置下进行评估。MCR数据集包含4980张图像,分辨率为1280×1024,用于训练和测试,包括3984张低光照RAW图像,498张单色图像和498张sRGB图像。它包括两种不同场景,室内和室外,设置了不同的曝光时间,室内场景为1/256秒到3/8秒,室外场景为1/4096秒到1/32秒。然而,没有提供RAW格式的地面真实图像,这对于训练提出的方法至关重要。因此,作者选择每个场景中曝光时间最长的图像作为RAW地面真值。作者将PSNR、SSIM和LPIPS分别作为像素评估、结构评估和感知评估的定量评估指标。
与SOTA方法对比
定量评估
如表1和表2所示,作者的方法在性能上远远超过了以前的最先进方法。在SID数据集上,DNF在PSNR和LPIPS得分上取得了最佳成绩,与第二佳方法相比,在Sony和Fuji子集上分别实现了0.97 dB和1.30 dB的PSNR提升,以及0.005和0.023的LPIPS提升。关于复杂性,DNF在参数和FLOPs方面明显少于当前最佳方法(即MCR和EEMEFN)。作者的网络使用的参数是MCR和EEMEFN的1/5和1/15,FLOPs分别是它们的3/5和1/13。在MCR数据集上,DNF在表2中展示了最佳PSNR和SSIM得分,与以前的最先进方法相比,分别提高了0.31dB和0.07,同时参数和FLOPs更少。
定性评估
图5和图6显示了在SID数据集上的定性结果。可以看出,与作者提出的方法相比,比较方法增强的结果存在严重的内容失真和伪影,这是由于它们有限的去噪能力造成的。此外,得益于DNF架构,颜色更准确地转换和增强,没有噪声干扰,因此展现出更好的颜色一致性以及更真实和生动的颜色渲染。作者的方法成功地抑制了强烈的噪声,同时保留了丰富的纹理细节。
消融实验
Domain-Specific Task Decoupling
为了更好地评估DNF的领域特定任务解耦的影响,作者对比了对去噪解码器的不同种类的中间监督方法,如表3所示。1)没有监督(w/o Sup.)无法将去噪先验反馈回共享RAW编码器,导致PSNR下降了0.14dB。2)sRGB监督(sRGB Sup.)将主任务分解为第一阶段增强和细节重建。第一阶段增强由于直接从嘈杂的RAW域学习到干净的sRGB域,导致领域歧义,PSNR下降了0.42dB。sRGB Sup.和w/o Sup.之间的比较(0.28dB↓)表明,领域歧义严重限制了网络的性能。
Denoising Prior Feedback
为了验证DNF的基于反馈机制的框架的有效性,作者首先检查了框架的单阶段和多阶段的变体,如表3所示。1)单阶段变体通过直接串联RAW编码器和sRGB解码器,导致PSNR下降了0.46dB。2)多阶段变体简单地串联了两个UNets,分别配备了CID块和MCC块。图像级数据流的损失严重恶化了性能,PSNR下降了0.30dB。
Gated Fusion Module
采用了三种其他的融合模块:传统的卷积层(下降0.22dB),没有门控机制的深度卷积(下降0.27dB),以及专门为特征融合设计的SKFF模块(下降0.25dB)。由于门控机制提供的逐像素选择,所提出的方法享有最佳性能。
Residual Switch Mechanism
如表4所示,与完全没有残差快捷连接相比,利用全局残差快捷连接可以提高性能(0.06dB提升)。然而,在颜色恢复阶段的全局快捷连接会通过引入领域歧义来限制性能(0.03dB下降)。有无所有局部快捷连接的实验都引入了功能上的矛盾,从而导致不同程度的性能下降(与DNF相比,分别为0.33dB下降和0.30dB下降)。与另一种实现残差开关机制的方法相比:在去噪时打开开关或在颜色恢复期间关闭开关,CID块在颜色恢复期间的局部快捷连接提供了更多关于图像内容的信息,从而带来了更高的性能。
Comparison with Other Feature-level Dataflow
如表5所示,与特征级数据流多阶段框架相比,作者的模型提供了最佳性能,这验证了残差开关机制(RSM)的有效性。特征级多阶段框架保留了门控融合模块,但涉及两个不同的RAW编码器。结果表明,使用作者提出的RSM,权重共享的编码器可以执行两种不同的功能。此外,噪声估计和信号重建这两个互补的功能相互补充,共同实现了更好的性能。
结论
针对RAW格式的独特属性,作者提出了一种用于基于RAW的低光照图像增强的DNF框架。作为一个通用的pipeline,所提出的DNF(Decouple and Feedback framework)克服了以前方法的固有限制。领域特定任务解耦消除了单阶段方法引入的领域模糊,而去噪先验反馈则取代了存在有损图像级数据流的多阶段方法。
限制:所提出的框架,也与大多数现有方法一样存在一个限制,即在极低光照条件下,根据曝光时间预定义输入图像的放大比例。在现实世界场景中,估计正常照明是至关重要且困难的。