Joint Image Compression and Denoising via Latent-Space Scalability

State-of-the-art classical image denoising methods are based on Non-local Self Similarity (NSS). In these methods, repetitive local patterns in a noisy image are used to capture signal and noise characteristics, and perform denoising. In BM3D [5], similar patches are first found by block matching. Then, they are stacked to form a 3D block. Finally, transform-domain collaborative filtering is applied to obtain the clean patch. [6] used adaptive filtering to improve BM3D. WNNM [7] performs denoisng by applying low rank matrix approximation to the stacked noisy patches. In [8], a patch group based NSS prior learning scheme to learn explicit NSS models from natural images is proposed. The denoising method in [9] used NSS priors in both the degraded images and the external clean images to perform denoising. CBM3D [10] and MCWNNM [11] are the extensions of BM3D and WNNM, respectively, created to handle color images.

More recently, (deep) learning-based denoising methods have gained popularity and surpassed the performance of classical methods.  [12] used a multi-layer perceptron (MLP) to achieve denoising results comparable to the state-of-the-art classic method. Among the learning-based denoisers, DnCNN [13] was the first Convolutional Neural Network (CNN) to perform blind Gaussian denoising. FFDNet [14] improved upon DnCNN by proposing a fast and flexible denoising CNN that could handle different noise levels with a single model. In [15], noise estimation subnetwork is added prior to the CNN-based denoiser to get an accurate estimate of the noise level in the real-world noisy photographs. A Generative Adversarial Networks (GAN)-based denoising method is proposed in [16]. The mentioned works are supervised methods where clean reference image is needed for training. In [17, 18], self-supervised denoising methods are proposed.

In recent years, there has been an increasing interest in the development of learning-based image codecs. Some of the early works [19, 20, 21] were based on Recurrent Neural Networks (RNNs), whose purpose was to model spatial dependence of pixels in an image. More recently, the focus has shifted to Convolutional Neural Network (CNN)-based autoencoders. [22] introduced Generalized Divisive Normalization (GDN) as a key component of the nonlinear transform in the encoder. The image codec based on GDN was improved by introducing a hyperprior to capture spatial dependencies and take advantage of statistical redundancy in the entropy model [23]. To further improve the coding gains, discretized Gaussian mixture likelihoods are used in [24] to parameterize the distributions of latent codes. Most recently, this approach has been extended using advanced latent-space context modelling [25] to achieve even better performance.

Most state-of-the-art learning-based image codding approaches [23, 24, 25] train different models for different bitrates, by changing the Lagrange multiplier that trades-off rate and distortion. Such approach is meant to explore the potential of learning-based compression, rather than be used in practice as is. There has also been a considerable amount of work on variable-rate learning-based compression [19, 26, 27, 28, 29], where a single model is able to produce multiple rate-distortion points. However, in terms of rate-distortion performance, “fixed-rate” approaches such as [24, 25] currently seem to have an advantage over variable-rate ones.

The mentioned learning-based codec are single-task models, where the task is the reconstruction of the input image, just like with conventional codecs. However, the real power of learning-based codecs is their ability to be trained for multiple tasks, for example, image processing or computer vision tasks, besides the usual input reconstruction. In fact, the goal of JPEG AI standardization is to develop such a coding framework that could support multiple tasks from a common compressed representation [30].

[31] proposed a scalable multi-task model with multiple segments in the latent space to handle computer vision tasks in addition to input reconstruction. The concept was based on latent-space scalability [32], where the latent space is partitioned in a scalable manner, from tasks that require less information to tasks that require more information. Our JICD framework is also based on latent-space scalability [32]. However, unlike these earlier works, the latent space is organized such that it supports image denoising from the base layer and noisy input reconstruction from the full latent space. In other words, the tasks are different compared to these earlier works.

Recently, [33] and [34] developed joint image compression and denoising pipelines built upon learning-based image codecs, where the pipeline is trained to take the input noisy image, compress it, and decode a denoised image. However, with these approaches, it is not possible to reconstruct the original noisy image, hence they are not multi-task models. Our proposed JICD performs the denoising task in its base layer, but it keeps the noise information in the enhancement layer, thereby also enabling noisy input reconstruction if needed.

The encoder employs an analysis transform to obtain a high fidelity latent-space representation for the input image. In addition, the encoder has blocks to efficiently encode the obtained latent-space tensor. The encoder’s analysis transform is borrowed from [24] due to its high compression efficiency. In addition to the analysis transform, we also adopted the entropy parameter (EP) module, the context model (CTX) for arithmetic encoder/decoder (AE/AD), synthesis transform and hyper analysis/synthesis without attention layers from [24].

