A Sub-band Approach to Deep Denoising Wavelet Networks and a Frequency-adaptive Loss for Perceptual Quality

In an $nxn$ discrete wavelet transform (DWT), a total of $n^{2}$ orthogonal wavelets are used to convolve an image and convolution results are downsampled by $n$. For example, in a $2x2$ DWT, $4$ wavelets are usually denoted as $f_{LL},f_{LH},f_{HL},f_{HH}$, where $L$ and $H$ denote low and high frequencies respectively. Then from a signal $x$, a sub-band is obtained as follows.

Due to the orthogonality property of the wavelets, the original signal $x$ can be recovered by an inverse discrete wavelet transform (IDWT).

Discrete wavelet transform (DWT) has been used for denoising in both traditional [21] and deep methods [18, 19]. A general advantage of DWT over Fourier Transform is its ability to keep both spatial and frequency information in the sub-bands.

For deep learning methods, DWT has another advantage: DWT both helps to increase the receptive field (due to downsampling) and at the same time prevents information loss (due to its invertible nature). In [18], DWT was utilized in a multi-level manner, where sub-bands after a DWT was concatenated in channel-dimension and passed through a convolutional layer. Resulting tensor is passed through another DWT and above process was repeated until desired level is achieved. In the inverse procedure, output of a convolutional layer is de-concatenated channel-wise and given to IDWT and this is repeated until the dimensions are back to image-level. In [19] a densely self-guided wavelet network was proposed that utilizes DWT following the same main approach as in [18].

All of the approaches utilizing DWT in neural networks, concatenate sub-bands in channel dimension, run convolutional layers, de-concatenate and then apply inverse DWT (IDWT) to go back to original image domain. We argue that the sub-band order information necessary for IDWT, may be lost after applying a convolution on concatenated sub-bands. Therefore, in Section 3.1, we propose to apply separate convolutional layers to each sub-band, thus keeping sub-band order information for IDWT. We find that this approach outperforms the common one.

L1 Loss is the most commonly used loss among deep learning based denoising methods [3, 1], as it has been shown to outperform L2 loss in terms of performance and convergence. Some works in [3, 1] have also employed a joint loss by combining L1 and MS-SSIM in order to achieve good results in both evaluations. Although distortion metrics such as L1, L2 and MS-SSIM are useful to objectively evaluate methods, they do not necessarily correlate well with visual quality which is the most important in the real use case. Especially, camera manufacturers have to take the perceptual quality into account as that is the single most important evaluation in terms of customer satisfaction.

Not only distortion and perception quality do not correlate well, moreover it was found in [7] that they actually go against each other in a non-invertible problem setting such as denoising. In other words, for a non-invertible problem there is a performance boundary that one cannot go beyond due to the nature of the problem, and among this boundary one can go either in favor of perception or distortion, and the other will suffer.

Although they are shown to outperform L2 also in terms of visual quality, L1 and MS-SSIM are still distortion metrics. In order to obtain a better visual quality, in Section 3.3, we propose an adaptive frequency loss that is dynamically changing according to frequency characteristics of image content and the performance of the neural network.

A residual dense block (RDB) combine properties of DenseNets and ResNets and was first proposed in a deep image super-resolution work. In an RDB bottleneck, several convolutional layers are densely connected and usually there is a single convolution at the end in order to convert to input channel dimension so that input and output can be added as in ResNets. In [26], RDB was proposed for image super-resolution and denoising tasks as a contiguous memory mechanism and was shown to outperform compared bottlenecks. In [22], RDB was used as the main bottleneck for a modified U-Net based denoising task. In [19], RDBs were used in a self-guided wavelet network for image denoising. In this paper we also choose RDB as our bottleneck, but we adapt the structure of RDB according to our proposed sub-band specialized convolutions as explained in Section 3.2.

As discussed in Section 2.1, image restoration neural networks that are using DWT, first apply the transform and then concatenate the generated sub-bands to a single tensor. Then, operations are performed on this single tensor. These operations can be convolutional blocks and for the sake of convenience, are simply denoted by ”blocks” in Fig. 1. Although the common usage of DWT in neural networks seems natural, we notice a potential problem which is explained as follows. The concatenation operation and subsequent convolutional blocks could potentially disrupt the ordering of the sub-bands as illustrated in upper branch of Fig. 1. Yet, for the de-concatenation which is necessary for IDWT – and thus going back to the original domain prior to DWT –, the sub-band ordering has to be preserved. Disrupting this ordering highly contradicts to the original motivation of using DWT: its invertible nature. We see this as a potential problem in learning, as a neural network is expected to learn also the ordering information in the convolutional filters during training.

In order to remedy this situation, we propose to simply apply separate convolutional blocks per each sub-band as shown in lower branch of Fig. 1. Separately applying convolutional layers to sub-bands discards inter-band communication, for example lower frequencies are processed separately without considering higher frequencies at all. Although this might seem like another potential problem, we avoid it with a simple trick in a specially designed bottleneck in Section 3.2.

