Learning Enriched Features for Fast Image Restoration and Enhancement

Introduction

Owing to the physical limitations of cameras or due to complicated lighting conditions, image degradations of varying severity are often introduced as part of image acquisition. For instance, smartphone cameras come with a narrow aperture and have small sensors with limited dynamic range. Consequently, they frequently generate noisy and low-contrast images. Similarly, images captured under the unsuitable lighting are either too dark or too bright. Image restoration aims to recover the original clean image from its corrupted measurements. It is an ill-posed inverse problem, due to the existence of many possible solutions.

Recent advances in image restoration and enhancement have been led by deep learning models, as they can learn strong (generalizable) priors from large-scale datasets. Existing CNNs typically follow one of the two architecture designs: 1) an encoder-decoder, or 2) high-resolution (single-scale) feature processing. The encoder-decoder models first progressively map the input to a low-resolution representation, and then apply a gradual reverse mapping to the original resolution. Although these approaches learn a broad context by spatial-resolution reduction, on the downside, the fine spatial details are lost, making it extremely hard to recover them in the later stages. On the other hand, the high-resolution (single-scale) networks do not employ any downsampling operation, and thereby recover better spatial details. However, these networks have limited receptive field and are less effective in encoding contextual information.

Image restoration is a position-sensitive procedure, where pixel-to-pixel correspondence from the input image to the output image is needed. Therefore, it is important to remove only the undesired degraded image content, while carefully preserving the desired fine spatial details (such as true edges and texture). Such functionality for segregating the degraded content from the true signal can be better incorporated into CNNs with the help of large context, e.g., by enlarging the receptive field. Towards this goal, we develop a new multi-scale approach that maintains the original high-resolution features along the network hierarchy, thus minimizing the loss of precise spatial details. Simultaneously, our model encodes multi-scale context by using parallel convolution streams that process features at lower spatial resolutions. The multi-resolution parallel branches operate in a manner that is complementary to the main high-resolution branch, thereby providing us more precise and contextually enriched feature representations.

One main distinction between our method and the existing multi-scale image processing approaches is how we aggregate contextual information. The existing methods process each scale in isolation. In contrast, we progressively exchange and fuse information from coarse-to-fine resolution-levels. Furthermore, different from existing methods that employ a simple concatenation or averaging of features coming from multi-resolution branches, we introduce a new selective kernel fusion approach that dynamically selects the useful set of kernels from each branch representations using a self-attention mechanism. More importantly, the proposed fusion block combines features with varying receptive fields, while preserving their distinctive complementary characteristics.

The main contributions of this work include:

A novel feature extraction model that obtains a complementary set of features across multiple spatial scales, while maintaining the original high-resolution features to preserve precise spatial details (Sec. 3).

A regularly repeated mechanism for information exchange, where the features from coarse-to-fine resolution branches are progressively fused together for improved representation learning (Sec. 3.1).

A new approach to fuse multi-scale features using a selective kernel network that dynamically combines variable receptive fields and faithfully preserves the original feature information at each spatial resolution (Sec. 3.1.1).

A preliminary version of this work has been published as a conference paper . The MIRNet model is expensive in terms of size and speed. In this work, we make several key modifications to MIRNet that allow us to significantly reduce the computational cost while enhancing model performance (see Table I). Specifically, in the proposed MIRNet-v2 , (a) We demonstrate feature fusion only in the direction from low- to high-resolution streams performs best, and the information flow from high- to low-resolution branches can be removed to improve efficiency. (b) We replace the dual attention unit with a new residual contextual block (RCB). Furthermore, we introduce group convolutions in RCB that are capable of learning unique representations in each filter group, while being more resource efficient than standard convolutions. (c) We employ progressive learning to improve training speed: the network is trained on small image patches in the early epochs and on gradually large patches in the later training epochs. (d) We show the effectiveness of the proposed design on a new task of dual-pixel defocus deblurring alongside the other image processing tasks of image denoising, super-resolution and image enhancement. Our MIRNet-v2 achieves state-of-the-results on all six datasets. Furthermore, we extensively evaluate our approach on practical challenges, such as generalization ability across datasets (Sec. 4)

In Table I, we compare MIRNet-v2 with MIRNet under the same training and inference settings. The results show that MIRNet-v2 is more accurate (improving PSNR from 39.72 dB to 39.84 dB), while reducing the number of parameters and FLOPs by $\sim 81\%$, convolutions by $36\%$, and activations by $69\%$. Furthermore, the training and inference speed is increased by $2.2\times$ and $3.6\times$, respectively.

