Multi-Stage Progressive Image Restoration

Syed Waqas Zamir, Aditya Arora, Salman Khan, Munawar Hayat, Fahad Shahbaz Khan, Ming-Hsuan Yang, Ling Shao

Introduction

Image restoration is the task of recovering a clean image from its degraded version. Typical examples of degradation include noise, blur, rain, haze, etc. It is a highly ill-posed problem as there exist infinite feasible solutions. In order to restrict the solution space to valid/natural images, existing restoration techniques explicitly use image priors that are handcrafted with empirical observations. However, designing such priors is a challenging task and often not generalizable. To ameliorate this issue, recent state-of-the-art approaches employ convolutional neural networks (CNNs) that implicitly learn more general priors by capturing natural image statistics from large-scale data.

The performance gain of CNN-based methods over the others is primarily attributed to its model design. Numerous network modules and functional units for image restoration have been developed including recursive residual learning , dilated convolutions , attention mechanisms , dense connections , encoder-decoders , and generative models . Nevertheless, nearly all of these models for low-level vision problems are based on single-stage design. In contrast, multi-stage networks are shown to be more effective than their single-stage counterparts in high-level vision problems such as pose-estimation , scene parsing and action segmentation .

Recently, few efforts have been made to bring the multi-stage design to image deblurring , and image deraining . We analyze these approaches to identify the architectural bottlenecks that hamper their performance. First, existing multi-stage techniques either employ the encoder-decoder architecture which is effective in encoding broad contextual information but unreliable in preserving spatial image details, or use a single-scale pipeline that provides spatially accurate but semantically less reliable outputs. However, we show that the combination of both design choices in a multi-stage architecture is needed for effective image restoration. Second, we show that naively passing the output of one stage to the next stage yields suboptimal results . Third, unlike in , it is important to provide ground-truth supervision at each stage for progressive restoration. Finally, during multi-stage processing, a mechanism to propagate intermediate features from earlier to later stages is required to preserve contextualized features from the encoder-decoder branches.

We propose a multi-stage progressive image restoration architecture, called MPRNet, with several key components. 1). The earlier stages employ an encoder-decoder for learning multi-scale contextual information, while the last stage operates on the original image resolution to preserve fine spatial details. 2). A supervised attention module (SAM) is plugged between every two stages to enable progressive learning. With the guidance of ground-truth image, this module exploits the previous stage prediction to compute attention maps that are in turn used to refine the previous stage features before being passed to the next stage. 3). A mechanism of cross-stage feature fusion (CSFF) is added that helps propagating multi-scale contextualized features from the earlier to later stages. Furthermore, this method eases the information flow among stages, which is effective in stabilizing the multi-stage network optimization.

A novel multi-stage approach capable of generating contextually-enriched and spatially accurate outputs. Due to its multi-stage nature, our framework breaks down the challenging image restoration task into sub-tasks to progressively restore a degraded image.

An effective supervised attention module that takes full advantage of the restored image at every stage in refining incoming features before propagating them further.

A strategy to aggregate multi-scale features across stages.

We demonstrate the effectiveness of our MPRNet by setting new state-of-the-art on ten synthetic and real-world datasets for various restoration tasks including image deraining, deblurring, and denoising while maintaining a low complexity (see Fig. 1). Further, we provide detailed ablations, qualitative results, and generalization tests.

Related Work

Recent years have witnessed a paradigm shift from high-end DSLR cameras to smartphone cameras. However, capturing high-quality images with smartphone cameras is challenging. Image degradations are often present in images either due to the limitations of cameras and/or adverse ambient conditions. Early restoration approaches are based on total variation , sparse coding , self-similarity , gradient prior , etc. Recently, CNN-based restoration methods have achieved state-of-the-art results . In terms of architectural design, these methods can be broadly categorized as single-stage and multi-stage.

Single-Stage Approaches. Currently, the majority of image restoration methods are based on a single-stage design, and the architectural components are usually based on those developed for high-level vision tasks. For example, residual learning has been used to perform image denoising , image deblurring and image deraining . Similarly, to extract multi-scale information, the encoder-decoder and dilated convolution models are often used . Other single-stage approaches incorporate dense connections .

Multi-Stage Approaches. These methods aim to recover clean image in a progressive manner by employing a light-weight subnetwork at each stage. Such a design is effective since it decomposes the challenging image restoration task into smaller easier sub-tasks. However, a common practice is to use the identical subnetwork for each stage which may yield suboptimal results, as shown in our experiments (Section 4).

Attention. Driven by its success in high-level tasks such as image classification , segmentation and detection , attention modules have been used in low-level vision tasks . Examples abound, including methods for image deraining , deblurring , super-resolution , and denoising . The main idea is to capture long-range inter-dependencies along spatial dimensions , channel dimensions , or both .

