All-in-one Multi-degradation Image Restoration Network via Hierarchical Degradation Representation

Image restoration is a fundamental task in computer vision, which aims to recover the degraded images to their original high-quality versions, including deblurring, denoising, inpainting, low-light enhancement, and so on. Traditional approaches focus on the exploration of the image prior, such as sparse (Luo et al., 2015; Mairal et al., 2007; Xu et al., 2013), low-rank (Gu et al., 2014; Xu et al., 2017), self-similarity (Dabov et al., 2007) etc..

Recently, with the support of a large number of collected paired images, many deep neural networks (DNN) methods (Abuolaim and Brown, 2020; Liang et al., 2021; Zamir et al., 2021; Chen et al., 2022a; Wei et al., 2018; Chen et al., 2021a; Zhang et al., 2020; Yan et al., 2023, 2019) have produced impressive results on each subtask of restoration. Those works emphasize the design of network architecture and loss functions. By utilizing images that are usually captured and categorized manually by humans according to degradation type, DNN-based methods focus on learning the implicit mapping between the distorted image and the high-quality image. Although some models (Zamir et al., 2021; Liang et al., 2021; Chen et al., 2022a, 2021a) can be adapted to handle multiple types of degraded images, they typically require separate training on specific degraded datasets, and they may not generalize well to other typed of degraded images without further adaptation. For instance, MPR (Zamir et al., 2021) trained well for image denoising has limited performance on image deblurring, which is not expected in practice. Therefore, it is vital to consider the fact that images are often corrupted by multiple degradations and devise all-in-one solutions to meet the requirement of real-world multi-task processing.

Currently, there has been increased interest in developing all-in-one models that can handle various degraded images in a single network after being trained. For instance, (Li et al., 2020) proposes a network with multiple encoders to process each degradation using a specific encoder. Similarly, Transweather (Valanarasu et al., 2022) introduces a decoder with learnable embeddings in a transformer architecture to address multiple degradation types.

Instead of only focusing on designing network structures, there are also techniques that incorporate contrastive learning to enable the network to differentiate between different types of corruptions and handle multi-degraded images. One such example is presented in (Chen et al., 2022b), which suggests using both soft and hard contrastive regularization to enhance the performance of both specific and multiple degradations. Other approaches, DASR (Wang et al., 2021) and AirNet (Li et al., 2022b) consider patches from the same image as positive samples and patches from different images as negative samples. However, this selection of positive and negative samples may not always be adequate, since it does not take into account situations where different degradations are related or where different images may belong to the same degradation, which could result in representation drift and performance drop of restoration network.

The representation of degradation is a crucial step in image restoration and serves as a prerequisite for accurately restoring a degraded image. Generally, the type of degradation is known before the estimation and the main task is to estimate the parameters of the degradation model. For example, in denoising (Liu et al., 2013; Guo et al., 2019), the core is to estimate the noise level when the noise type is known. And in deblurring, many methods (Gong et al., 2017; Pan et al., 2016; Chen et al., 2021b; Hu et al., 2014; Yan et al., 2017) usually estimate the blur kernel before non-blind deblurring. Additionally, there exist approaches that implicitly represent the degradation by learning a feature vector that serves as a proxy for the degradation and is subsequently fed into the restoration method. DASR (Wang et al., 2021) and AirNet (Li et al., 2022b) realize the learning of representation by contrastive learning. (Li et al., 2022c) proposes to learn degradation representations with a blurry-sharp cycle framework. (Li et al., 2022a) proposes to learn a latent representation space for degradations in super-resolution. Our method falls into this category, in which we guide the network to restore clear images by constructing hierarchical degradation representations of degraded images and adapting to different degraded image features.

Our solution is a two-stage All-in-One approach, the first stage is designed to construct the hierarchical representation of degraded images, and the second stage is to remove degradation artifacts to produce a high-quality image with the supervision of degradation representation. The overview of our method is depicted in Figure 2. Despite serving different purposes, the two stages share a common network architecture. The network we propose comprises of two main components: a Degradation representation sub-Network (DRN) and a restoration sub-network (RN). Initially, a corrupted image $I$ is fed into the encoder of the degradation representation sub-network to extract its features $z$, which are then transformed by projectors to yield the hierarchical degradation representation $r$ in latent space. Subsequently, the degradation representation $r$ is integrated into the restoration via the newly proposed feature transform block. With the assistance of the degradation representation $r$, our network is able to recover a high-quality image from the input distorted image.