Related Work

Rapidly growing image content necessitates the need to develop effective image restoration and enhancement algorithms. In this paper, we propose a new method capable of performing dual-pixel defocus deblurring, image denoising, super-resolution, and image enhancement. Unlike existing works for these problems, our approach processes features at the original resolution in order to preserve spatial details, while effectively fuses contextual information from multiple parallel branches. Next, we briefly describe the representative methods for each of the studied problems.

Images captured with wide camera aperture have shallow depth of field (DoF), where the scene regions that lie outside the DoF are out-of-focus. Given an image with defocus blur, the goal of defocus deblurring is to generate an all-in-focus image. Existing defocus deblurring approaches either directly deblur images , or first estimate the defocus dispartiy map and then use it to guide the deblurring procedure . Modern cameras are equipped with dual-pixel sensor that has two photodiodes at each pixel location, thereby generating two sub-aperture views. The phase difference between these views is useful in measuring the amount of defocus blur at each scene point. Recently, Abuolaim et al. presented a dual-pixel deblurring dataset (DPDD) and a new method based on encoder-decoder design. In this paper our focus is also on deblurring images directly using the dual-pixel data as in . Previous defocus deblurring works employ the encoder-decoder that repeatedly uses the downsampling operation, thus causing significant fine detail loss. Whereas the architectural design of our approach enables preservation of desired textural details in the restored image.

2 Image Denoising

Classic denoising methods are mainly based on modifying transform coefficients or averaging neighborhood pixels . Although the classical approaches perform well, the self-similarity based algorithms, e.g., NLM and BM3D , demonstrate promising denoising performance. Numerous patch-based schemes that exploit redundancy (self-similarity) in images are later developed . Recently, deep learning models make significant advances in image denoising, yielding favorable results than those of the hand-crafted methods.

3 Image Super-Resolution

Prior to the deep-learning era, numerous super-resolution (SR) algorithms have been proposed based on the sampling theory , edge-guided interpolation , natural image priors , patch-exemplars and sparse representations . Currently, deep-learning techniques are being actively explored as they provide dramatically improved results over conventional algorithms. The data-driven SR approaches differ according to their architecture designs . Early methods take a low-resolution (LR) image as input and learn to directly generate its high-resolution (HR) version. In contrast to directly producing a latent HR image, recent SR networks employ the residual learning framework to learn the high-frequency image detail, which is later added to the input LR image to produce the final result. Other networks designed to perform SR include recursive learning , progressive reconstruction , dense connections , attention mechanisms , multi-branch learning , and generative adversarial networks (GANs) .

4 Image Enhancement

Oftentimes, cameras generate images that lack vivid details or contrast. A number of factors contribute to the low quality of images, including unsuitable lighting conditions and physical limitations of camera devices. For image enhancement, histogram equalization is the most commonly used approach. However, it frequently produces under- or over-enhanced images. Motivated by the Retinex theory , several enhancement algorithms mimicking human vision have been proposed in the literature . Recently, CNNs have been successfully applied to general, as well as low-light, image enhancement problems . Notable works employ Retinex-inspired networks , encoder-decoder networks , and GANs .

Proposed Method

A schematic of the proposed MIRNet-v2 is shown in Fig. 1. We first present an overview of the proposed MIRNet-v2 for image restoration and enhancement. We then provide details of the multi-scale residual block, which is the fundamental building block of our method, containing several key elements: (a) parallel multi-resolution convolution streams for extracting (fine-to-coarse) semantically-richer and (coarse-to-fine) spatially-precise feature representations, (b) information exchange across multi-resolution streams, (c) attention-based aggregation of features arriving from different streams, and (d) residual contextual blocks to extract attention-based features.

where $\mathbf{I}^{*}$ denotes the ground-truth image, and $\varepsilon$ is a constant which we empirically set to $10^{-3}$ for all the experiments.

To encode context, existing CNNs typically employ the following architecture design: (a) the receptive field of neurons is fixed in each layer/stage, (b) the spatial size of feature maps is gradually reduced to generate a semantically strong low-resolution representation, and (c) a high-resolution representation is gradually recovered from the low-resolution representation. However, it is well-understood in vision science that in the primate visual cortex, the sizes of the local receptive fields of neurons in the same region are different . Therefore, a similar mechanism of collecting multi-scale spatial information in the same layer is more effective when incorporated with in CNNs . Motivated by this, we propose the multi-scale residual block (MRB), as shown in Fig. 1. It is capable of generating a spatially-precise output by maintaining high-resolution representations, while receiving rich contextual information from low-resolutions. The MRB consists of multiple (three in this paper) fully-convolutional streams connected in parallel that operate on varying resolution feature maps (ranging from low to high). It allows contextualized-information transfer from the low-resolution streams to consolidate the high-resolution features. Next, we describe the individual components of MRB.