Multi-Stage Progressive Restoration

The proposed framework for image restoration, shown in Fig. 2, consists of three stages to progressively restore images. The first two stages are based on encoder-decoder subnetworks that learn the broad contextual information due to large receptive fields. Since image restoration is a position-sensitive task (which requires pixel-to-pixel correspondence from the input to output), the last stage employs a subnetwork that operates on the original input image resolution (without any downsampling operation), thereby preserving the desired fine texture in the final output image.

Instead of simply cascading multiple stages, we incorporate a supervised attention module between every two stages. With the supervision of ground-truth images, our module rescales the feature maps of the previous stage before passing them to the next stage. Furthermore, we introduce a cross-stage feature fusion mechanism where the intermediate multi-scale contextualized features of the earlier subnetwork help consolidating the intermediate features of the latter subnetwork.

Although MPRNet stacks multiple stages, each stage has an access to the input image. Similar to the recent restoration methods , we adapt the multi-patch hierarchy on the input image and split the image into non-overlapping patches: four for stage- $1$ , two for stage- $2$ , and the original image for the last stage, as shown in Fig. 2.

At any given stage ${S}$ , instead of directly predicting a restored image $\mathbf{X}_{S}$ , the proposed model predicts a residual image $\mathbf{R}_{S}$ to which the degraded input image $\mathbf{I}$ is added to obtain: $\mathbf{X}_{S}=\mathbf{I}+\mathbf{R}_{S}$ . We optimize our MPRNet end-to-end with the following loss function:

where $\mathbf{{Y}}$ represents the ground-truth image, and $\mathcal{L}_{char}$ is the Charbonnier loss :

with constant $\varepsilon$ empirically set to $10^{-3}$ for all the experiments. In addition, $\mathcal{L}_{edge}$ is the edge loss, defined as:

where $\Delta$ denotes the Laplacian operator. The parameter $\lambda$ in Eq. (1) controls the relative importance of the two loss terms, which is set to $0.05$ as in . Next, we describe each key element of our method.

Existing single-stage CNNs for image restoration typically use one of the following architecture designs: 1). An encoder-decoder, or 2). A single-scale feature pipeline. The encoder-decoder networks first gradually map the input to low-resolution representations, and then progressively apply reverse mapping to recover the original resolution. While these models effectively encode multi-scale information, they are prone to sacrificing spatial details due to the repeated use of downsampling operation. In contrast, the approaches that operate on single-scale feature pipeline are reliable in generating images with fine spatial details . However, their outputs are semantically less robust due to the limited receptive field. This indicates the inherent limitations of the aforementioned architecture design choices that are capable of generating either spatially accurate or contextually reliable outputs, but not both. To exploit the merits of both designs, we propose a multi-stage framework where earlier stages incorporate the encoder-decoder networks, and the final stage employs a network that operates on the original input resolution.

Encoder-Decoder Subnetwork. Figure 3(a) shows our encoder-decoder subnetwork, which is based on the standard U-Net , with the following components. First, we add channel attention blocks (CABs) to extract features at each scale (See Fig. 3(b) for CABs). Second, the feature maps at U-Net skip connections are also processed with the CAB. Finally, instead of using Transposed convolution for increasing spatial resolution of features in the decoder, we use bilinear upsampling followed by a convolution layer. This helps reduce checkerboard artifacts in the output image that often arise due to the Transposed convolution .

Original Resolution Subnetwork. In order to preserve fine details from the input image to the output image, we introduce the original-resolution subnetwork (ORSNet) in the last stage (see Fig. 2). ORSNet does not employ any downsampling operation and generates spatially-enriched high-resolution features. It consists of multiple original-resolution blocks (ORBs), each of which further contains CABs. The schematic of ORB is illustrated in Fig. 3(b).

2 Cross-stage Feature Fusion

In our framework, we introduce the CSFF module between two encoder-decoders (see Fig. 3(c)), and between encoder-decoder and ORSNet (see Fig. 3(d)). Note that the features from one stage are first refined with $1\times 1$ convolutions before propagating them to the next stage for aggregation. The proposed CSFF has several merits. First, it makes the network less vulnerable by the information loss due to repeated use of up- and down-sampling operations in the encoder-decoder. Second, the multi-scale features of one stage help enriching the features of the next stage. Third, the network optimization procedure becomes more stable as it eases the flow of information, thereby allowing us to add several stages in the overall architecture.

3 Supervised Attention Module

Recent multi-stage networks for image restoration directly predict an image at each stage, which is then passed to the next consecutive stage. Instead, we introduce a supervised attention module between every two stages, which facilitates achieving significant performance gain. The schematic diagram of SAM is shown in Fig. 4, and its contributions are two-fold. First, it provides ground-truth supervisory signals useful for the progressive image restoration at each stage. Second, with the help of locally supervised predictions, we generate attention maps to suppress the less informative features at the current stage and only allow the useful ones to propagate to the next stage.

