DPMambaIR: All-in-One Image Restoration via Degradation-Aware Prompt State Space Model

TL;DR

DPMambaIR addresses the challenge of restoring images under diverse, unknown degradations with a single model. It introduces a fine-grained Degradation Extractor and a Degradation-Aware Prompt State Space Model (DP-SSM) that dynamically modulates state-space parameters using degradation embeddings, coupled with a High-Frequency Enhancement Block to recover details. The approach demonstrates state-of-the-art performance on a mixed seven-degradation dataset (PSNR $27.69$ dB, SSIM $0.893$) and strong results on degradation-specific tasks, while ablations confirm the contributions of the degradation prompts and frequency guidance. This work advances practical All-in-One restoration by enabling fine-grained degradation modeling and efficient global reasoning, with potential impact on real-world imaging systems subject to multiple degradations.

Abstract

All-in-One image restoration aims to address multiple image degradation problems using a single model, offering a more practical and versatile solution compared to designing dedicated models for each degradation type. Existing approaches typically rely on Degradation-specific models or coarse-grained degradation prompts to guide image restoration. However, they lack fine-grained modeling of degradation information and face limitations in balancing multi-task conflicts. To overcome these limitations, we propose DPMambaIR, a novel All-in-One image restoration framework that introduces a fine-grained degradation extractor and a Degradation-Aware Prompt State Space Model (DP-SSM). The DP-SSM leverages the fine-grained degradation features captured by the extractor as dynamic prompts, which are then incorporated into the state space modeling process. This enhances the model's adaptability to diverse degradation types, while a complementary High-Frequency Enhancement Block (HEB) recovers local high-frequency details. Extensive experiments on a mixed dataset containing seven degradation types show that DPMambaIR achieves the best performance, with 27.69dB and 0.893 in PSNR and SSIM, respectively. These results highlight the potential and superiority of DPMambaIR as a unified solution for All-in-One image restoration.

Figures (8)

• Figure 1: Comparison with SOTA All-in-One Image Restoration Methods on mixed Dataset. Our DPMambaIR achieves consistently better performance.
• Figure 2: The overall framework of DPMambaIR consists of a shallow feature extractor and multiple DPSS blocks, with an upsampler implemented by pixel shuffle and a downsampler by pixel unshuffle. Additionally, a degradation extractor is included to extract degradation embeddings, which are then inserted into the DPSS blocks.
• Figure 3: (a) Real-World Image Degradation Modeling Process. (b) Architecture details of the Large version degradation extractor and reconstructor. (c) Architecture details of the normal version degradation extractor and reconstructor.
• Figure 4: The Architecture of (a) Degradation-aware Prompt Selective Scan (DPSS) and (b) High-frequency Enhancement Block (HEB), with Three Inputs $F_1$, $F_2$, and Degradation Embedding $\boldsymbol{E}_{d}$.
• Figure 5: We compare frequency-domain differences between clear and restored images under three configurations. The All-in-One setting shows a notable high-frequency gap compared to the Degradation-specific setting, which is significantly reduced by adding the High-frequency Enhancement Block (HEB), as shown by the blue line.