Deep learning-based Edge-aware pre and post-processing methods for JPEG compressed images

Various post-processing modules have been proposed in the literature. Modeling of the image-prior, based on maximum a posteriori (MAP) criteria is used in [3], but it is found to be quite more expensive than handcrafted models. The post-processing method in [4] based on the adaptive DCT method is quite good, but failed to generate sharp images. The notion of Wiener filtering for denoising used in [5] was novel, however, the images generated contained highly visible noise at the edge regions. Similarly, the modeling of image prior and quantization is found to be expensive in [6]. Ren et al. [7] combined the local sparsity image prior model and non-local similarity model for artifacts reduction using low-rank minimization and patch clustering proposed in [3] and [6] respectively. Dictionary learning and variational regularization used in [8] for restoring images outperformed the other decompression method but consumed substantial computational time. Weighted nuclear norm minimization (WNNM) concept used in [9] being complex, outperformed Dabov et al. [5] approach while preserving the local structure and reducing visual artifacts taking more computation time. The deblocking method, given in [10] using a non-convex low rank constrained model was a good optimization of the problem. The denoising method proposed by Zhang et al. [11] helped in reducing artifacts but with the lack of retaining edge information in reconstructed images. The benchmark deblocking method ARCNN, proposed by Dong et al. [12] outperformed the above-discussed methods used for image compression. All these deblocking and post-processing models aimed to improve the reconstruction quality while ignoring the joint optimization of the traditional codec with encoder-decoder structure.

Recent work [13, 14] has demonstrated promising results with deep learning methodologies for compression. Specifically, the work in [13] proposed an “end-to-end” compression framework wherein pre & post-processing DNN is introduced around a standard compression codec. However, the training procedure did not take care of high-frequency components (i.e. edges) due to optimization with mean square error (MSE) resulting in smooth reconstructed images. Moreover, the author did not present the rate-distortion analysis for the JPEG codec, which is needed to compare the compression performance based on bit-rates. Like [13], the approach in [14] uses jointly trained CNNs for the pre-processing and post-processing stages. These stages incorporate resolution reduction only for low JPEG quality factor (QF) and maintain resolution for higher QF. On the other hand, the work in [15] is based on investigating different downsampling networks and super-resolving or upsampling networks along with the different training strategies. The sequential training of downsampling CNN with upsampling CNN only has been considered, without the consideration of any codec in-between.

In 2016, Toderici et al. [16] introduced an RNN based full resolution compression exploiting the principle of progressive encoding-decoding of images. In 2017, Balle et al. [17] proposed to use generalized divisive normalization based local gain control method with the use of additive uniform noise to make non-differentiable quantization function a continuous function. Theis et al. [18] developed an algorithm to use a single autoencoder for providing variable bit-rate image compression by just tweaking with the size of latent space representation before quantization. Basically, the work in [16, 19, 20] is based on the calculation of the probability of binary representation for context-adaptive coding. Basically, Johnston et al. [20] extracted binary information for maximizing the entropy. In 2018, Li et al. [21] came up with a compression model based on content prediction with the help of importance map generation through CNN. Later on, Balle et al. extended the work in [17] by exploiting spatial dependencies for variable rate image compression in [22]. Similarly, Lee et al. [23] extended the work in [22] by proposing a context-adaptive entropy model by incorporating bit-consuming and bit-free contexts. On the other hand, the approach reported by Minnen et al. [24] was based on weighted auto-regressive and hierarchical priors for contextual image compression. Mentzer et al. [25] proposed to focus on rate-distortion trade-offs by learning a conditional probability model of latent space distribution in a 3D-CNN autoencoder. Later on in 2019's, Choi et al. [26] proposed a variable bit-rate image compression in the hierarchy of having one compression module for providing multiple bit-rates. In early 2020's, the model proposed by Yang et al. [27] was based on multi-layer feature modulation which reduced the problem of lower performance at low bit-rates. Chen et al. [28] proposed to learn a single CNN model at high bit-rates and then scaling the bottleneck by some factor for low bit-rates.

Motivated by the discussions on downsampling/upsampling based CNN methods, we utilize an “end-to-end” deep compression framework compatible with the traditional codec by utilizing pre & post-processing DNNs. Different from prior works, however, we:

• •

  introduce an edge-aware loss function during compression, which significantly improves compression performance. The edge-aware loss function is motivated by the need to minimize edge distortion in compression, which the MSE loss function in [13] does not emphasize.

• •

  for the low bit-rate regime, utilize a super-resolution network in the post-processing stage.

• •

  prefer to use direct image-to-image learning over patch-to-patch learning criteria which is generally practiced in previous compression algorithms. The later scheme degrades the training performance, prediction efficiency, and lead to distortions in the output image.

