Learning Weighting Map for Bit-Depth Expansion within a Rational Range

Let $\mathbf{I}^{L}$, $\mathbf{I}^{H}$ be the LBD and restored HBD images, respectively. Our method aims to find the missing information from lost bits and keep the given high-order bits. Let $\mathbf{I}^{Res}$ be the residual image of the replenished low-order bits, then there is

The residual image is gained based on the learned weighting map $\mathbf{I}^{WM}$. Let $\mathbf{N}^{b}$ be the theoretical upper bound of the missing bits. Then, $\mathbf{N}^{b}$ is a 3-channel tensor denoting the RGB color channels separately. The value of $\mathbf{N}^{b}$ is the decimal representation of the missing low-order bits. For example, if the number of missing bits $b=4$, each value of $\mathbf{N}^{b}$ is $2^{4}=16$. Thus, the residual image $\mathbf{I}^{Res}$ is

where $*$ denotes the point-wise multiplication, and $I^{WM}$ is the weighting map with each value in the range of $[0,1]$.

Further, we design a bit restoration network (BRNet), jointly considering $\mathbf{I}^{L}$ and the upper bound $\mathbf{N}^{b}$ to infer the adaptive weighting map $I^{WM}$. Let $f^{BRNet}(\cdot)$ represent the network, we can formulate the learning process as

where $[\cdot]$ is the concatenation operation.

Fig. 2 illustrates the entire framework of the proposed BRNet. BRNet is mainly composed of encoder and decoder that follows an UNet-style design [30]. The encoder contains an input convolutional layer and three optimization blocks. The input layer converts the input tensor to a series of feature maps for optimization. The optimization blocks (OptBlock) explore the feature distribution in different scale and find a feasible solution in the feature space. The optimized feature maps are down-scaled successively by MaxPooling and explored by the next stage. The smaller size of feature equivalently increases the receptive field of network and learns different scales of feature distributions. The decoder part integrates the multi-scale solutions and produces the desired weighting map.

The design of optimization block (OptBlock) follows the mathematical analysis of BDE in the optimization perspective. Let $\mathbf{F}^{L}$ be the LBD feature, the task of OptBlock is to find a restored HBD feature $\mathbf{F}^{H}$ satisfying

where $\mathcal{D}_{b}(\cdot)$ is the operation that removes the information of last $b$ bits, $\Phi(\cdot)$ is the prior term, $\lambda$ is a weighting factor. It is worth noting that Eq. 4 is an unconstrained optimization issue, and the solution space of $\mathbf{F}^{H}$ has no extra boundary. By converting the feature into weighting map, we can find a suitable mapping within the rational range.

An ideal prior $\Phi(\cdot)$ should be a function related to the number of bits in $\mathbf{F}^{H}$, which can be regarded as a statistical-relative indicator. However, it is difficult for derivation. Thus, we employ proximal gradient descent algorithm to solve Eq. 4 in an iterative formulation by

where

and

$\mathbf{z_{1}}$, $\mathbf{z_{2}}$, and $\mathbf{z_{3}}$ are placeholders.

Eq. 6 has a similar formulation to traditional image restoration problem, and thereby we design a residual block ProxBlock composed of three convolutional layers and two ReLU activation to find the solution, as demonstrated in Fig. 3(a).

Further, $\mathcal{D}_{b}(\cdot)$ can be regarded as a quantization step to remove the information of last $b$ bits in the feature space, which is difficult to find an explicit quantization expression. To address this issue, we denote $\mathbf{F}^{Prox}_{k}=\mathbf{F}^{H}_{k}-t_{k}\nabla g(\mathbf{F}^{H}_{k})$, $h_{k}=-t_{k}$, and $G(\mathbf{u},\mathbf{v})=\nabla g(\mathbf{v})$, then there is

where $\mathbf{u}$, $\mathbf{v}$ are placeholders.

Eq. 8 holds a similar formulation to the Euler method for finding an approximation of ordinary differential equation (ODE). Thus, we consider Eq. 8 as a dynamical system, and find the solution by a one-step approximation over an interval $h_{k}$. To find a more accurate numerical solution, we utilize fourth-order Runge-Kutta (RK-4) method [11] to calculate the result and there is

where

Fig. 3(b) shows the design for finding the solution by RK-4 (RK-4 Block). We emulate the $G(\cdot)$ of Eq. 10 by the color blocks of Fig. 3(b) where each block consists of two convolutional layers and a ReLU activation.

As shown in Fig. 3(c), each optimization step consists of a group of one RK-4 Block and one ProxBlock, that optimizes the feature map according to Eq. 9 and Eq. 5, separately. After $m$ optimization steps in the OptBlock, we design a residual block with two convolutional layers and a ReLU activation for better gradient transmission.

