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

Jingbo Lin, Zhilu Zhang, Wenbo Li, Renjing Pei, Hang Xu, Hongzhi Zhang, and Wangmeng Zuo

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.

TODO

• [x] 2024.12.31: Paper and supplement files are release on ArXiv.
• [ ] Inference code and pre-trained models release.
• [ ] Datasets, training code release.

Method

Our method has three steps, constructing multi-granularity degradation set, train multi-granularity MoE restoration model, train degradation and granularity estimation-based routing.

Comparisons

All-in-One Image Restoration (single-degradation) (click to expand)

All-in-One Image Restoration (mixed-degradation) (click to expand)

Compare to Single-task Models (click to expand)

More comparison resuslts can be found in supplement material or our website.

Datasets

Following instructions in DATA.md to download training and testing data.

Contact & Acknowledgement

If you have any questions, please contact [email protected].

We thank to the following image restoration works for their awesome backbones and code repos:

• Restormer
• RetinexFormer
• DehazeFormer
• MiOIR
• OneRestore

Our code is based on BasicSR and KAIR.

Citation

If our work is helpful, you can cite our work as follows:

@article{lin2024unirestorer,
  title={UniRestorer: Universal Image Restoration via Adaptively Estimating Image Degradation at Proper Granularity},
  author={Lin, Jingbo and Zhang, Zhilu and Li, Wenbo and Pei, Renjing and Xu, Hang and Zhang, Hongzhi and Zuo, Wangmeng},
  journal={arXiv preprint arXiv:2412.20157},
  year={2024}
}