The analysis transform converts the input image $\mathbf{X}$ into $\mathbfcal{Y}\in\mathbb{R}^{N\times M\times C}$, with $C=192$ as in [35, 24]. Unlike [35, 24], the latent representation $\mathbfcal{Y}$ is split into two separate sub-latents $\mathbfcal{Y}=\mathbfcal{Y}_{1}\cup\mathbfcal{Y}_{2}$, $\mathbfcal{Y}_{1}\cap\mathbfcal{Y}_{2}=\emptyset$, where $\mathbfcal{Y}_{1}$ is the base layer containing $i$ channels, $\mathbfcal{Y}_{1}=\{\mathbf{Y}_{1},\mathbf{Y}_{2},...\mathbf{Y}_{i}\}$, and $\mathbfcal{Y}_{2}$ is the enhancement layer containing $C-i$ channels, $\mathbfcal{Y}_{2}=\{\mathbf{Y}_{i+1},\mathbf{Y}_{i+2},...,\mathbf{Y}_{C}\}$. This allows the latent representation to be used efficiently for multiple purposes, namely denoising (from $\mathbfcal{Y}_{1}$) and noisy input reconstruction (from $\mathbfcal{Y}_{1}\cup\mathbfcal{Y}_{2}$). Since denoising requires only $\mathbfcal{Y}_{1}$, it can be accomplished at a lower bitrate compared to decoding the full latent space. The sub-latents are then quantized to produce $\widehat{\mathbfcal{Y}}_{1}$ and $\widehat{\mathbfcal{Y}}_{2}$, respectively, and then coded using their respective context models to produce two independently-decodable bitstreams, as discussed in [31, 32]. The side bitstream shown in Fig. 1 is considered to be a part of the base layer and its rate is included in bitrate calculations for the base layer bitstream in the experiments.

Two task-specific decoders are constructed: one for denoised image decoding and one for noisy input image reconstruction. The hyperpriors used in both decoders are reconstructed from the side bitstream which, as mentioned above, is considered to be a part of the base layer. Quantized base representation $\widehat{\mathbfcal{Y}}_{1}$ is reconstructed in the base decoder by decoding the base bitstream, and used to produce the denoised image $\widehat{\mathbf{X}}$. Unlike [32, 31], where the base layer was dedicated to object detection/segmentation, our decoder does not require latent space transformation from $\widehat{\mathbfcal{Y}}_{1}$ into another latent space; the synthesis transform (Fig. 1) produces the denoised image $\widehat{\mathbf{X}}$ directly from $\widehat{\mathbfcal{Y}}_{1}$. Quantized enhancement representation $\widehat{\mathbfcal{Y}}_{2}$ is decoded only when noisy input reconstruction is needed. The reconstructed noisy input image $\widehat{\mathbf{X}}_{n}$ is produced by the second decoder using $\widehat{\mathbfcal{Y}}=\widehat{\mathbfcal{Y}}_{1}\cup\widehat{\mathbfcal{Y}}_{2}$.

Although not pursued in this work, it is worth mentioning that the proposed JICD framework can be extended to perform various computer vision tasks as well, such as image classification or object detection. These tasks typically require clean images, so one can think of the processing pipeline described by the following Markov chain: $\mathbf{X}_{n}\to\widehat{\mathbfcal{Y}}_{1}\to\widehat{\mathbf{X}}\to T$, where $T$ is the output of a computer vision task, for example a class label or object bounding boxes. Applying the DPI to this Markov chain we have

which implies that a subset of information from $\widehat{\mathbfcal{Y}}_{1}$ is sufficient to produce $T$. Hence, if such tasks are required, the encoder’s latent space can be further partitioned by splitting $\widehat{\mathbfcal{Y}}_{1}$, in a manner similar to [32, 31], to support such tasks at an even lower bitrate than our base layer.

The proposed multi-task model is trained from scratch using the randomly cropped $256\times 256$ patches from the CLIC dataset [37]. The noisy images are obtained using additive white Gaussian noise (AWGN) with three noise levels $\sigma=\{15,25,50\}$, clipping the resulting values to [0, 255] and quantizing the clipped values to mimic how noisy images are stored in practice. The batch size is set to 16. Training is ran for 300 epochs using the Adam optimizer with initial learning rate of $1\times 10^{-4}$. The learning rate is reduced by factor of $0.5$ when the training loss plateaus. We trained 6 different models by changing the value of $\lambda$ in (9). The list of different values for $\lambda$ is shown in Table I. For all the models we used $w=0.05$ in (11). We trained the model for the first rate point (lowest $\lambda$) from scratch. However, for the remaining rate points we fine-tune the model starting from the previous rate point’s weights

