Selective Hourglass Mapping for Universal Image Restoration Based on Diffusion Model

TL;DR

This paper introduces DiffUIR, a diffusion-based universal image restoration method with a selective hourglass mapping strategy that effectively handles multiple degradation types, achieving state-of-the-art results with a lightweight model.

Contribution

Proposes a novel selective hourglass mapping strategy for diffusion models that unifies multiple image restoration tasks into a shared framework with high performance.

Findings

Achieves state-of-the-art results on five image restoration tasks.

Operates effectively in zero-shot generalization settings.

Uses a lightweight model with only 0.89M parameters.

Abstract

Universal image restoration is a practical and potential computer vision task for real-world applications. The main challenge of this task is handling the different degradation distributions at once. Existing methods mainly utilize task-specific conditions (e.g., prompt) to guide the model to learn different distributions separately, named multi-partite mapping. However, it is not suitable for universal model learning as it ignores the shared information between different tasks. In this work, we propose an advanced selective hourglass mapping strategy based on diffusion model, termed DiffUIR. Two novel considerations make our DiffUIR non-trivial. Firstly, we equip the model with strong condition guidance to obtain accurate generation direction of diffusion model (selective). More importantly, DiffUIR integrates a flexible shared distribution term (SDT) into the diffusion algorithm elegantly and naturally, which gradually maps different distributions into a shared one. In the reverse process, combined with SDT and strong condition guidance, DiffUIR iteratively guides the shared distribution to the task-specific distribution with high image quality (hourglass). Without bells and whistles, by only modifying the mapping strategy, we achieve state-of-the-art performance on five image restoration tasks, 22 benchmarks in the universal setting and zero-shot generalization setting. Surprisingly, by only using a lightweight model (only 0.89M), we could achieve outstanding performance. The source code and pre-trained models are available at https://github.com/iSEE-Laboratory/DiffUIR