ENHANCEMENT OF A CNN-BASED DENOISER BASED ON SPATIAL AND SPECTRAL ANALYSIS

The whole network architecture is shown in Fig. 1. The network takes a noisy observation of size $W\times H\times C$ as input, where $W$, $H$, and $C$ are the width, the height, and the channels, respectively, i.e. $C=1$ for grayscale images and $C=3$ for color images. After the discrete wavelet transform (DWT), it becomes a tensor $\bm{t}$ of size $W_{0}\times H_{0}\times 4C$, where $W_{0}\leq W$ and $H_{0}\leq H$. $W_{0}$ and $H_{0}$ are determined by the size of the input image and the length of the wavelet filter. Each sub-band tensor $\bm{t}_{i}$ is then concatenated with a uniform noise level map, whose elements are all equal to $\sigma_{n}$, and the size of the sub-band tensors becomes $W_{0}\times H_{0}\times(C+1)$. We send these sub-band tensors to the corresponding branches of the band normalization module for training.

Since WDnCNN is a fully convolutional network without pooling layers, and the output of WDnCNN is of the same size as the input, we set all convolutional filters in WDnCNN to the size of $3\times 3$ and the padding size to $1$. Different from DnCNN and FFDNet, we do not apply any batch normalization (BatchNorm). The reason for this will be discussed in Sec. 2.3. Considering the scales and the semantic intensity of different filter bands, we set the number of convolutional layers in BNM to $3$, $2$, $2$, and $1$ for the LL, LH, HL, and HH bands, respectively. Considering both effectiveness and efficiency of the network, we set the number of convolutional layers in the mapping module to 16 and 13 for grayscale and color images, respectively. In terms of the number of channels for feature maps, we set 72 for grayscale images and 108 for color images, because the color images contain more information and require more feature maps for description. Inspired by DnCNN, we adopt residual learning to train the network. The residual learning aims to learn the pixel-wise noise for each sub-band. It is applied because the high-frequency bands contain more information about noise, which provides a better capacity for the network to predict the residual.

By using handcrafted image priors for denoising, we aim to solve the problem as follow:

$$\mathbf{\hat{w}}=\mathop{\arg\min}_{\mathbf{w}}||\mathbf{u}-\mathbf{w}||^{2}+\lambda\Theta(\mathbf{w}),$$

(2)

where $\mathbf{w}$ and $\mathbf{u}$ can be the clean and noisy image pair in any domain. $\Theta(\mathbf{w})$ is the regularization term, which relates to the image prior, and $\lambda$ is the weight that controls the balance between noise removal and detailed texture protection [5]. It has been proven that the weight-control parameter $\lambda$ is highly related to the noise level $\sigma_{n}$. Therefore, it justifies that concatenating a noise level map in FFDNet to guide the pixel-wise noise reduction is reasonable. However, for denoising in the spectral domain, each of the frequency spectra should be handled unequally, by setting $\lambda$ to different values. Thus, the above problem can be changed to

$$\mathbf{\hat{w}}_{k}=\mathop{\arg\min}_{\mathbf{w}_{k}}||\mathbf{u}_{k}-\mathbf{w}_{k}||^{2}+\lambda_{k}\Theta(\mathbf{w}_{k}),$$

(3)

where $\mathbf{w}_{k}$ and $\mathbf{u}_{k}$ are the clean and noisy pair of the $k$-th frequency component, and $\lambda_{k}$ is the weight of the $k$-th frequency component. To discriminatively guide each band for regression, we propose an unequal-depth structure in the band normalization module to generate the different $\lambda_{k}$.

In the band normalization module, unequal-depth branches are applied to the different band coefficients. Therefore, using more convolutional layers for the LL band and less convolutional layers for the HH band can benefit scale unification. Semantically, we need more convolutional layers for the lower frequency components to learn their noise distribution. By introducing the band normalization module with an unequal-depth structure, we achieve to use i) a single denoising network for different noise levels, and ii) unequal weights in different sub-bands for better texture preservation.

In DnCNN [4], internal covariate shift (ICS) [10] was addressed as the main problem for not obtaining sufficient regression. Therefore, batch normalization is applied to reduce the impact of ICS. Other networks based on DnCNN (e.g. FFDNet and MWCNN) also inherit this structure. However, in [11], it was found that the relation between BatchNorm and ICS reduction is tenuous, and BatchNorm might even increase ICS in training. The main contribution made by BatchNorm is that it makes the loss gradients more reliable. Therefore, considering the efficiency of our denoising model, we do not add any batch normalization layers, but propose a training criterion, namely band discriminative training (BDT), to enhance the regression of WDnCNN.