We trained models under two different settings. In the first setting, a given model is trained for each noise level. For this case, the number of enhancement channels $C-i$ is chosen according to the strength of the noise. For stronger noise, we allocate more channels to the enhancement layer, so that it can capture enough information to reconstruct the noise. The number of enhancement channels is reduced as the noise gets weaker. Specifically, the number of enhancement channels is empirically set to 32, 12, and 2 for $\sigma=50$, $\sigma=25$, and $\sigma=15$, respectively. The second training setting is to train a single model with different noise levels $\sigma\in\{50,25,15\}$ simultaneously, and use the final trained model to perform denoising for all noise levels. This is beneficial when the noise level information is not given. In this model, we used 180 base channels and 12 enhancement channels. $\sigma$ at each training iteration is uniformly chosen from $\{50,25,15\}$.

To evaluate the performance of the proposed JICD framework, four color image datasets are used: 1) CBSD68 [38], 2) Kodak24 [39], 3) McMaster [40] and 4) JPEG AI testset [30], which is used in the JPEG AI exploration experiments. The mentioned datasets contain 68, 24, 18, and 16 images, respectively. The resolution of the images in the Kodak24 and McMaster dataset is fixed to $500\times 500$. CBS68 dataset contains the lowest-resolution images among the four datasets, with the height and width of images ranging between 321 and 481. The images in the JPEG AI testset are high-resolution images with the height varying between 872 and 2456 pixels and width varying between 1336 and 3680 pixels. The results are reported for two sets of noisy images. In the first set, we added synthesized AWGN to the testing images with three noise levels: $\sigma=\{15,25,50\}$ and tested the results with the quantized noisy images. In the second set, we used the synthesized noise obtained from the noise simulator in [41], which was also used to generate the final test images for the denoising tasks in the ongoing JPEG AI standardization. This type of noise was not used during the training of the proposed JICD framework. Hence, the goal of testing with this second set of images is to evaluate how well the proposed JICD generalizes to the noise that is not seen during the training.

The denoising performance of the proposed JICD framework is compared against well-established baselines: CBM3D [10] and FFDNet [14]. CBM3D is a NSS-based denoising method, and FFDNet belongs to the learning-based denoising category. FFDNet was trained using AWGN with different noise levels during the training. At inference time, FFDNet needs the variance of the noise as input. FFDNet-clip [14] is a version of FFDNet that is trained with quantized noisy images. Since our focus is on practical settings with quantized noisy images, we used FFDNet-clip as a baseline in the experiments. We also tested the DRUNet denoiser [42], which is one of the latest state-of-the-art denoisers. DRUNet assumes that the noise is not quantized, and when tested with quantized noise, it performs worse than FFDNet-clip. As a result, we did not include it in the experiments.

Two baselines are established by applying CBM3D and FFDNet-clip directly on noisy images, without compression. However, to assess the interaction of compression and denoising, we establish one more baseline. In this third baseline, the noisy image is first compressed using the end-to-end image compression model from [24] (the “Cheng model”) with an implementation from CompressAI [43], and then decoded. Then FFDNet-clip is used to denoise the decoded noisy image. We call this cascade denoising approach as Cheng+FFDNet-clip. It is worth mentioning that Cheng+FFDNet-clip, similar to the proposed JICD framework, is able to obtain both the reconstructed noisy images and denoised images, hence it could be considered a multi-task approach.

We evaluate the baselines and the proposed JICD method using the quantized noisy images obtained using AWGN with three noise levels, $\sigma\in\{15,25,50\}$. The test results with the strongest noise ($\sigma=50$) across the four datasets (CBSD68, Kodak24, McMaster, and JPEG AI) are shown Fig. 3 in terms of rate vs. Peak Signal-to-Noise Ratio (PSNR) and in Fig. 4 in terms of rate vs. Structural Similarity Index Measure (SSIM). The horizontal lines in the figure correspond to applying CBM3D and FFDNet-clip to the raw (uncompressed) noisy images. The blue curve shows the results for Cheng+FFDNet-clip. The six points on this curve correspond to the six Cheng models from CompressAI [43]. For JICD, two curves are shown. The orange curve shows the results obtained from the models trained for $\sigma=50$ with 160 base feature channels and 32 enhancement channels. The yellow curve corresponds to the results obtained using the model that was trained with variable $\sigma$ values and has 180 base and 12 enhancement channels. The six points on the orange and yellow curves correspond to the six JICD models we trained with $\lambda$ values shown in Table I.