One fundamental property of neurons present in the visual cortex is their ability to change receptive fields according to the stimulus . This mechanism of adaptively adjusting receptive fields can be incorporated in CNNs by using multi-scale feature generation (in the same layer) followed by feature aggregation and selection. The most commonly used approaches for feature aggregation include simple concatenation or summation. However, these choices provide limited expressive power to the network, as reported in . In MRB, we introduce a nonlinear procedure for fusing features coming from different resolution streams using a self-attention mechanism. Motivated by , we call it selective kernel feature fusion (SKFF).

1.2 Residual Contextual Block

While the SKFF block fuses information across multi-resolution branches, we also need a distillation mechanism to extract useful information from within a feature tensor. Motivated by the advances of recent low-level vision methods which incorporate attention mechanisms , we propose the residual contextual block (RCB) to extract features in the convolutional streams. The schematic of RCB is shown in Fig. 3. The RCB suppresses less useful features and only allows more informative ones to pass further. The overall process of RCB is summarized as:

2 Progressive Training Regime

When considering the image patch size for network training, there is a trade-off between the training speed and test-time accuracy . On large patches, CNNs capture fine image details to provide improved results, but they are slower to train. Whereas, training on small image patches is faster, but comes at the cost of accuracy drop. To strike the right balance between the training speed and accuracy, we propose a progressive learning method where the network is trained on smaller image patches in the early epochs and on gradually larger patches in the later training epochs. This approach can also be understood as a curriculum learning process where the network sequentially moves from learning a simpler task to a more complex one (where modeling of fine details is required). The progressive learning strategy on mixed-size image patches not only improves the training speed but also enhances the model performance at test time where the input images can be of different sizes (which is common in image restoration problems).

Experiments

In this section, we perform qualitative and quantitative assessments of the results produced by our MIRNet-v2 and compare it with the state-of-the-art methods. Next, we describe the datasets, and then provide the implementation details. Finally, we report results for (a) dual-pixel defocus deblurring, (b) image denoising, (c) image super-resolution and (d) image enhancement, on six real image datasets.

Dual-pixel defocus deblurring. DPDD dataset contains 500 indoor/outdoor scenes captured with a DSLR camera. Each scene consists of two defocus blurred sub-aperture views captured with a wide camera aperture, and the corresponding all-in-focus ground truth image captured with a narrow aperture. The DDPD dataset is divided into 350 images for training, 74 images for validation and 76 images for testing.

Image denoising. (1) DND consists of $50$ images captured with four consumer cameras. Since the images are of very high-resolution, the dataset providers extract $20$ crops of size $512\times 512$ from each image, yielding $1000$ patches in total. All these patches are used for testing (as DND does not contain training or validation sets). The ground-truth noise-free images are not released publicly, therefore the image quality scores in terms of PSNR and SSIM can only be obtained through an online server . (2) SIDD is collected with smartphone cameras. Due to the small sensor and high-resolution, the noise levels in smartphone images are much higher than those of DSLRs. SIDD contains $320$ image pairs for training and $1280$ for validation.

Super-resolution. RealSR contains real-world LR-HR image pairs of the same scene captured by adjusting the focal-length of the cameras. RealSR has both indoor and outdoor images taken with two cameras. The number of training image pairs for scale factors $\times 2$, $\times 3$ and $\times 4$ are $183$, $234$ and $178$, respectively. For each scale factor, $30$ test images are also provided in RealSR.

Image enhancement. (1) LoL is created for low-light image enhancement problem. It provides 485 images for training and 15 for testing. Each image pair in LoL consists of a low-light input image and its corresponding well-exposed reference image. (2) MIT-Adobe FiveK contains $5000$ images of various indoor and outdoor scenes captured with DSLR cameras in different lighting conditions. The tonal attributes of all images are manually adjusted by five different trained photographers (labelled as experts A to E). Similar to , we also consider the enhanced images of expert C as the ground-truth. Moreover, the first 4500 images are used for training and the last 500 for testing.