To train our model, we collect 4,744 images from the Waterloo Exploration Database [12], 400 images from the Berkeley segmentation dataset [13], and 600 images from the validation set of ImageNet [14], to generate the whole training set for grayscale image denoising. As for color image denoising, we replace 400 images from the Berkeley segmentation dataset with another 400 images from the validation set of ImageNet. For each epoch, we randomly select 1,000 noise-free images from the training set and then randomly crop $128\times 2,000$ patches, with patch size of $50\times 50$, from them. To obtain pairs of clean and noisy patches, additive white Gaussian noise (AWGN) is added to the selected patches. This is because the noise in a local region of a real noisy image can be assumed as AWGN. Moreover, we densely sample the noise level $\sigma_{n}$ in a range of $[0,75]$ for each synthetic noisy patch. The mini-batch size is set to 128 for all epochs, and the rotation-flip operations are employed for data augmentation.

We apply the ADAM [15] algorithm to the optimization process for minimizing the following loss function

$$\mathcal{L}(\Theta)=\frac{1}{2KN}\sum_{k=1}^{K}\sum_{i=1}^{N}\mu_{k}||\mathcal{F}_{k}(\mathbf{u}_{ik},\bm{M}_{ik};\Theta)-\mathbf{w}_{ik}||^{2},$$

(4)

where $K$ is the number of filter bands, and $\mu_{k}$ represents the weight for the $k$-th band. For the wavelet transformation, we simply choose the discrete Meyer wavelet (dmey) filter, because of its rapid decay and compact support in the frequency domain. Moreover, it does not cause any aliasing problems or distortions in the reconstruction process.

To achieve a better regression for each sub-band, we initialize the convolutional filters with Kai_ming_normal [16] in Pytorch. Due to the unequal importance of the different bands, conventional training procedures cannot lead to sufficient regression. To tackle this problem, we first train the initialized model with $\mu_{k}=1.5$, $2.5$, $2.5$, and $5$ for the LL, LH, HL, and HH bands, respectively. Because the wavelet coefficients of the clean image are unnormalized, we use these inversely proportional weights for keeping the target coefficients in balance. The learning rate is kept at $10^{-4}$ until the model fully converges. Then, we modify $\mu_{k}$ as shown in Table 1, for each 50 epochs, and fine-tune our model with a learning rate decreasing from $10^{-4}$ to $10^{-7}$. Each fine-tuning process can be divided into a reversal part and a regression part. A comparison of BDT and Non-BDT is shown in Fig. 2. It is clear that BDT provides a more stable training process and a better regression performance.

We evaluate our denoising model on the Set12, CBSD68 [13], McMaster [17], and Kodak24 [18] datasets for synthetic noise reduction. The tested noisy images are obtained from the above datasets corrupted by AWGN with noise level $\sigma_{n}$ equal to 25, 50, and 75. We compared our proposed method with other state-of-the-art methods, such as BM3D [2] (or CBM3D for color images), WNNM [3], and FFDNet [5]. Table 2 and Table 3 tabulate the results, in terms of PSNR(dB), of the different methods on grayscale and color images, respectively. The state-of-the-art denoisers selected for our experiments can handle different noise levels with a single model.

It can be observed from the results that WDnCNN outperforms the benchmark BM3D by a large margin on all testing sets. Compared to WNNM, WDnCNN can obtain a PSNR gain of about 0.3dB for denoising grayscale images. Although WDnCNN only achieves comparable results with FFDNet when the noise level is low($\sigma_{n}\leq 25$), it shows great ability in distinguishing texture information from noise, contributing to a relatively larger PSNR improvement on C.man, Peppers, and House. Furthermore, WDnCNN can regress the noise distribution more discriminatively on different sub-bands, making it more effective when noise is stronger and able to achieve a PSNR gain of about 0.2dB, with $\sigma_{n}=75$, compared to FFDNet. Considering the fact that we used a long wavelet filter, the performance of WDnCNN can be further improved by using Sym8 for a larger receptive field.

To further evaluate WDnCNN, we tested it on some real-world noisy images. Since the ground-truth clean images are unavailable, visual comparison is adopted for assessment. One example is shown in Fig. 3. It can be seen that WDnCNN produces a better, detailed texture preservation, as the grain and wrinkles on the flower stalk are more vivid. In summary, WDnCNN establishes state-of-the-art visual quality for the restored images.