As mentioned above and shown in Figure 1, the degradation has a characteristic of hierarchical subordination. The characteristic indicates that the images with adverse distortions can be classified into different categories in each layer of the hierarchy. An instance could be a motion-blurred image and a defocused image both categorized under ”blur”, however, their classification in the subgroup of blur types is different. To model the commonality and the distinction between multiple distortions, naturally, we propose a novel tree-structured representation, which corresponds to the categories of the hierarchical structure of degradations. The tree-structured representation $T$ can be flattened into a vector when performing a level-order traversal of the hierarchy, which is formulated as follows:

where $L$ and $c_{L}$ denote the number of levels in the tree structure and the number of nodes in each level, respectively. The binary value $v_{i,j}$ in $j_{th}$ node of $i_{th}$ level represents whether the image belongs to that node or not. In our method, we predefine the structure of the tree as a 4-layer binary tree, so that it can be used consistently during both training and testing. After establishing the specific form of the hierarchical representation, we need to take into account how to construct the representation for each corrupted image without additional degradation information and learn it with a degradation representation network.

Representation Construction. We recommend a progressive strategy to construct the degradation representation from the top layer to the bottom layer in the hierarchy, as described in Algorithm 1. Given the corrupted images $I$, clear images $Y$, the number of hierarchy $L$, and the number of clusters in each layer $c_{1},c_{2},...,c_{L}$, our goal is to build a tree-structured representation $T$. The degradation representation of each input image is initially assigned to the root node and updated with the increasing depth of the tree during the outer loop. In each iteration of a layer, a node is sequentially selected as the current node, and the samples belonging to the current node are located to obtain their degradation representation using DRN. After clustering, the clustering results are concatenated with the original representations to update the degradation representations.

Representation Learning. To learn the degradation representation and facilitate usage during testing, a sub-network DRN is employed to extract degradation features and transform them into a low-dimensional embedding space under the supervision of the hierarchical tree-structured representation. Our DRN includes an encoder and two parallel projectors. The encoder $\mathcal{F}$ is composed of multiple convolutional layers with residual connections, activation layers, and pooling layers. The encoder’s objective is to extract features about degradations from degraded image $I$, and the process is expressed as $z=\mathcal{F}(I)$. The projectors are based on MLPs, which are composed of fully-connected layers and activation functions. One of the projectors, named mask projector $\mathcal{P}_{m}$, is designed to predict the binary mask $r_{m}=\mathcal{P}_{m}(z)$. The length of vector $r_{m}$ is always consistent with the current clustering layer in the construction process of tree-structured representation and will increase with the number of nodes in the degradation tree.

Given a degraded image $I$ and its degradation $r$, we calculate Cross Entropy loss to optimize the parameters of the encoder $\mathcal{F}$ and the mask projector $\mathcal{P}_{m}$:

And the other parallel projector, named attribute projector $P_{a}$ is only used in the first training stage, its output has the same size as the output of the first projector. The value on each vector dimension represents the attribute value at the corresponding node. In the first training stage, the degradation representation can be described as a product of mask $r_{m}$ and attribute value $r_{a}=\mathcal{P}_{a}(z)$ :

Feature Transform Block. To leverage the hierarchical degradation representation in image restoration, we introduce a feature transform block to modulate the image feature transformation in latent feature space according to the distortion information. The structure of this block is depicted in Figure 3. Inspired by Domain-Specific Batch Normalization (DSBN) (Chang et al., 2019) in domain adaptation, which transforms domain-specific information into domain-invariant representation using the parameters of BN, we employ degradation-specific parameters in layer normalization (LN) to refine the degraded image features. LN is a widely used technique in restoration networks, which is expressed as

where $\mu$ and $\sigma$ denote the mean and standard deviation of the image feature $x\in\mathbb{R}^{C\times H\times W}$, and $\epsilon$ is a small constant to avoid dividing by zero, and $\gamma$ and $\beta$ are learnable affine parameters. In order to make the layer normalization adapt to various degraded image features, we propose a Degradation-Specific Layer Normalization (DSLN). Formally, DSLN allocates degradation-specific affine parameters $\gamma(r)$ and $\beta(r)$ for different degradations.

Hence, DSLN can be written as:

where $\gamma(r)$ and $\beta(r)$ are connected with the degradation representation $r$ by using linear transformation matrixes, which is realized by a fully-connection layer in the implementation,i.e., $\gamma(r)=Wr+b$.