Experiments and Analysis

We evaluate our method for several image restoration tasks, including (a) image deraining, (b) image deblurring, and (c) image denoising on $10$ different datasets.

Quantitative comparisons are performed using the PSNR and SSIM metrics. As in , we report (in parenthesis) the reduction in error for each method relative to the best performing method by translating PSNR to RMSE ( $\textrm{RMSE}\propto\sqrt{10^{-\textrm{PSNR}/10}}$ ) and SSIM to DSSIM ( $\textrm{DSSIM}=(1-\textrm{SSIM})/2$ ). The datasets used for training and testing are summarized in Table 1 and described next.

Image Deraining. Using the same experimental setups of the recent best method on image deraining , we train our model on $13$ , $712$ clean-rain image pairs gathered from multiple datasets , as shown in Table 1. With this single trained model, we perform evaluation on various test sets, including Rain100H , Rain100L , Test100 , Test2800 , and Test1200 .

Image Deblurring. As in , we use the GoPro dataset that contains $2$ , $103$ image pairs for training and $1$ , $111$ pairs for evaluation. Furthermore, to demonstrate generalizability, we take our GoPro trained model and directly apply it on the test images of the HIDE and RealBlur datasets. The HIDE dataset is specifically collected for human-aware motion deblurring and its test set contains $2$ , $025$ images. While the GoPro and HIDE datasets are synthetically generated, the image pairs of RealBlur dataset are captured in real-world conditions. The RealBlur dataset has two subsets: (1) RealBlur-J is formed with the camera JPEG outputs, and (2) RealBlur-R is generated offline by applying white balance, demosaicking, and denoising operations to the RAW images.

Image Denoising. To train our model for image denoising task, we use $320$ high-resolution images of the SIDD dataset . Evaluation is conducted on $1$ , $280$ validation patches from the SIDD dataset and $1$ , $000$ patches from the DND benchmark dataset . These test patches are extracted from the full resolution images by the original authors. Both SIDD and DND datasets consist of real images.

2 Implementation Details

Our MPRNet is end-to-end trainable and requires no pre-training. We train separate models for three different tasks. We employ $2$ CABs at each scale of the encoder-decoder, and for downsampling we use $2<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"><semantics><mrow><mo>×</mo></mrow><annotation encoding="application/x-tex">\times</annotation></semantics></math>×2$ max-pooling with stride $2$ . In the last stage, we employ ORSNet that contains $3$ ORBs, each of which further uses $8$ CABs. Depending on the task complexity, we scale the network width by setting the number of channels to $40$ for deraining, $80$ for denoising, and $96$ for deblurring. The networks are trained on $256<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"><semantics><mrow><mo>×</mo></mrow><annotation encoding="application/x-tex">\times</annotation></semantics></math>×256$ patches with a batch size of $16$ for $4<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"><semantics><mrow><mo>×</mo></mrow><annotation encoding="application/x-tex">\times</annotation></semantics></math>×10^{5}$ iterations. For data augmentation, horizontal and vertical flips are randomly applied. We use Adam optimizer with the initial learning rate of $2<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"><semantics><mrow><mo>×</mo></mrow><annotation encoding="application/x-tex">\times</annotation></semantics></math>×10^{-4}$ , which is steadily decreased to $1<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"><semantics><mrow><mo>×</mo></mrow><annotation encoding="application/x-tex">\times</annotation></semantics></math>×10^{-6}$ using the cosine annealing strategy .

3 Image Deraining Results

For the image deraining task, consistent with prior work , we compute image quality scores using the Y channel (in YCbCr color space). Table 2 shows that our method significantly advances state-of-the-art by consistently achieving better PSNR/SSIM scores on all five datasets. Compared to the recent best algorithm MSPFN , we obtain a performance gain of $1.98$ dB (average across all datasets), indicating $20\%$ error reduction. The improvements on some datasets are as large as ${4}$ dB, e.g., Rain100L . Further, our model has ${3.7}\times$ fewer parameters than MSPFN , while being ${2.4}\times$ faster.

Figure 5 shows visual comparisons on challenging images. Our MPRNet is effective in removing rain streaks of different orientations and magnitudes, and generates images that are visually pleasant and faithful to the ground-truth. In contrast, other approaches compromise structural content (first row), introduce artifacts (second row), and do not completely remove rain streaks (third row).

4 Image Deblurring Results

We report the performance of evaluated image deblurring approaches on the synthetic GoPro and HIDE datasets in Table 3. Overall, our model performs favorably against other algorithms. Compared to the previous best performing technique , our method achieves $9\%$ improvement in PSNR and $21\%$ in SSIM on the GoPro dataset, and a $11\%$ and $13\%$ reduction in error on the HIDE dataset . It is worth noticing that our network is trained only on the GoPro dataset, but achieves the state-of-the-art results (+ $0.98$ dB) on the HIDE dataset, thereby demonstrating its strong generalization capability.

