Selective Hourglass Mapping for Universal Image Restoration Based on Diffusion Model

通用图像恢复是一个实际且潜在的计算机视觉任务，适用于现实世界的应用。这项任务的主要挑战是同时处理不同的退化分布。现有方法主要利用任务特定条件（例如提示）来指导模型分别学习不同的分布，称为多部分映射。然而，它不适用于通用模型学习，因为它忽略了不同任务之间的共享信息。在这项工作中，我们提出了一种基于扩散模型的先进选择性沙漏映射策略，称为DiffUIR。两个新颖的考虑使我们的DiffUIR非常重要。首先，我们为模型配备了强条件指导，以获取扩散模型的准确生成方向（选择性）。更重要的是，DiffUIR将一个灵活的共享分布项（SDT）巧妙而自然地集成到扩散算法中，逐渐将不同的分布映射到一个共享分布中。在反向过程中，结合SDT和强条件指导，DiffUIR迭代地将共享分布引导到具有高图像质量的任务特定分布（沙漏）。仅通过修改映射策略，我们在五个图像恢复任务、通用设置和零-shot泛化设置的22个基准测试中实现了最先进的性能，而且仅使用轻量级模型（仅0.89M），表现出色。源代码和预训练模型可在https://github.com/iSEE-Laboratory/DiffUIR上获得。

本文旨在解决通用图像修复中的挑战，即如何同时处理不同的退化分布。

本文提出了一种基于扩散模型的高级选择性沙漏映射策略，称为DiffUIR，它将灵活的共享分布项（SDT）优雅地自然地集成到扩散算法中，并将不同的分布逐渐映射到共享分布中。

本文通过修改映射策略，在只使用轻量级模型的情况下，在五个图像修复任务中，包括22个通用设置和零样本泛化设置的基准测试上实现了最先进的性能。作者提供了源代码和预训练模型。

最近在这个领域中，还有一些相关的研究，如：Multi-Task Learning for Image Restoration, Deep Image Prior, Universal Denoiser等。