Cat-AIR: Content and Task-Aware All-in-One Image Restoration

TL;DR

Cat-AIR proposes a content- and task-aware all-in-one image restoration framework that unifies multiple degradation types under a single model. It introduces an alternating spatial-channel attention scheme with cross-layer channel attention and cross-feature spatial attention to adapt computation to content and task difficulty, paired with a smooth learning strategy for scalable task extension via task-specific prompts and EMA updates. The approach achieves state-of-the-art performance on denoising, deraining, dehazing, and extended tasks like deblurring and low-light enhancement, while reducing FLOPs compared to prior methods. This combination of adaptive computation and scalable task extension offers significant practical benefits for real-world restoration where multiple degradations can co-exist and evolve over time.

Abstract

All-in-one image restoration seeks to recover high-quality images from various types of degradation using a single model, without prior knowledge of the corruption source. However, existing methods often struggle to effectively and efficiently handle multiple degradation types. We present Cat-AIR, a novel \textbf{C}ontent \textbf{A}nd \textbf{T}ask-aware framework for \textbf{A}ll-in-one \textbf{I}mage \textbf{R}estoration. Cat-AIR incorporates an alternating spatial-channel attention mechanism that adaptively balances the local and global information for different tasks. Specifically, we introduce cross-layer channel attentions and cross-feature spatial attentions that allocate computations based on content and task complexity. Furthermore, we propose a smooth learning strategy that allows for seamless adaptation to new restoration tasks while maintaining performance on existing ones. Extensive experiments demonstrate that Cat-AIR achieves state-of-the-art results across a wide range of restoration tasks, requiring fewer FLOPs than previous methods, establishing new benchmarks for efficient all-in-one image restoration.

Figures (16)

• Figure 1: Illustration of content and task-aware image restoration. Left) Image patches with varying texture complexity and degradation levels. Right) We dynamically apply self-attention for complex regions/tasks, while using convolution for simple regions/tasks.
• Figure 2: Average performance (PSNR vs. FLOPs) across three degradations (denoising, deraining, dehazing) and five degradations (including deblurring and low-light enhancement). Our methods achieve state-of-the-art results with lower computations.
• Figure 3: Overview of Cat-AIR. Our design features alternating channel and spatial attention mechanisms, where channel attention complexity scales across layers and spatial attention complexity adapts based on features. Prompt modules between decoder blocks are inserted to identify degradations, following potlapalli2024promptir.
• Figure 4: Architectural diagrams of (a) cross-layer channel attention and (b) cross-feature spatial attention mechanisms.
• Figure 5: Visual comparison on five image restoration tasks. We compare against both specialized models (Restormer restormer and NAFNet chen2022simple, which use separate models for each degradation type) and all-in-one methods (TAPE liu2022tape, AirNet Li_2022_CVPR, and IDR zhang2023ingredient).