Moreover, we incorporate a Gating Mechanism (GM) to activate the channels of image features according to degradation. The gating mechanism is formulated as the element-wise product of image feature $x$ and $\phi(r)$, and $\phi(r)=GELU(W^{1}r+b^{1})(W^{2}r+b^{2})$. The gating mechanism controls the information flow with the guidance of degradation representation $r$, thereby allowing the network to focus on the degradation-specific channels. Overall, with DSLN and gating mechanism, the FTB has the capability to align degradation domains and refine the image features under the guidance of degradation representation $r$, allowing the restoration network to produce high-quality images.

Training and Loss. Our restoration network is constructed based on NAFNet (Chen et al., 2022a), which is a variant Unet with symmetric encoder-decoder architecture. In the encoder of our restoration network, the FTB serves as the fundamental module and modifies the degradation-related image features. In the decoder, efficient NAFBlocks (Chen et al., 2022a) are employed to transform modulated image features. The encoder features are concatenated with the decoder features via skip connections to avoid the gradient vanishing. Finally, a convolution layer is applied to generate a residual image. The specific structure of our restoration network can be referred to in Figure 2.

The restoration network takes the degraded image $I$ and degradation representation $r$ as input and performs adaptive processing to generate a high-quality image $R$. In the training phase, to optimize the parameters in the network, we adopt two commonly used losses, smoothL1 (Girshick, 2015) and SSIM loss, which can measure the discrepancy of the restored result $R$ and ground truth image $Y$ at the pixel-wise level and patch-wise level, respectively. The optimization objective of our restoration network is the combination of two losses, which is formulated as follows:

where $\alpha$ is a hyper-parameter to balance the two losses. In the first stage, the DRN is required to be optimized by cross-entropy loss, therefore, the total loss in our method can be summarized as follows:

In the second stage, the degradation representation $r$ is produced by detached DRN, and we focus on the optimization of the restoration network. Hence, the total loss in the second stage is $\mathcal{L}_{res}$.

Datasets. As our work primarily focuses on multi-degradation restoration, the dataset used should contain adverse types of degraded images. For this purpose, we train our network on a combination of different degradation datasets, including RED4(Nah et al., 2019) for deblurring and jpeg-compression removal, SIDD(Abdelhamed et al., 2018) for denoising, LOL (Wei et al., 2018) for low-light enhancement, and DPDD (Abuolaim and Brown, 2020) for defocus deblurring. The training data consists of 1920 images uniformly sampled from the four datasets. And we randomly sample 539 images for testing to validate the effectiveness of our method. Moreover, to measure the network’s performance on images with different degradation levels, we follow (Li et al., 2022b) to train our network on WED (Ma et al., 2016) and test it on synthetic noisy images with different noise levels from CBSD68 (Martin et al., 2001).

Implementation Details. We adopt NAFNet (Chen et al., 2022a) as our restoration backbone, as it has demonstrated remarkable performance and computational advantages across multiple restoration tasks. We implement our approach using Pytorch framework and train the network on two NVIDIA A100 GPUs in a distributed manner. An AdamW (Loshchilov and Hutter, 2017) optimizer is adopted to optimize network parameters. The learning rate is initialized to 5e-4 and decreased with the CosineAnnealingLR decay strategy. The network is trained for 600 epochs. The batch size is set to 28 and the patch size is 256. In the first stage, the clustering algorithm is conducted every 150 epochs.

Compared Methods. In this section, we evaluate our method with the state-of-the-art methods of image restoration. Initially, we selected some restoration approaches that are tailored for individual restoration tasks, including deblurring, denoising, low-light image enhancement, and defocus. Those methods like MPR (Zamir et al., 2021), NAFNet (Chen et al., 2022a), which are CNN-based methods, and SwinIR (Liang et al., 2021) is a transformer-based method. Besides, we also compare the performance of our method with other all-in-one approaches, like TransWeather (Valanarasu et al., 2022) and AirNet (Li et al., 2022b). For a fair competition with the compared methods, we follow them to conduct experiments with default optimal hyperparameters and settings. During the experiments, all the comparative methods were trained and tested on multi-degraded datasets in order to explore and compare their performance in handling multi-degraded images.