2 Implementation Details

The proposed architecture is end-to-end trainable and requires no pre-training of sub-modules. We train four different networks for four different restoration tasks. For the dual-pixel defocus deblurring, we concatenate the left and right sub-aperture images and feed them as input to the network. The training parameters, common to all experiments, are the following. We use 4 RRGs, each of which further contains $2$ MRBs. The MRB has $3$ parallel streams with channel dimensions of $80,120,180$ at resolutions $1,\frac{1}{2},\frac{1}{4}$, respectively. Each stream in MRB has $2$ RCBs with shared parameters. The models are trained with the Adam optimizer ($\beta_{1}=0.9$, and $\beta_{2}=0.999$) for $3\times 10^{5}$ iterations. The initial learning rate is set to $2\times 10^{-4}$. We employ the cosine annealing strategy to steadily decrease the learning rate from initial value to $10^{-6}$ during training. For progressive training, we use the image patch sizes of 128, 144, 192, and 224. The batch size is set to $64$ and, for data augmentation, we perform horizontal and vertical flips.

3 Dual-Pixel Defocus Deblurring

We compare the performance of the proposed MIRNet-v2 with the conventional defocus deblurring methods (EBDB and JNB ) as well as the learning-based approaches (DMENet , DPDNet , and RDPD ). Table II shows that our method achieves state-of-the-art results for both the indoor and outdoor scene categories. In particular, our MIRNet-v2 achieves 0.86 dB PSNR improvement over the previous best method RDPD on indoor images and 0.77 dB on outdoor images. When both scene categories are combined, our method shows performance gains of 0.81 dB over RDPD and 1.07 dB over the second best method DPDNet .

In Fig. 4, we provide defocus-deblurred results produced by different methods for both indoor and outdoor scenes. It is noticeable that our method effectively removes the spatially varying defocus blur and produces images that are more sharper and visually faithful to the ground-truth than those of the compared approaches.

4 Image Denoising

In this section, we demonstrate the effectiveness of the proposed MIRNet-v2 for image denoising. We train our network only on the training set of the SIDD and directly evaluate it on the test images of both SIDD and DND datasets. Quantitative comparisons in terms of PSNR and SSIM metrics are summarized in Table III. Our MIRNet-v2 performs favourably against the data-driven, as well as conventional, denoising algorithms. Specifically, when compared to the recent best methods, our algorithm demonstrates a performance gain of $0.32$ dB over CycleISP on SIDD and $0.11$ dB over DAGL on DND. Furthermore, it is worth noting that CycleISP uses additional training data, yet our method yields considerably better results.

Fig. 5 shows a visual comparisons of our results with those of other competing algorithms. The MIRNet-v2 is effective in removing real noise and produces perceptually-pleasing and sharp images. Moreover, it is can maintain the spatial smoothness of the homogeneous regions without introducing artifacts. In contrast, most of the other methods either yield over-smooth images and thus sacrifice structural content and fine textural details, or produce images with chroma artifacts and blotchy texture.

Generalization capability. The DND and SIDD datasets are acquired with different sets of cameras having different noise characteristics. Since the DND benchmark does not provide training data, setting a new state-of-the-art on DND with our SIDD trained network indicates the good generalization capability of our approach.

5 Super-Resolution

We compare our MIRNet-v2 against the state-of-the-art SR algorithms (VDSR , SRResNet , RCAN , LP-KPN ) on the testing images of the RealSR for upscaling factors of $\times 2$, $\times 3$ and $\times 4$. Note that all the benchmarked algorithms are trained on the RealSR dataset for a fair comparison. In the experiments, we also include bicubic interpolation , which is the most commonly used method for generating super-resolved images. Here, we compute the PSNR and SSIM scores using the Y channel (in YCbCr color space), as it is a common practice in the SR literature . The results in Table IV show that the bicubic interpolation provides the least accurate results, thereby indicating its low suitability for dealing with real images. Moreover, the same table shows that the recent method LP-KPN achieves marginal improvement of only $\sim 0.04$ dB over the previous best method RCAN . In contrast, our method significantly advances state-of-the-art and consistently achieves better image quality scores than other approaches for all three scaling factors. Particularly, compared to LP-KPN , our method leads to performance gains of $0.48$ dB, $0.73$ dB, and $0.24$ dB for scaling factors $\times 2$, $\times 3$ and $\times 4$, respectively. The trend is similar for the SSIM metric as well.

