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.
- 2024.12.31: Paper and supplement files are release on ArXiv.
- Inference code and pre-trained models release.
- Datasets, training code release.
Our method has three steps, constructing multi-granularity degradation set, train multi-granularity MoE restoration model, train degradation and granularity estimation-based routing.
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.
Following instructions in DATA.md to download training and testing data.
If you have any questions, please contact [email protected].
We thank to the following image restoration works for their awesome backbones and code repos:
Our code is based on BasicSR and KAIR.
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}
}