In the proposed approach, we utilize the enhanced deep super-resolution (EDSR) network [29], which is a multi-scale network designed for specific super-resolution scaling that reduces model size and training time using residual learning compared with alternative approaches. The residual learning used in the super-resolution network helps in preserving the detail components. Because the quantization in the codec represents a non-differentiable operation, we use progressive training that co-optimizes the pre & post-processing modules in an “end-to-end” manner by appropriately approximating the codec. The super-resolution is chosen to improve the quality of the JPEG-compressed output by reducing blocking artifacts at the time of post-processing. In-fact, super-resolution which is a process of converting a low-resolution image into its high-resolution counterparts is exactly equivalent to image decompression which aims to improve the quality of the decoded image by adding more high-frequency details. Hence, the image super-resolution mechanism can be called a special case of an image compression-decompression algorithm. Different from [13, 14], the proposed network uses an edge-aware loss function to optimize the pre-processing module. In addition to that, it uses EDSR module instead of VDSR [30] module (used in [13]) & 5 layer encoder-decoder network (used in [14]) as post-processing module. The proposed algorithm is quite different from the work reported in [14, 15] in terms of the pre-processing module and post-processing module exploited. Moreover, the algorithm presented in [14, 15] exploited switching between one network to another when investigating compression performance from low bit-rates to high bit-rates, which does not seem optimal. Also, the author did not compare the work reported in [15] with any DNN based image compression algorithms.

The remaining sections of the paper are organized as follows. Section II presents the proposed approach covering details of the post-processing DNN & the edge-aware loss function [31, 32] that we introduce. Section III discusses for the datasets used and other experimental setup details. In Section IV, comparison between different experiment strategies and modules has been shown. Section V throws light on the comparison of the final proposed algorithm with the state-of-the-art algorithms. Section VI finally summarizes the effectiveness of the proposed algorithm reported in this paper.

The proposed compression-decompression model is trained on 400 images ($256\times 256$) of the ImageNet dataset [34]. In order to compare the results with [14], we have assessed the algorithm on various datasets namely Live 1 [35], Set 14 [36], Set 5 [37] & Set 7 (same as built up in [14]). These datasets are having different attributes which also help in the better generalization of the proposed scheme. Live 1 dataset has twenty-nine diversified generic PNG images with sizes ranging from $438\times 634$ to $720\times 480$. These images contain high resolution and high-quality, sharp images. Set 14 contains fourteen low-frequency PNG images with size ranging from $250\times 361$ to $768\times 512$. Set 5 contains five edge centered luminance images. These are images with high-frequency domination where size varies from $280\times 280$ to $512\times 512$. Apart from these three standard test sets, a test of 7 images named Set 7 is prepared (as in [14]) to compare the results with Zhao's approach [14]. This special set contains both smooth and sharp images with size $256\times 256$, $512\times 512$ & $512\times 768$, and reported as challenging for image compression application. To compare with the recent techniques, the proposed algorithm is also assessed on Kodak benchmark dataset [38] and CLIC challenging dataset [39]. Kodak dataset contains uncompressed true-color PNG images, with a size of $768\times 512$ or $512\times 768$. CLIC 2019 consists of high-resolution and high-quality images with size ranging from $512\times 768$ to $2048\times 1646$. Moreover, to compare the deblockiness effect of the proposed algorithm, we have also tested Classic 5 [40] & General 100 [41] datasets. Classic 5 contains five BMP gray images of $512\times 512$ size. General 100 contains a hundred images in BMP format with zero compression, which makes it highly suitable for the compression task. It contains good quality images with clear sharp edges but fewer smooth regions. The image dimensions range from $131\times 112$ to $710\times 704$.

Both the CNN modules are trained for five iterations in each epoch. Adam optimizer is used to optimize the weights of the network that has shown very good performance as reported in [42]. Initially, $\alpha=0.01$ was selected as the learning rate, but the loss observed was very high. Then we decreased the learning rate to $\alpha=0.001$, due to which loss also decreased substantially. On further decreasing the learning rate, we did not get much improvement. Hence, we finalized the learning rate at $\alpha=0.001$. The momentum parameters are selected as $\beta_{1}=0.9$, $\beta_{2}=0.999$ & $\epsilon=1\times 10^{-7}$, which are reported to yield good results as in [13, 29, 43] to find the global minima. The batch size depends upon the computational platform used for the implementation of the CNN based algorithm. It should fit into the memory specification of CPU or GPU based architecture. Generally, training on small batch sizes speeds up the learning process at the cost of low accuracy on account of noise in the training process. On the other hand, with large batch size, the learning process converges slowly with better prediction or estimation. Accordingly, we have set a mini-batch size of 10, suitable for the used platform. We have finalized training epochs to be 50, as, after that, the performance was not improving. All the experiments have been conducted using Python with Tensorflow & Keras deep learning framework on TITAN X (Pascal)/PCIe/SSE2 16 GB CPU RAM, 12 GB GPU enabled system.