We evaluate our MPRNet on the real-world images of a recent RealBlur dataset under two experimental settings: 1). apply the GoPro trained model directly on RealBlur (to test generalization to real images), and 2). train and test on RealBlur data. Table 4 shows the experimental results. For setting $1$ , our MPRNet obtains performance gains of $0.29$ dB on the RealBlur-R subset and $0.28$ dB on the RealBlur-J subset over the DMPHN algorithm . A similar trend is observed for setting $2$ , where our gains over SRN are $0.66$ dB and $0.38$ dB on RealBlur-R and RealBlur-J, respectively.

Figure 6 shows some deblurred images by the evaluated approaches. Overall, the images restored by our model are sharper and closer to the ground-truth than those by others.

5 Image Denoising Results

In Table 5, we report PSNR/SSIM scores of several image denoising methods on the SIDD and DND datasets. Our method obtains considerable gains over the state-of-the-art approaches, i.e., $0.19$ dB over CycleISP on SIDD and $0.21$ dB over SADNet on DND. Note that the DND dataset does not contain any training images, i.e., the complete publicly released dataset is just a test set. Experimental results on the DND benchmark with our SIDD trained model demonstrates our model generalizes well to different image domains.

Fig. 7 illustrates visual results. Our method is able to remove real noise, while preserving the structural and textural image details. In contrast, the images restored by other methods contain either overly smooth contents, or artifacts with splotchy textures.

6 Ablation Studies

Here we present ablation experiments to analyze the contribution of each component of our model. Evaluation is performed on the GoPro dataset with the deblurring models trained on image patches of size $128<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"><semantics><mrow><mo>×</mo></mrow><annotation encoding="application/x-tex">\times</annotation></semantics></math>×128$ for $10^{5}$ iterations, and the results are shown in Table 6.

Number of stages. Our model yields better performance as the number of stages increases, which validates the effectiveness of our multi-stage design.

Choices of subnetworks. Since each stage of our model could employ different subnetwork design, we test different options. We show that using the encoder-decoder in the earlier stage(s) and the ORSNet in the last stage leads to improved performance ( $29.7$ dB) as compared to employing the same design for all the stages ( $29.4$ dB with U-Net+U-Net, and $29.53$ dB with ORSNet+ORSNet).

SAM and CSFF. We demonstrate the effectiveness of the proposed supervised attention module and cross-stage feature fusion mechanism by removing them from our final model. Table 6 shows a substantial drop in PSNR from $30.49$ dB to $30.07$ dB when SAM is removed, and from $30.49$ dB to $30.31$ dB when we take out CSFF. Removing both of these components degrades the performance by a large margin from $30.49$ dB to $29.86$ dB.

Resource Efficient Image Restoration

CNN models generally exhibit a trade-off between accuracy and computational efficiency. In the pursuit of achieving higher accuracy, deeper and complex models are often developed. Although large models tend to perform better than their smaller counterparts, the computational cost can be prohibitively high. As such, it is of great interest to develop resource-efficient image restoration models. One solution is to train the same network by adjusting its capacity every time the target system is changed. However, it is tedious and oftentimes infeasible. A more desirable approach is to have a single network that can make (a) early predictions for compute efficient systems and (b) latter predictions to obtain high accuracy. A multi-stage restoration model naturally offers such functionalities.

Table 7 reports the stage-wise results of our multi-stage approach. Our MPRNet demonstrates competitive restoration performance at each stage. Notably, our stage-1 model is light, fast, and yields better results than other sophisticated algorithms such as SRN and DeblurGAN-v2 . Similarly, when compared to a recent method DMPHN , our stage-2 model shows the PSNR gain of 0.51 dB while being more resource-efficient ( $\sim$ ${2\times}$ fewer parameters and ${13\times}$ faster).

Conclusion

In this work, we propose a multi-stage architecture for image restoration that progressively improves degraded inputs by injecting supervision at each stage. We develop guiding principles for our design that demand complementary feature processing in multiple stages and a flexible information exchange between them. To this end, we propose contextually-enriched and spatially accurate stages that encode a diverse set of features in unison. To ensure synergy between reciprocal stages, we propose feature fusion across stages and an attention guided output exchange from earlier stages to the later ones. Our model achieves significant performance gains on numerous benchmark datasets. In addition, our model is light-weighted in terms of model size and efficient in terms of runtime, which are of great interest for devices with limited resources.

Acknowledgments. M.-H. Yang is supported in part by the NSF CAREER Grant 1149783. Special thanks to Kui Jiang for providing image deraining results.