In this section, we define the bottleneck block used in our method. First, we review Residual Dense Blocks (RDBs) as we will use a similar bottleneck. RDBs are widely used in many image restoration problems [26, 22, 19]. In RDB, the output of each operation (feature) is concatenated to the subsequent features as in DenseNets. At the end, there is a single convolution that operates on concatenation of all features. Moreover, the input to RDB is added to the output as in ResNets. In this work, we are going to use a variant of RDB utilizing our DWT approach, therefore we first take the dense block (DB) part which is RDB without the skip connection and is illustrated in Fig. 2.

We will transfer the last convolution and skip connection in RDB to our bottleneck block which is defined next. For convenience we will address our bottleneck as Residual Sub-band Dense Block (RSDB) block. DWT is applied on a tensor and 4 sub-bands are obtained. We assumed a DWT with 2x2 wavelets – hence 4 sub-bands –, but the idea can be easily extended to larger wavelets. Then, a dense block (DB) is applied on each sub-band. The output of DB for each sub-band is given to inverse DWT operation. Finally, a convolution is applied and the result is added to the input of the block. Besides converting densely concatenated features to input feature dimension, final convolution also addresses the no inter-band communication issue explained in Section 3.1. By applying IDWT and applying convolution afterwards, we make sure that the bottleneck also performs an operation on original domain, where all bands’ information is present. The Residual Sub-band Dense Block is illustrated in Fig. 3.

As explained in [7], for non-invertible problems there is a trade-off between the perceptual quality and distortion metric (such as PSNR), when it comes to extreme ends. L1 loss is the dominant loss function that is used in many deep image restoration works [3, 1]. Although advantage of L1 over L2 is proven in many previous studies, L1 is still largely a PSNR oriented loss.

Here we propose a loss that largely keeps the PSNR accuracy and also increase the perceptual quality. We first take the Discrete Cosine Transform (DCT) of both the ground truth $y_{true}$ and the prediction $y_{pred}$, in order to decompose them to their frequency-related coefficients. We chose DCT due to its computational simplicity. Also, DCT is commonly used in image compression as a rough measure of perceptual quality as in JPEG [14]. Next, we take the absolute error between each frequency coefficient of $y_{true}$ and $y_{pred}$. These errors are sorted in a descending order and the loss is calculated by taking the mean of the top $k\%$ in the sorted absolute errors as in Eq. 2.

In Eq. 2, $K$ is the set of elements of $DCT_{error}$ that are in top $k\%$. During training, the top $k\%$ frequency loss dynamically changes according to prediction at each iteration, since the frequencies that fall into top $k\%$ changes as the network learns. This ensures that the network is concentrated on the frequencies that it fails to restore the most. For example, a common data imbalance scenario may be that the bias term (zero frequency) dominates the others in a dataset where flat-like areas are more common. In such a case, a neural network with uniform contributions of DCT coefficients to loss function would dominantly be trained for recovering the zero-frequency term, and may easily discard not-so-populated frequencies. That would result in bad visual quality in some locations in an image. In the same scenario, when using our loss, once a network learns the bias-term well enough, zero frequency would not pop up in the top $k\%$ frequency errors anymore, therefore the network would concentrate on recovering frequencies that it fails the most: Possibly the minority frequency components that are less populated in the data.

Using the top $k\%$ frequency loss may result in a slight decrease in the performance in terms of distortion, since the goal is not to reduce the overall error. The goal is to reduce the error in a more balanced way in frequency domain in order to provide a better visual quality.

The entire neural network architecture is illustrated in Fig. 4. First, the input is transformed to a higher-dimensional space via consecutive convolutional and dense blocks. Then, multiple RSDB blocks are applied and the features are transferred back to image dimensions via consecutive dense and convolutional blocks. We also utilize a big skip connection between the input and the output of the neural network.

In order to verify the effectiveness of our DWT approach to the common DWT approach and no-DWT approach, we have conducted experiments with compatible blocks as shown in Fig. 5. In particular, we have experimented on three blocks: (a) our proposed RSDB block shown in Fig. 3, (b) a compatible block based on common DWT approach as shown in Fig. 5(a), and (c) a compatible block where no DWT is applied as shown in Fig. 5(b). Note that for (c) we have used space-to-depth (s2d) and depth-to-space(d2s) operators in order to be compatible with downsampling and upsampling procedures in (a) and (b) due to DWT and IDWT. The general neural network architecture for all experiments are exactly the same, only the bottleneck RSDB in Fig. 4 is replaced with corresponding bottlenecks for each experiment. The number of filters in compared bottlenecks are arranged slightly in order to keep the total number of parameters in the network similar.