As seen in Fig. 3, for $\sigma=50$, the quality of the images denoised by CBM3D is considerably lower compared to those obtained using FFDNet-clip. It was shown in [14] that CBM3D and FFDNet-clip achieve comparable performance for non-quantized noisy images. Our results show that CBM3D’s performance is degraded when the noise deviates (due to clipping and quantization) from the assumed model, at least at high noise levels.

The comparison of the results obtained by JICD and Cheng+FFDNet-clip reveal that JICD is able to reduce the bitrate substantially while achieving the same denoising performance as Cheng+FFDNet-clip. This is due to the fact that the Cheng model allocates the entire latent representation to noisy input reconstruction, whereas the proposed method uses a subset of the latent features to perform denoising. The results of JICD trained with variable $\sigma$ are also shown in the curves. Since the number of base channels is larger in this model compared to the model trained for $\sigma=50$, its denoising performance is improved.

To summarize the differences between the performance-rate curves, we compute Bjøntegaard Delta-rate (BD-rate) [44]. The BD-rate of the proposed JICD compared to Cheng+FFDNet-clip on the four datasets is given in the first two rows of Table II for PSNR, and Table III for SSIM. It can be seen that the proposed method achieves up to 80.2% BD-rate savings compared to Cheng+FFDNet-clip. Both JICD and Cheng+FFDNet-clip denoising methods outperform CBMD3D for all the tested rate points at $\sigma=50$. Using the proposed JICD method, we are able to denoise images at a quality close to what FFDNet-clip achieves on raw images, and at the same time compress the input.

We repeat the denoising experiment for $\sigma=25$, and the results are shown in Fig. 5 for PSNR and Fig. 6 for SSIM. As seen in the figures, the gap between the CBM3D and FFDNet-clip performance is now reduced, and the compression-based methods now outperform CBM3D only at the higher rates. The gap between the curves corresponding to JICD and Cheng+FFDNet-clip is also reduced. However, JICD still achieves a considerable BD-rate saving compared to Cheng+FFDNet-clip, as shown in the third row of Table II and Table III. JICD trained with variable $\sigma$ has slightly better PSNR performance compared to the noise-specific model on three datasets, and a slightly worse performance (by 0.3%) on the low-resolution CBSD68 dataset.

At the lowest noise level ($\sigma=15$), the gap between CBM3D and FFDNet-clip shrinks further. It can be seen in the denoised rate-PSNR curves in Fig. 7 and rate-SSIM curves in Fig. 8 that when the noise is weak, applying denoising to the raw images achieves high PSNR, and the compression-based methods cannot outperform either CBM3D, or FFDNet-clip at the tested rates. The gap between JICD and Cheng+FFDNet-clip curves is also reduced compared to the higher noise levels. This can also be be seen from the BD-rates in the fourth row of Table II and Table III. JICD trained for $\sigma=15$ outperforms Cheng+FFDNet-clip on three datasets, but it suffers a 1% (4.5% for SSIM) loss on the low-resolution CBSD68.

As seen above, the performance of the proposed JICD framework is lower on the low-resolution CBSD68 dataset than on other datasets. The reason is the following. The processing pipeline USED in JICD expects the input dimensions to be multiples of 64. For images whose dimensions do not satisfy this requirement, the input is padded up to the nearest multiple of 64. At low resolutions, the padded area may be somewhat large in relation to the original image, which causes noticeable performance degradation. At high resolutions, the padded area is insignificant compared to the original image, and the impact on JICD’s performance is correspondingly smaller. It is worth mentioning that for $\sigma=15$, the JICD trained with variable $\sigma$ has a weaker denoising performance compared to the model trained specifically for $\sigma=15$. This is because the number of base channels in the variable-$\sigma$ model (180) is smaller than the number of base channels in the noise-specific model (190). At low noise levels, fewer channels are needed to hold noise information, which means the number of base channels could be higher. Hence, the structure chosen for the noise-specific model is better suited for this case. However, we show in the next subsection that the variable-$\sigma$ model is more useful when the noise parameters are not known.