The original image $\mathbf{f}(x,y)$ has been compared with $\hat{\mathbf{f}}(x,y)$ for the assessment of the proposed compression-decompression algorithm. We have used both objective and subjective performance metrics i.e., PSNR and SSIM for a comprehensive assessment. Along with the consideration of structural information, a metric should be chosen which incorporates image details at various resolutions. Therefore, MS-SSIM, a measure based on structural information at different scales is chosen for assessment. In the same direction, we have also calculated PSNRB [47] which incorporates the effect of blocking artifacts (by incorporating the blocking effect factor) while calculating the PSNR. To check for the edge-preserving property of edge-aware loss function we have calculated intersection over union (IoU) to check the overlap regions between true edges and retained edges. Bjontegaard metric calculation [48] has also been conducted to compare the coding gain achieved with the proposed algorithm with the other methods. The BD metric facilitate the comparison of area under rate-distortion curves in terms of the average PSNR improvement or the average percent bit-rate saving. A negative (positive) value of BD-Rate indicates a decrease (increase) of bit-rate for the same PSNR. A positive (negative) value of BD-PSNR indicates an increase (decrease) of PSNR for the same bit-rate.

Many edge detection algorithms are available in the literature. CED is one of the good classical edge detection algorithms which uses the idea of change in intensity on each pixel in an image, i.e., the intensity is very high on the edges. CED identifies edges in 4 steps: noise removal, gradient calculation, non-maximal suppression, and hysteresis thresholding. Since the CED only focuses on local changes and it has no semantic understanding (understanding the content of the image), it has limited accuracy. Semantic understanding is crucial for edge detection that is why learning-based detectors generate better results than CED. Hence, we have separately tested the classical Canny Edge Detection [49] and learning-based edge detector. For learning-based edge detection, we have chosen the Holistically-Nested Edge Detection technique (HED), which has been set as a benchmark technique for edge detection after the development in the CNN architectures. HED is a learning-based end-to-end edge detection system that uses a trimmed VGG-like CNN for an image to image prediction task. It is based on multi-scale and multi-level feature learning and claims to learn rich hierarchical features which are crucial. However, Canny edges are not directly connected and produce some spatial shifts and inconsistencies. It makes use of the side outputs of intermediate layers to learn edges synchronizing with the contextual information leaving undesired and false edges. Since the feature maps generated at each layer are of different sizes, it effectively looks at images at different scales. Generally, edge detectors are supposed to assume 90% of the ground truth as non-edge, but HED uses a class balancing weight to make a balance between edge & non-edge region, to find the true & strong edges much effectively as shown in Fig. 3. To implement it, we have used the code provided at ⁴⁴4https://github.com/senliuy/Keras_HED_with_model for exploiting edge map predictions obtained through training the modified VGGNet with BSDS500 [50]. HED seems to produce better compression performance than CED as shown in Table I.

In [13], the VDSR module has been used as the post-processing module. However, we have used the EDSR module to remove the blocking effects generated after JPEG decoding. To show the efficacy of the EDSR module, we have implemented the proposed algorithm considering the VDSR post-processing module. The compression with the EDSR module achieves significant improvement over compression with the VDSR module as shown in Table II.

To check the individual contribution of each sub-module i.e., pre & post-processing module, we have trained the network considering $P_{r}N$ & $P_{o}N$ individually. In fact, these individual implementations are the special cases of proposed $P_{E}~{}(FR)$. From Table III, it is inferred that traditional codec compression followed by the post-processing network $P_{o}N$ produces better results than pre-processing module $P_{r}N$. Moreover, adding both these modules with the codec compression added further gain in compression performance.

The training of $P_{r}N$ on edge-aware loss forced the network to retain structural and sharp details while compressing the images. Also, the tuning parameter $\alpha=0.75$ worked best for training the network to preserve edges in the reconstructed images. The first term in Eq. (2) ensures high PSNR by minimizing MSE while preserving edges. Preserving edges via the second term in Eq. (2) ensures the preservation of structural information and hence maximize SSIM & MS-SSIM. High mIoU value in Table IV indicate to retain edges with edge loss (in $P_{E}$) as compared to MSE loss function (in $P_{M}$), ultimately ensuring the edge-preserving property. Also, it indicates that an FR network can recover images with quite a higher mIoU than a CR representation network. In addition to this, we have trained both the modules with MSE and the results are given in Table V. The results show that training of $P_{r}N$ with edge-aware loss function leads to better results than training with MSE. Apart from that, Li et al. [15] use a resolution decrement parameter (layer with stride 2) at the first CNN layer. But, we believe it should be placed at the last CNN layer (as used in the proposed algorithm) because gradual or delayed downsampling help in learning more features as compared to sudden downsampling just after the first layer. This sudden downsampling may lose learning of some features which are crucially required. To support this statement, a small experiment is conducted first by taking stride 2 in the first layer, then in last and the results are shown in Table VI. With multiple sets of experiments, we finalize the training of $P_{r}N$ & $P_{o}N$ module on edge-aware & MSE loss functions, respectively. It was found that training with edge-aware loss functions performed well with comparatively more bit saving.

The results in Fig. 4 stressed upon the point of more appropriate learning of features showing the superiority of the FR over CR image representation in terms of compression gain and performance measures. From the highlighted region in the image given in Fig. 5, it is quite evident that at low bit-rates, the proposed $P_{E}~{}(CR)$ & $P_{E}~{}(FR)$ reduces the blocking artifacts more effectively compared to JPEG, due to super-resolution network introduced to post-process the image. However, it is noted that better performance at low bit-rates using CR network representation is obtained at the cost of performance saturation at high bit-rates which can be seen from Fig. 6.