极市导读
本文提出了 UniRestore，这是一个基于扩散模型先验的统一图像恢复框架，通过互补特征恢复模块和任务特征适配器，有效弥合了感知式与任务导向图像恢复之间的差距，在多项任务中展现了领先性能和良好的可扩展性。 >>加入极市CV技术交流群，走在计算机视觉的最前沿
-
组织机构：台湾大学, Microsoft, UC Merced, Google Research -
论文：https://arxiv.org/abs/2501.13134 -
主页：https://unirestore.github.io
簡介：
现有图像恢复技术大致可分为两类：感知式图像恢复(Perceptual Image Restoration, PIR)与任务导向图像恢复(Task-oriented Image Restoration, TIR)。前者以提升画面视觉质量为目标，强调图像清晰度与观感，常用于图像修复与增强等应用；后者则关注是否增强图像中的关键语义信息，从而提升下游任务如分类、识别或分割的准确率。以往的研究多聚焦于其中一类，使得模型难以兼顾视觉质量与任务性能，难以满足实际应用中多样化的需求。
本论文提出 UniRestore，一种基于扩散模型先验(Diffusion Prior)的统一图像恢复框架，能够有效弥合PIR与TIR任务之间的差距。为此，我们引入两个关键模块：互补特征恢复模块(Complementary Feature Restoration Module, CFRM)，用于增强编码器的图像表征恢复能力；以及任务特征适配器(Task Feature Adapter, TFA)，通过轻量化的任务提示(Task Prompt)调节扩散模型的解码特征，能够根据不同任务需求动态融合恢复与生成的表征。 此外，UniRestore 在无需修改主模型结构与原始数据的前提下，仅通过新增提示向量(Prompt)即可快速扩展至其他下游任务，具有良好的可扩展性。实验结果表明，UniRestore在视觉质量、分类、语义分割等多项任务中均取得领先性能，并可有效泛化至未知数据集与未知下游模型，展现出极高的实用价值与鲁棒性。
方法：
Overview of UniRestore
UniRestore基于预训练扩散模型，并结合两个关键模块进行任务适配：编码端通过互补特征恢复模块(Complementary Feature Restoration Module, CFRM)从退化图像中恢复干净特征，同时提供增强特征作为补偿信息及扩散生成条件；解码端则结合任务特征适配器(Task Feature Adapter, TFA)，利用轻量级任务提示词(Task Prompts)控制恢复与生成特征的融合，输出对应下游任务的还原结果。该设计实现了统一建模，同时具备良好的扩展性，整体训练流程分为三个阶段：
Stage 1: Complementary Feature Restoration Module & ControlNet
第一阶段旨在强化编码器的特征表征能力。CFRM通过四个步骤依次处理退化输入：特征增强(Feature Enhancement)、组内通道注意(Intra-group Channel Attention)、组间通道整合(Inter-group Channel Integration)与特征重建(Feature Recovery)，从而恢复出兼具清晰度与语义的补偿特征，用于后续解码。此外，最后一层的输出也作为输入提供给 Controller 与 SC-Tuner，通过联合训练共同控制预训练扩散模型的生成过程，使其更符合还原目标与任务需求。
Stage 2: Task Feature Adapter
第二阶段的目标是调适扩散模型的生成能力，以满足不同下游任务的需求。本阶段仅更新任务特征适配器(TFA)，其余模块参数保持不变。TFA负责将 CFRM 恢复的特征与解码器中的扩散特征融合，输出符合任务需求的还原结果。为避免每个任务都需单独设计适配器，TFA 借鉴 prompt tuning 和 LSTM 的设计思路，采用轻量化的任务提示词(Task Prompt)作为调适核心，在各层中逐步更新并引导特征融合。该设计不仅提升了模型的可扩展性，也能有效应对多任务场景，实现高效的模型共享与快速适应。
Introducing Additional Tasks
在新增任务时，仅需引入对应的新任务提示词(Task Prompt)，并使用新任务的数据与目标进行训练，无需重新访问原始任务的数据或损失函数，便可高效扩展至新的任务需求。
实验结果：
本文针对感知式图像还原(PIR)与任务导向图像还原(TIR)进行定性与定量分析，并与具有代表性的PIR与TIR方法进行比较。在原始设置与UniRestore多任务设置下分别训练，全面评估各方法在视觉质量与下游任务表现上的差异与优势。
感知式图像还原比较（Perceptual Image Restoration Comparison）
本文在多个数据集上进行了测试，涵盖已知与未知场景，以验证模型在不同退化条件下的泛化能力与视觉质量。相比传统感知式方法，UniRestore在去噪与细节还原方面表现更加出色，能够有效减少瑕疵与模糊，同时保留更多纹理与边缘信息，展现出优异的视觉一致性与真实感。
任务导向图像还原比较（Task-oriented Image Restoration Comparison）
消融实验(Ablation Study)
本实验比较了四种设置：(i) 仅使用预训练 Stable Diffusion 与控制器的 Baseline；(ii) 移除 CFRM，直接使用原始编码特征；(iii) 移除 TFA，仅使用解码器生成特征；(iv) 完整的 UniRestore 模型。实验结果表明，CFRM 与 TFA 两个模块在感知式影像还原（PIR）与任务导向影像还原（TIR）任务中皆具有关键作用，能显著提升整体性能。
为验证 UniRestore 的扩展能力，论文额外引入目标检测任务，仅需新增对应的任务提示词（Task Prompt），并以RetinaNet作为识别模型进行训练。实验结果显示，无需重新训练整体模型，UniRestore即可顺利适应新任务，展现 TFA 模块设计在多任务场景中的高弹性与优异扩展性。
结语
本论文提出了UniRestore模型，首次实现感知式与任务导向影像还原的统一处理。通过引入互补特征还原模块(CFRM)与任务特征适配器(TFA)，有效整合扩散模型与多种下游任务。实验结果验证了其在视觉质量与任务性能方面的优势，并展现出高度的可扩展性，为多任务图像还原领域提供了全新的解决方案与研究方向。
（文：极市干货）