All of the neural networks used in this work were trained on original SIDD-Medium dataset with provided training and validation splits in [3]. Except for the last convolutional layer, we have always used the same number of filters in all convolutional layers (64). We have used 8 RSDB blocks in total. We have trained all neural networks for 300 epochs in total, using RADAM [17] optimizer with initial learning rate of $0.0002$ and we have reduced the learning rate by 10 after 200 epochs. For the experiments that we have used top $k\%$ loss, we have gradually decreased $k$ by $1$ in each epoch, starting from $100$ to a minimum of $10$. All networks are trained to process RGGB Bayer data, where 1-channel RGGB image is converted to 4-channel by space-to-depth operator as neural network input. Neural network output is converted back to 1-channel RGGB image by depth-to-space operator. All experiments use horizontal flip, vertical flip and rotation as data augmentation during training. For all reported results in this paper, we have taken an ensemble of 8 results on all permutations of horizontal flip, vertical flip and rotation. While doing the augmentation and the ensemble we have followed the approach of [16] in order to preserve the Bayer order for all samples, particularly cropping or padding respective locations to preserve RGGB Bayer order for the neural network.

In this section, we compare PSNR results on validation sets of SIDD and SIDD+ datasets of our method and the baselines. For convenience, we call our method SDWT referring to separately applying operations on sub-bands (S) of DWT outputs.

We first analyze the training curves for compared algorithms in Fig. 6. It can be observed that our method learns faster at an already high PSNR in the beginning of the training compared to other methods. Moreover as the training continues, SDWT generalization keeps getting better, whereas for others validation PSNR seems to saturate at lower values compared to SDWT.

At the time of implementing the method, we did not yet have access to SIDD+ dataset, therefore validation was made based on SIDD dataset. Later, we had access to SIDD+ dataset too. Hence, we share our results in SIDD+ validation set as well. Similar to SIDD validation set, as it can be observed from Table 1, our SDWT outperforms the baselines in SIDD+ validation set too.

As expected, the network trained with top-$k\%$ loss ($SDWT_{top-k}$ ) has slightly lower PSNR then SDWT. Yet, as discussed earlier, the main advantage of this loss is perceptual quality. Next, we share some images from SIDD+ validation dataset to visually compare methods. After denoising, for better visual comparison, the images were processed by the ISP pipe provided by NTIRE-2020 Challenge organizers.

In Fig. 7, we show some enlarged crops from the SIDD+ validation dataset. SDWT’s advantage over other methods are clearly observed on all examples. The perceptual improvement of top-$k\%$ loss ($SDWT_{top-k}$) over L1 loss (SDWT) is best observed in high-frequency areas where top-$k\%$ loss helps recovering the edges better. The letters on images on rows 2 and 4 in Fig. 7 are clear examples for this.

For further demonstration of the visual quality improvement that $SDWT_{top-k}$ brings, we have conducted experiments on our own modules in Section 4.5, where the improvement is more evident.

In this section, we compare our method’s performance to the state-of-the-art on SIDD+ Test Set. We have evaluated the PSNR and MS-SSIM metrics of our method on a copied evaluation server of NTIRE-2020 Denoising Challenge - rawRGB Track, thanks to the original challenge organizers. We have evaluated two variants of our method with SDWT bottleneck: (a) 8 SDWT blocks and 64 filters ($SDWT_{8,64}$) (b) 12 SDWT blocks and 96 filters ($SDWT_{12,96}$). As the goal of the challenge was not perceptual quality, but distortion metrics, for these evaluations, we have only used L1 loss in order to aim for the best PSNR for each variant.

The relevant information in Table 2 are directly copied from [1]. Our heavier model $SDWT_{12,96}$ model ranks among the top-performing methods, whereas our model $SDWT_{8,64}$ gives a good speed and PSNR compromise. We would like to note here that the run-time of our method is calculated by considering the ensembles, for example $SDWT_{12,96}$ operates at 6.63 seconds per MPixel on 8 ensembles, which means per ensemble it runs with 0.83 seconds per MPixel.

In order to fully investigate the method in real use case, we have also conducted experiments in our own camera modules. In these experiments, the only difference from above experimental settings is that we have simulated the noisy data from SIDD ground truths according to our module’s noise model. Although we do not make use of the real noisy images in SIDD, the dataset is still irreplaceable due to the noise-free (ground truth) images it provides, which enables us to create realistically simulated noisy images. Note that in all experiments in this section, we have used our own ISP to generate images for visualization.

As the most important criterion for camera manufacturers, we have focused on visual quality improvement in these experiments. We first visually compare our method $SDWT_{top-k}$ to a baseline method NLM [8]. NLM was algorithmically adopted to process Bayer data and signal-dependent noise model, and tuned for best performance in our cameras. The comparison is made in test images captured with our camera at lab and in real world. It can be observed from Fig. 8 that $SDWT_{top-k}$ clearly outperforms NLM in representative samples from many cases: flat areas, dense lines and various textures.

We also compare $SDWT$ to $SDWT_{top-k}$ to highlight the differences visual quality that our proposed loss function brings. As it can be observed from Fig. 9, $SDWT_{top-k}$ in Fig. 9(b) provides a good restoration of the Siemens star lines whereas $SDWT$ in Fig. 9(a) over-smoothens the image, especially when lines gets denser. This shows that $SDWT_{top-k}$ indeed preserves more details which is an important factor for perceptual quality.