ICML 2024 | 川大发布用于开集图像复原的测试时退化适应框架

©PaperWeekly 原创 · 作者 | 缑元彪单位 | 四川大学彭玺教授课题组研究方向 | 图像复原

论文标题：

Test-Time Degradation Adaptation for Open-Set Image Restoration

收录会议：

ICML 2024, Spotlight

论文地址：

https://openreview.net/pdf/cbe8a535cb6ad39d7f4315b6eaedd1bcc36a0a4d.pdf

Github地址：

https://github.com/XLearning-SCU/2024-ICML-TAO

背景

图像复原旨在消除图像中的退化现象（如噪声、模糊和雨雾），提升图像的视觉效果和信息内容。近年来，图像复原方法取得了显著进展，能够处理多种多样的退化问题，例如高斯噪声、运动模糊、低分辨率等。
尽管这些方法取得了不错的性能，但大都面向封闭场景，即假设测试时的图像退化与训练时的图像退化是相似的。然而，这种假设在实际应用中往往是不成立的。实际场景中的图像退化情况复杂多变，可能由于设备差异、环境因素或人为干扰等，导致在测试时出现的退化是未知的，并且是在训练数据中没有见到过的。

▲ 图1. 不同图像复原任务之间的区别。（a）经典任务[1]主要面向一种封闭场景，即训练和测试的图像退化是相同且已知的，需要为每种退化定制一个专门的模型；（b）多合一任务 [2] 也面向一种封闭场景，即训练和测试的退化集合是相同的，但需要通过单个模型解决集合中包含的多种未知退化；（c）零样本任务专注于从单张退化图像中直接进行复原 [3]，但通常是面向特定退化进行设计，并利用特定的先验知识进行辅助；而本文提出的（d）开集任务主要面向开放场景，其中测试退化是未知的并与训练时的退化不同，需要通过单个模型去解决多种未知且训练时未见过的图像退化。

方法针对上述问题，本文研究了一个更具挑战性且未被涉足的问题：开集图像复原（OIR）。不同于封闭场景下的方法仅解决训练阶段遇到的特定退化，OIR 要求模型去处理训练数据中不存在的未知退化。
为此，我们在该工作中研究并揭示了 OIR 问题的本质，即测试和训练数据之间的未知分布偏移。针对这一问题，测试时适应（TTA）作为一种有效方法应运而生，以解决测试数据与训练数据之间固有的差异 [4]。简单来说，它在测试阶段根据测试数据来调整预训练模型的参数，使模型能够在更广泛的输入场景中表现更好。
基于此，我们提出了一个用于开集图像复原的测试时退化适应框架 TAO。它利用 TTA 的思想为 OIR 问题提供了一种可行的解决方案。具体地，该框架包括三个组成部分：一个预训练的图像扩散模型（PDM）、一个测试时退化适配器（TDA）和适配器引导的图像复原（AIR）。
PDM 被采用为解决 OIR 问题的基础模型主要基于以下考虑。首先，PDM 具备生成各种高质量视觉场景的丰富知识，可以被视为产生干净图像的通用预训练模型。其次，PDM 是退化无关的，测试数据中的任何退化都可被视为未见过的。在 PDM 的每个去噪步骤之后，TDA 和 AIR 被依次执行，分别用于适应开集场景和引导图像复原。
具体地，TDA 在测试阶段采用一个可学习的适配器来适配 PDM 到测试退化图像上。这个适配器被设计用于域对齐，将 PDM 的生成域与测试图像的退化域对齐。这样，生成的干净图像可以被转换为相应的退化图像，然后可以进一步被引导更新以生成对应的干净图像。AIR 被设计来进行这种监督更新，通过在 PDM 不同的去噪步骤中动态地调整引导更新策略，以实现更好的图像恢复。

▲ 图2. 所提出的 TAO 框架利用（i）预训练扩散模型（PDM）作为OIR的通用预训练模型，并在每个去噪步后，首先执行（ii）测试时退化适配器（TDA）以适应开放场景中未知和未见过的退化，然后进行（iii）适配器引导的图像复原（AIR）以引导生成图像成为复原出的干净图像。注意：雪花表示图像或模型是固定的，火焰表示图像或模型将通过梯度进行更新。

实验我们在多种退化类型上验证了 TAO 框架在处理 OIR 问题时的有效性，包括图像去雾、低光照图像增强和图像去噪任务。结果显示，其在多个评价指标上与为任务专门设计的方法相比，取得了同等甚至更优的性能。部分实验结果如下，更多实验请参考论文。

▲ 表1. 在 HSTS 数据集上的图像去雾结果。我们的方法优于零样本方法，并获得了与经典有监督学习方法相当甚至更好的结果。比较的所有方法都是专门为图像去雾设计的。

▲ 表2. 在 LOL 数据集上低光照图像增强结果。我们的方法在零样本方法中获得了最佳的 SSIM 值和第二好的 PSNR 值，并只有监督学习的 MBLLEN 在这两个指标上能同时优于我们的方法。

▲ 图3. 图像去雾的视觉结果，从中可以观察到现有方法过度去雾导致图像变暗和/或产生伪影。相比之下，我们的方法获得了更清晰、更接近自然真实的结果。

▲ 图4. 低光照图像增强的结果，从中可以看到我们的结果既不像 MBLLEN 那样平滑，也不像 ZDCE 那样暗淡。尽管存在轻微的色差，我们的方法实现了对低光照图像的合理照明。

该工作从实际出发，正式定义了开集图像复原问题（测试退化是未知的并与训练时的退化不同，且通过单个模型去解决多种未知且训练时未见过的图像退化）。为了解决这个问题，本工作从分布偏移的角度揭示了其本质，并发现测试时适应是解决这种内在偏移的有效方法。
因此，我们提出了一个用于开集图像复原的测试时退化适应框架，其巧妙之处在于以下几点：首先，它考虑了一个预训练的图像扩散模型作为解决各种图像复原任务的通用预训练模型；其次，它引入了一个在测试阶段优化的适配器，用于使预训练模型适应未知和未见过的测试退化；第三，它根据去噪过程动态调整引导策略，获得了更好的复原结果。通过在多种退化上的实验，我们展示了这些设计的有效性。
自 2022 年本课题组（四川大学计算机学院彭玺教授课题组）在国际上率先提出多合一图像复原问题（即训练和测试的退化集合是相同的，但需要通过单个模型解决集合中包含的多种未知退化），该研究方向已成为影像复原和增强领域近年来最受关注的新方向之一。开集图像复原是本课题组提出的多合一影像复原任务面向真实开放世界的进一步延伸和靠近，有望给领域提供新的洞见和理解。

参考文献

[1] Yuanbiao Gou, Boyun Li, Zitao Liu, Songfan Yang, and Xi Peng*, CLEARER: Multi-Scale Neural Architecture Search for Image Restoration, Neural Information Processing Systems (NeurIPS), 2020.

[2] Boyun Li, Xiao Liu, Peng Hu, Zhongqin Wu, Jiancheng Lv, Xi Peng*, All-In-One Image Restoration for Unknown Corruption, IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), New Orleans, LA. Jun. 19-25, 2022.

[3] Boyun Li#, Yuanbiao Gou#, Shuhang Gu, Jerry Zitao Liu, Joey Tianyi Zhou, and Xi Peng*, You Only Look Yourself: Unsupervised and Untrained Single Image Dehazing Neural Network, International Journal of Computer Vision (IJCV), 2021.

[4] Wang, D., et al. Tent: Fully test-time adaptation by entropy minimization. In Proceedings of the International Conference on Learning Representations, 2021.

更多阅读