Learning Frequency-Aware Dynamic Transformers for All-In-One Image Restoration

Zenglin Shi,Tong Su,Pei Liu,Yunpeng Wu,Le Zhang,Meng Wang

2024-06-30

Abstract:This work aims to tackle the all-in-one image restoration task, which seeks to handle multiple types of degradation with a single model. The primary challenge is to extract degradation representations from the input degraded images and use them to guide the model's adaptation to specific degradation types. Recognizing that various degradations affect image content differently across frequency bands, we propose a new all-in-one image restoration approach from a frequency perspective, leveraging advanced vision transformers. Our method consists of two main components: a frequency-aware Degradation prior learning transformer (Dformer) and a degradation-adaptive Restoration transformer (Rformer). The Dformer captures the essential characteristics of various degradations by decomposing inputs into different frequency components. By understanding how degradations affect these frequency components, the Dformer learns robust priors that effectively guide the restoration process. The Rformer then employs a degradation-adaptive self-attention module to selectively focus on the most affected frequency components, guided by the learned degradation representations. Extensive experimental results demonstrate that our approach outperforms the existing methods on four representative restoration tasks, including denoising, deraining, dehazing and deblurring. Additionally, our method offers benefits for handling spatially variant degradations and unseen degradation levels.

Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

The paper aims to address multiple challenges in the full-function image restoration task with the following specific objectives:

1. **Handling Multiple Degradation Types**: Propose a single model approach capable of handling various image degradations to overcome the limitations of existing methods designed for specific degradation tasks when faced with different types or severities of degradation.
2. **Extracting Degradation Representations**: Capture the essential characteristics of various degradation types by decomposing different frequency components of the input image, thereby learning robust prior knowledge to guide the restoration process.
3. **Adaptive Restoration Mechanism**: Develop a degradation-adaptive self-attention module that enables the model to selectively focus on the most severely affected frequency components in different types of degradation tasks.

To achieve the above objectives, the authors propose a new method comprising two main components:
- **Frequency-Aware Degradation Representation Learning Transformer (Dformer)**: Responsible for estimating degradation representations from the input degraded image.
- **Degradation-Adaptive Restoration Transformer (Rformer)**: Utilizes the learned degradation representations to guide the restoration process and improves restoration performance by adaptively focusing on the most severely affected frequency components.

Experimental results show that this method outperforms existing methods in four representative restoration tasks: denoising, deraining, dehazing, and deblurring. It also excels in handling spatially varying degradations and unseen degradation levels.