Quantitative Comparison. In this comparison, we adopt three commonly used in image restoration reference-based image quality assessment metrics, PSNR, SSIM, and LPIPS (Zhang et al., 2018a) to evaluate the quality of restored images. Better results are indicated by higher values of PSNR and SSIM, while LPIPS is the opposite. During the test, we calculate the metrics not only on all testing degraded images but also on each type of degradation according to the dataset type. As illustrated in Table 1, our method performs better than other methods in terms of the average metrics on all types of degraded data. Additionally, when classified by types of degradation, our method achieves the best or second-best results on different types of degraded images, suggesting that the performance improvement of our method can effectively distinguish and represent each type of degradation, rather than relying on the improvement on a single type of degradation. We also measured the computational costs of different models using a 3x256x256 image as input. It is evident that our method exhibits significantly reduced computational cost compared to AirNe t(Li et al., 2022b), MPR (Zamir et al., 2021), and SwinIR (Liang et al., 2021). While there is an increase in computational cost compared to NAFNet (Chen et al., 2022a) and TransWeather (Valanarasu et al., 2022), our method holds an advantage in handling multiple degradations.

Visual Comparison. We conduct a visual comparison of our method and several state-of-the-art techniques on the aforementioned types of degraded images. Figure4 presents a visual comparison against other image restoration methods on RED4 (Nah et al., 2019) and DPDD (Abuolaim and Brown, 2020) datasets, and it demonstrates that our method is capable of restoring the intricate details of the image while the image content is corrupted seriously by blur. Meanwhile, the results on LOL (Wei et al., 2018) dataset are shown in Figure 5, our method effectively removes artifacts from low-light images when enhancing them, ensuring that the processed images have good visual quality. Figure 6 is the comparison of denoising results on SIDD (Abdelhamed et al., 2018) dataset, our approach can remove the noise in distorted images and recover a high-quality image. There are more results about the multi-degradation experiment, which can be found in the supplement material.

To evaluate our network’s performance on distorted images with different degradation levels, we compare the denoising results of our method and other state-of-the-art denoising methods, including BM3D (Dabov et al., 2007), DnCNN (Zhang et al., 2017a), IRCNN (Zhang et al., 2017b), DL (Fan et al., 2019), FFDNet (Zhang et al., 2018b), and MPR (Zamir et al., 2021). The noisy images are synthetic by adding White Gaussian noise with different levels (i.e. $\sigma=15,25,50$) to clear images. The comparison result is reported in Table 2. Our approach exhibits superior performance compared to other denoising methods and delivers optimal results in quantitative evaluations. This highlights the versatility of our method in restoring images degraded at varying levels.

Impact of Supervision on Degradation Types. To investigate the impact of degradation category supervision on the restoration network, we retrain some separate networks. As listed in Figure 7, the backbone model is a restoration network that does not utilize any degradation information. And backbone+CN is a restoration network that uses a classification network trained with degradation-type supervision and utilizes the classification results. The comparative results in Figure 7 show that the performance of networks that use degradation representation for processing multiple degradations has been significantly improved. However, even with the supervision of degradation types in the classification network, it cannot surpass our proposed method of using hierarchical degradation representation for restoration. This implies that degradation types are not mutually exclusive in multi-degraded image restoration, and exploring the relationships between degradations and leveraging them for restoration is highly necessary.

Impact of Hierarchical Representation Layers. We evaluate the performance of networks with different layers in hierarchical degradation representation. The visual results are shown in Figure 8. From the results shown in Table 3, it can be observed that as the number of layers in the tree structure increases, the restoration performance also improves with the enhancement of representation capacity. When using the two-stage strategy, the degradation representation provided by the fixed DRN allows the training to focus on RN. As a result, the performance of the restoration network is further improved.

Impact of FTB. We also conduct an ablation study to validate the effectiveness of FTB in our network. We remove FTB, DSLN, and GM separately in our model, to observe the effect of each component. The results are shown in Table 3. Removing either DLSN or GM leads to a performance drop in our network. However, as both have the ability to refine degraded image features by utilizing degradation information, the absence of one does not hinder the other from being effective. FTB’s ability to modulate image features is maximized when both are used.

We randomly sample degraded images from the training set and extract their degradation features through DRN. By utilizing t-SNE on the degradation representations, we can evaluate the performance of degradation representation in low-dimension embedding space. The visualized result is presented in Figure 8. From the illustration, we can see that the hierarchical degradation representations generated by our method exhibit greater inter-class separation and intra-class compactness compared to the results of AirNet (Li et al., 2022b). This result is consistent with our previous conjecture that tighter constraints on samples with similar degradation types can help improve the performance of the restoration network.