UniRestorer: Universal Image Restoration via Adaptively Estimating Image Degradation at Proper Granularity

TL;DR

UniRestorer addresses the challenge of all-in-one image restoration by explicitly leveraging degradation priors while remaining robust to estimation errors. It introduces a fine-grained degradation extractor and hierarchical clustering to form a multi-granularity degradation set, paired with a multi-granularity MoE restoration model. Routing uses both degradation and granularity estimates to select appropriate experts, mitigating the impact of degradation-estimation errors. Across single-degradation and mixed-degradation benchmarks, UniRestorer outperforms existing all-in-one methods and closely approaches, or matches, dedicated single-task models, with strong generalization to real-world and unseen degradations. This work offers a practical path toward universal restoration by integrating degradation-aware priors with robust, multi-granularity routing.

Abstract

Recently, considerable progress has been made in all-in-one image restoration. Generally, existing methods can be degradation-agnostic or degradation-aware. However, the former are limited in leveraging degradation-specific restoration, and the latter suffer from the inevitable error in degradation estimation. Consequently, the performance of existing methods has a large gap compared to specific single-task models. In this work, we make a step forward in this topic, and present our UniRestorer with improved restoration performance. Specifically, we perform hierarchical clustering on degradation space, and train a multi-granularity mixture-of-experts (MoE) restoration model. Then, UniRestorer adopts both degradation and granularity estimation to adaptively select an appropriate expert for image restoration. In contrast to existing degradation-agnostic and -aware methods, UniRestorer can leverage degradation estimation to benefit degradation specific restoration, and use granularity estimation to make the model robust to degradation estimation error. Experimental results show that our UniRestorer outperforms state-of-the-art all-in-one methods by a large margin, and is promising in closing the performance gap to specific single task models.

Figures (21)

• Figure 1: Illustration of representative all-in-one image restoration frameworks. (a) Degradation-agnostic methods wang2021realesrganBSRGAN train a shared backbone using data from all tasks and are limited in leveraging degradation-specific restoration. (b) PEFT-based methods AirNetIDRPromptIRTransweatherOneRestoremioirMPerceiver apply learnable prompts or adapters to the backbone to adapt various tasks. (c) Expert-based methods HQ_50KRestoreAgentgridstgsrD2CSRADMSHAIRinstructirDASRU_WADNWGWSNet train specific MoE, LoRA, or filters for specific tasks. However, both (b) and (c) suffer from the inevitable error in degradation estimation. (d) We construct a multi-granularity degradation set and a multi-granularity MoE restoration model. By adaptively estimating image degradation at proper granularity, our UniRestorer can be effective in leveraging degradation-specific restoration while being robust to degradation estimation error.
• Figure 2: Comparisons with task-agnostic methods, all-in-one methods, and single-task models. Ours and Ours† denote auto and instruction modes, which are respectively used for a fair comparison with all-in-one and specific single-task models. (a) Comparisons on single-degradation all-in-one setting. (b) Comparisons with specific single-task models. (c) Comparisons on mixed-degradation all-in-one (in-of-distribution) setting. (d) Comparisons on mixed-degradation all-in-one (out-of-distribution) setting.
• Figure 3: Illustration of our proposed UniRestorer. We develop a multi-granularity degradation set by hierarchical clustering on extracted DRs at different granularities. Based on the multi-granularity degradation set, we train a multi-granularity MoE restoration model. Besides vanilla degradation estimation, we introduce granularity estimation to indicate the degree of degradation estimation error. Adopting both degradation and granularity estimation, we train routers to adaptively allocate an expert to unknown corrupted input.
• Figure 4: Visual comparison of results on All-in-One image restoration (single-degradation). Our method can effectively remove the degradation pattern (i.e., rainstreak, haze, and noise), restore clearer texture (i.e., 'branches' and 'eyes' pointed by red arrows) and closer color (i.e., 'bridge'). More results can be seen in Suppl.
• Figure A: Degradation Pipeline.