Visual comparisons in Fig. 6 show that our MIRNet-v2 can effectively recover content structures . In contrast, VDSR , SRResNet and RCAN reproduce results with noticeable artifacts. Furthermore, LP-KPN is not able to preserve structures (see near the right edge of the crop). Several more examples are provided in Fig. 7 to further compare the image reproduction quality of our method against the previous best method . It can be seen that LP-KPN has a tendency to over-enhance the contrast (cols. 1, 3, 4) and in turn causes loss of details near dark and high-light areas. In contrast, the proposed MIRNet-v2 successfully reconstructs structural patterns and edges (col. 2) and produces images that are natural (cols. 1, 4) and have better color reproduction (col. 5).

6 Image Enhancement

In this section, we demonstrate the effectiveness of our algorithm by evaluating it for the image enhancement task. We report PSNR/SSIM values of our method and several other techniques in Table V and Table VI for the LoL and MIT-Adobe FiveK datasets, respectively. It can be seen that our MIRNet-v2 achieves significant improvements over previous approaches. Notably, when compared to the recent best methods, MIRNet-v2 obtains $3.44$ dB performance gain over KinD++ on the LoL dataset and $0.93$ dB improvement over DeepUPENote that the quantitative results reported in are incorrect. The correct scores are later released by the original authors [link]. on the Adobe-Fivek dataset.

We show visual results in Fig. 8 and Fig. 9. Compared to other techniques, our method generates enhanced images that are natural and vivid in appearance and have better global and local contrast.

7 Ablation Studies

We study the impact of each of our architectural components and design choices on the final performance. All the ablation experiments are performed for the super-resolution task with $\times 3$ scale factor. The ablation models are trained on image patches of size $128<span class="katex-display"><span class="katex"><span class="katex-mathml"><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></span><span class="katex-html" aria-hidden="true"><span class="base"><span class="strut" style="height:0.6667em;vertical-align:-0.0833em;"></span><span class="mord">×</span></span></span></span></span>128$ for $10^{5}$ iterations. Table VII shows that removing skip connections causes the largest performance drop. Without skip connections, the network finds it difficult to converge and yields high training errors, and consequently low PSNR. Furthermore, the information exchange among parallel convolution streams via SKFF is helpful and leads to improved performance. Similarly, RCB contributes positively towards the final image quality.

Table VIII shows that the proposed RCB provides favorable performance gain over the baseline Resblock from EDSR . Moreover, removing the transform part from RCB causes drop in accuracy. Table VIII also shows that replacing the group convolutions with regular convolutions in RCB increases the PSNR score, but at the cost of significant increase in parameters and FLOPs. Therefore, we opt for RCB with group convolutions (g=2) as a balanced choice.

Next, we analyze the feature aggregation strategy in Table IX. It shows that the proposed SKFF generates favorable results compared to summation and concatenation. Note that our proposed SKFF module uses $\sim 5\times$ fewer parameters than concatenation. Table X shows that the progressive learning strategy on mixed-size image patches yields PSNR similar to the model trained on large image patches (ps=224), but takes less time for training. Finally, in Table XI we study how the number of convolutional streams and columns (RCB blocks) of MRB affect the image restoration quality. We note that increasing the number of streams provides significant improvements, thereby justifying the importance of multi-scale features processing. Moreover, increasing the number of columns yields better scores, thus indicating the significance of information exchange among parallel streams for feature consolidation.

Concluding Remarks

Conventional image restoration and enhancement pipelines either stick to the full resolution features along the network hierarchy or use an encoder-decoder architecture. The first approach helps retain precise spatial details, while the latter one provides better contextualized representations. However, these methods can satisfy only one of the above two requirements, although real-world image restoration tasks demand a combination of both conditioned on the given input sample. In this work, we propose a novel architecture whose main branch is dedicated to full-resolution processing and the complementary set of parallel branches provides better contextualized features. We propose novel mechanisms to learn relationships between features within each branch as well as across multi-scale branches. Our feature fusion strategy ensures that the receptive field can be dynamically adapted without sacrificing the original feature details. Consistent achievement of state-of-the-art results on six datasets for four image restoration and enhancement tasks corroborates the effectiveness of our approach.

Acknowledgements

Ming-Hsuan Yang is supported by NSF CAREER grant 1149783. Ling Shao is is partially supported by the National Natural Science Foundation of China (grant no. 61929104). Munawar Hayat is supported by the ARC DECRA Fellowship DE200101100.

References