There are also three stages in the decoder. For each decoder stage, a $1\times 1$ convolutional layer and a residual block (ResBlock) integrate the information from the current stage of the encoder and the former decoder stage. The ResBlock follows the same design as EDSR [17] which is composed of two convolutional layers and a ReLU activation. The up-scaling method in the decoder is devised by a $1\times 1$ convolutional layers and a sub-pixel convolution [31]. The $1\times 1$ convolution expands the number of channels and the sub-pixel convolution reshapes the feature maps and increases the resolution. At the end of the decoder, there is a block group to restore the weighting map from the feature. Two convolutional layers, a ReLU activation and a Sigmoid activation are utilized to explore the decoded feature maps and generate the weighting map $\mathbf{I}^{WM}$.

In this paper, the filter number is set as 64 for all convolutional layers, except for the $1\times 1$ convolutions and the last convolutional layer. For the first and second stages of the encoder, the number of optimization step is set as 1. For the third stage, there are 6 optimization steps in the OptBlock. The BRNet is designed for bit-depth expansion on standard dynamic range (SDR) image with 8 bits. For HDR bit-depth expansion (with 16 bits), we use two parallel block groups at the end of the decoder for restoration, since it has double information than SDR. Similar to SDR, each group aims to restore at most 8-bit information.

It is obvious that replenishing a small number of bits is easier than restoring many of bits. Inspired by the human’s learning behavior from easy to hard, we propose a progressive training strategy by gradually increasing the missing bits of the training data to boost the learning step [2]. Let $\mathcal{S}=\{\mathcal{X}_{1},...\mathcal{X}_{P}\}$ be $P$ datasets with different difficulty. The difficulty rises with the increase of the indicator $P$. Then, the training steps are regarded as a sequence $\mathcal{C}=<\mathcal{T}_{1},...,\mathcal{T}_{P}>$, where $\mathcal{T}_{P}=(\theta,\mathcal{X}_{P})$, meaning to update the set of network parameter $\theta$ with dataset $\mathcal{X}_{P}$. The indicator $P$ of training steps is sequentially increasing, and the set of network parameters $\theta$ is progressively optimized. Specially, the difficulty of $\mathcal{X}_{P}$ is defined as the upper bound of missing bits. With the increase of missing bits, the information suffers severer loss, and the restoration becomes harder.

The detail of progressive training is as follows. For $E$-th epoch, we generate the LBD image $\mathbf{I}^{L}_{E}$ from the original image $\mathbf{I}^{Ori}_{E}$ for training. Specifically, we increase the difficulty of training data for every 20 epochs. Let $b^{max}$ be the maximum bit depth of image, the upper bound of missing bits for $E$-th epoch $b^{ub}_{E}$ is calculated as $b^{ub}_{E}=\min(4+\lfloor E/20\rfloor,b^{max}-1)$. $b^{max}$ is set as 8 for SDR image, and 16 for HDR image. Then, we randomly remove $b_{E}^{pt}$ bits for progressive training from $\mathbf{I}^{Ori}_{E}$ and generate the LBD image $\mathbf{I}^{L}_{E}=\mathcal{D}_{b^{pt}_{E}}(\mathbf{I}^{Ori}_{E})$, where $b_{E}^{pt}={\rm rand}(1,b^{ub}_{E})$. The overall progressive learning procedure is outlined in Algorithm 1. Overall, we gradually increase the training difficulty to simulate the human’s learning behavior for better performance.

Besides the network design, how to evaluate the performance of restored HBD image is also a vital issue for BDE [27]. By delving into the evaluation method of the human visual system, we hold the hypothesis that if the restored image has the color distribution closer to the natural image, it is more coincident with visual experience. However, the existing PSNR and SSIM are both pixel-wise evaluations [26, 39], having no capability of reflecting the distribution of the whole image. Thus, a distribution distance metric is more suitable for visual performance evaluation.

We use Wasserstein distance (W-dis) to indicate the difference between restored image $\mathbf{I}^{H}$ and the ground truth. Wasserstein distance describes the physical transportation cost between two distributions. Similar to PSNR and SSIM, Wasserstein distance is a symmetric metric, which is more suitable for image quality assessment than KL-divergence and other asymmetrical methods. Fig. 8 shows comparisons between PSNR/SSIM and W-dis. In the first row of Fig. 8, The PSNR of BE-CALF is 0.1 dB higher than RMFNet, but there are more artifacts from BE-CALF in the background of the parrot. RMFNet holds lower W-dis than BE-CALF and has fewer unnatural textures, indicating that W-dis consists with the image quality.

