[Paper Review] Multi-Stage Progressive Image Restoration

Syed Waqas Zamir, Aditya Arora+5 more|arXiv (Cornell University)|Feb 4, 2021

Advanced Image Processing Techniques100 references52 citations

TL;DR

MPRNet presents a three-stage, progressive image restoration framework that couples encoder-decoder contextual learning with a final high-resolution stage, using supervised attention and cross-stage feature fusion to achieve state-of-the-art deraining, deblurring, and denoising across ten datasets.

ABSTRACT

Image restoration tasks demand a complex balance between spatial details and high-level contextualized information while recovering images. In this paper, we propose a novel synergistic design that can optimally balance these competing goals. Our main proposal is a multi-stage architecture, that progressively learns restoration functions for the degraded inputs, thereby breaking down the overall recovery process into more manageable steps. Specifically, our model first learns the contextualized features using encoder-decoder architectures and later combines them with a high-resolution branch that retains local information. At each stage, we introduce a novel per-pixel adaptive design that leverages in-situ supervised attention to reweight the local features. A key ingredient in such a multi-stage architecture is the information exchange between different stages. To this end, we propose a two-faceted approach where the information is not only exchanged sequentially from early to late stages, but lateral connections between feature processing blocks also exist to avoid any loss of information. The resulting tightly interlinked multi-stage architecture, named as MPRNet, delivers strong performance gains on ten datasets across a range of tasks including image deraining, deblurring, and denoising. The source code and pre-trained models are available at https://github.com/swz30/MPRNet.

Motivation & Objective

• Break down image restoration into manageable sub-tasks via a multi-stage architecture.
• Balance contextual information and fine spatial details by combining encoder-decoder and original-resolution processing.
• Enable progressive learning with stage-wise supervision and feature fusion across stages.
• Preserve information flow with cross-stage connections to stabilize training.

Proposed method

• Use an encoder-decoder subnetwork to learn multi-scale contextual features in early stages.
• In the last stage, operate at original image resolution to preserve fine details.
• Introduce a supervised attention module (SAM) between stages to reweight features using ground-truth guidance.
• Incorporate cross-stage feature fusion (CSFF) to propagate multi-scale contextual features across stages.
• Predict residuals at each stage and add to the degraded input to obtainX_S.
• Train end-to-end with a joint loss combining Charbonnier loss and an edge loss.
• Employ channel attention blocks and ORSNet with original-resolution blocks for detail preservation.

Experimental results

Research questions

• RQ1 Can a multi-stage architecture that fuses multi-scale contextual features with high-resolution detail improve restoration quality across deraining, deblurring, and denoising?
• RQ2 Does supervising intermediate stages and reweighting features via per-pixel attention improve progressive restoration?
• RQ3 What is the impact of cross-stage feature fusion on information flow and training stability?
• RQ4 How well does the model generalize across synthetic and real-world degradation datasets?

Key findings

• MPRNet achieves state-of-the-art performance on ten synthetic and real-world datasets across deraining, deblurring, and denoising.
• The supervised attention module and cross-stage feature fusion provide substantial gains; removing SAM drops PSNR by about 0.42–0.63 dB in ablations, and removing CSFF drops PSNR by about 0.16–0.18 dB in ablations.
• Compared to MSPFN, MPRNet attains approximately 20% relative PSNR improvement on deraining with fewer parameters and faster runtime.
• The last stage operating at original resolution helps preserve fine spatial details while earlier encoder-decoder stages capture context.
• Stage-wise guidance allows a light, fast first stage and progressively stronger later stages, enabling resource-efficient scalability.