To further investigate the superiority of W-dis than traditional metrics, we choose two representative examples to compare PSNR/SSIM and W-dis, as shown in Figure 5. In the first row of the figure, we focus on the color of the ship. The result of BE-CALF is more different to the original image than the result of RMFNet. Correspondingly, although the PSNR of BE-CALF is higher than RMFNet, the W-dis of RMFNet is lower. Similar situation can be observed in the second row of Figure 5. In the sky area, we can find that L-BDEN generates more artifacts than the RMFNet. Although the SSIM of L-BDEN is higher, the W-dis of RMFNet is lower, and the perceptual experience is better. In this point of view, W-dis can well describe the perceptual difference between two images.

We compare our BRNet with 6 state-of-the-art methods, including a recent traditional method (IPAD [22]) and five CNN-based methods: BitNet [5], LBDEN [47], BE-CALF [19], and RMFNet [20]. It should be noticed that vanilla RMFNet, BE-CALF and LBDEN have no solution for different bit depths. For a fair comparison, we re-train these models with our whole training datasets covering the same bit expansion situations with our model.

Following the comparison in previous works, we first investigate the performances of state-of-the-art methods on restoring 4-bit image to 8-bit on Kodak dataset. Tab. III shows the PSNR/SSIM/W-dis results on Kodak dataset. The best results are shown in bold. “Zero Padding” means to set the missing bits as zero. In the table, our method achieves the best performance than other works in all evaluation terms. Specifically, our network gains 0.3 dB higher than BitPlane and 2 dB higher than RMFNet in term of PSNR. The W-dis of our method is 0.2 lower than BitPlane and RMFNet, which indicates a better visual performance.

Fig. 7 shows the visual comparisons with two competitive methods on restoring 4-bit image to 8-bit. As can be observed, BRNet generates more natural textures on the sky and cloud while other works suffer from artifacts and false contours. In the zoomed-in sky area, our result contains fewer noises and produces more correct colors than other works, which is the closest to the ground truth. It is worth noting that the PSNR/SSIM results of RMFNet is lower than the BE-CALF. In contrast, RMFNet holds better W-dis result than BE-CALF with higher visual quality, indicating W-dis is a suitable metric for visual quality assessment.

To show the general applicability of BRNet, we compare the performances in different bit-depth situations. We set the bit depth as 1, 3, 5, and, 7, and demonstrate the PSNR/SSIM/W-dis results in Tab. III. In the table, BRNet outperforms the state-of-the-art methods in all bit expansion situations. In general, with the increase of bit depth, the superiority of BRNet becomes more obvious. It should be noticed that BitPlane has no capability on restoring 1-bit image. IPAD is a traditional method that intrinsically satisfies different bit-depth situations. Compared with IPAD, BRNet achieves near 4 dB improvements in term of PSNR when the bit-depth is 5.

Fig. 8 shows the visual comparisons on restoring 2-bit to 8-bit. In the figure, our method achieves the best visual performance than other works with superior PSNR/SSIM/W-Dis results. We can find that BRNet can still generate satisfactory results when suffering severe information loss. In the flat area where banding effect is obvious, such as the background of the parrot and the sky over sea, BRNet can remove the false contours and estimate the correct colors and textures more effectively.

To further investigate the effectiveness of BRNet, we compare the color distributions restored by our network and RMFNet respectively, as shown in Fig. 9. The histograms are constructed in red, blue and green channels respectively. The distributions recovered by our network are much close to the ground truth, while the histograms of RMFNet has huge difference with ground truth. The overlapping area between our distribution and the ground truth is larger than that of RMFnet, which explains why we achieve the best W-dis performance.

We also compare the performances on Set5, Set14, and BSD100 datasets, where BRNet outperforms other works near 2 dB, 3 dB, 1.5 dB separately in term of PSNR on restoring 4-bit image to 8-bit. Table IV, V and VI shows the PSNR [35]/SSIM [37]/W-dis [32] comparisons on different datasets separately. We can find that our BRNet achieves the best performance on all testing benchmarks with different settings of bit depth. With the increase of bit depth, BRNet outperforms other works with much higher PSNR results. Specially, when the bit-depth is 7, our BRNet gains more than 5 dB, 10 dB, and 3 dB PSNR improvement than other works on Set5, Set14, and B100 datasets separately.

Figure 10 shows the visual comparisons on restoring images from 2-bit to 8-bit on B100 dataset. In the first row of the figure, we can find that our BRNet produces the most satisfactory visual results. Especially in the sky area, BRNet contains fewer false contours and more correct colors than other methods. In the second row, other works are suffering from the false blurs on the snow, while BRNet can effectively suppress the unpleasant artifacts.

Figure 11 shows the visual comparison on restoring 2-bit image to 8-bit on Set14 dataset. In the figure, BRNet outputs fewer artifacts around the logo than BE-CALF, while LBDEN and RMFNet suffer from the severe false contours. All the competitive methods exhibit obvious banding effect on the blue circle within the green rectangle, while our BRNet has perfect color transition with accurate textures.