VQCNIR: Clearer Night Image Restoration with Vector-Quantized Codebook

Recent advances in deep learning techniques have greatly impacted the field of computer vision. A large number of deep learning methods have been proposed for both single image and video deblurring tasks (2014; 2017; 2018; 2019; 2019; 2020; 2021), and have demonstrated superior performance. With the introduction of large training datasets for deblurring tasks (2009; 2017; 2019), many researchers (2009; 2019) have adopted end-to-end networks to directly recover clear images. Despite the fact that end-to-end methods outperform traditional approaches, they may not be effective in cases with severe blurring. To improve network performance, some methods (2017; 2018; 2021) use multi-scale architectures to enhance deblurring at different scales.

However, the limited ability of these methods to capture the correct blur cues in low-light conditions, particularly in dark areas, has hindered their effectiveness in handling low-light blurred images. To tackle this issue, Zhou et al. (2022) introduce a night image blurring dataset and develop an end-to-end UNet architecture that incorporates a learnable non-linear layer to effectively enhance dark regions without overexposing other areas. Although the method achieves good performance, it is not capable of generalizing enough in real-world scenes, indicating that there is still room for improvement in handling low-light blurred images.

Recent years have witnessed the impressive success of deep learning-based low-light image enhancement (LLIE) since the first pioneering work (2022). Many end-to-end methods (2017; 2019) have been proposed for enhancing image illumination using an encoder-decoder framework. To further improve the performance of LLIEs, researchers have developed deep Retinex-based methods (2018; 2019; 2021) inspired by Retinex theory, which employs dedicated sub-networks to enhance the illuminance and reflectance components and achieve better recovery performance. However, such methods have limitations, as the enhancement results strongly depend on the characteristics of the training data. To improve the generalization ability of the network, researchers (2020; 2021; 2021) propose a number of unsupervised methods. For example, Jiang et al (2021) introduced self-regularization and unpaired training into LLIE with EnlightenGAN. Additionally, Guo et al (2020) propose a fast and flexible method for estimating image enhancement depth curves that do not require any normal illumination reference images during the training process.

VQVAE (2017) is the first to introduce vector quantization (VQ) techniques into an autoencoder-based generative model to achieve superior image generation results. Specifically, the encoded latent variables are quantized to their nearest neighbors in a learnable codebook, and the resulting quantized latent variables are used to reconstruct the data samples. Building upon VQVAE, subsequent work has proposed various improvements to codebook learning. For instance, VQGAN (2021) utilizes generative adversarial learning and refined codebook learning to further enhance the perceptual quality of reconstructed images. The well-trained codebook can serve as a high-quality prior that can be leveraged for various image restoration tasks such as image super-resolution and face restoration. To this end, Chen et al. (2022a) introduce a VQ codebook prior for blind image super-resolution, which matches distorted LR image features with distortion-free HR features from a pre-trained HR prior. Furthermore, Gu et al. (2022) explore the impact of internal codebook properties on reconstruction performance and extended discrete codebook techniques to face image restoration. Drawing inspiration from these works, we apply the high-quality codebook prior to night image restoration.

To improve the recovery of high-quality images with realistic textures and normal illumination from night image $x$ containing complex degradation, we introduce a Vector-Quantized codebook as high-quality prior information to design a night image restoration network (VQCNIR). The overview of the VQCNIR framework is illustrated in Figure 2. VQCNIR comprises an encoder $E$, an adaptive illumination enhancement module, a high-quality codebook $\mathcal{Z}$, and two decoders $G$ and $D$. Decoder G is a pre-trained decoder from VQGAN with fixed parameters. Decoder D represents the primary decoder, which progressively recovers fine details by fusing high-quality features in decoder G.

VQ Codebook: We first briefly describe the VQGAN (2021) model and its codebook, and more details can be referenced in (2021). Given a high-quality image $x_{h}\in\mathbb{R}^{H\times W\times 3}$ with normal light, the encoder $E$ maps the image $x_{h}$ to its spatial latent representation $\hat{z}=E(x)\in\mathbb{R}^{h\times w\times n_{z}}$, where $n_{z}$ is the dimension of latent vectors. Then, each element $\hat{z}_{i}\in\mathbb{R}^{n_{z}}$ Euclidean distance nearest vector $z_{k}$ in the codebook is found as a VQ representation $z_{\textbf{q}}$ by the element-by-element quantization process $\textbf{q}(\cdot)$. It is shown as follows:

where the codebook is $\mathcal{Z}=\{z_{k}\}_{k=1}^{K}\in\mathbb{R}^{K\times n_{z}}$ with $K$ discrete codes. Then, the decoder $G$ maps the quantized representation $z_{\textbf{q}}$ back into sRGB space. The overall reconstruction process can be formulated as follows:

Subsequently, we explore the effectiveness of VQGAN trained on high-quality images for the reconstruction of degraded images at night. As shown in Figure 3 bottom, the restoration image is unable to recover to normal illumination due to the inconsistent illumination of the input image and the training set of VQGAN. Moreover, we found that VQGAN further deteriorates blurred textures and produces artifacts. This was attributed to the difficulty of the network in matching the correct VQ codebook features, which resulted in the inability of VQGAN under high-quality image training to recover from low illumination and blur. Therefore, we design an adaptive illumination enhancement module and a deformable bi-directional cross-attention for the mentioned low light and blur problems respectively.

Based on the previous observations and analysis, we design an Adaptive Illumination Enhancement Module (AIEM) to solve the problem of illumination inconsistency between the quantized features and the latent features obtained from the encoder, as shown in Figure 2. This module consists of two parts: Hierarchical Information Extraction (HIE) and Illumination Mutual Attention Enhancement (IMAE).

As previously described and analyzed, the high-quality quantized features obtained from the codebook are not flawless. The structural warping and textural distortion leads to a more severe misalignment between high-quality VQ codebook features and original degraded features. Therefore, we propose the Deformable Bi-directional Cross-Attention (DBCA) to fuse high-quality VQ codebook features and degraded features.

Unlike the conventional cross-attention method (2021), our DBCA aims to integrate two different features using a bi-directional cross-attention mechanism and employs deformable convolutions to effectively correct the blurring degradation in the degraded feature. As shown in Figure 5, given the input decoder $D$ and $G$ features $F_{D}$ and $F_{G}$, they are first mapped to corresponding $\textbf{Q}_{D}=W_{D}^{p}\text{LN}(F_{D})$ and $\textbf{Q}_{G}=W_{G}^{p}\text{LN}(F_{G})$ via normalization and linear layers. We further utilize linear layers to map these features to corresponding values $\textbf{V}_{D}$ and $\textbf{V}_{G}$. We reshape the aforementioned $\textbf{Q}_{D}$, $\textbf{Q}_{G}$, $\textbf{V}_{D}$, and $\textbf{V}_{G}$ into the shape of $(B,C,H*W)$ and fuse the two features using the following bi-directional cross-attention formula:

where $\text{Softmax}(\cdot)$ denotes the softmax function. $\text{A}_{D}$ and $\text{A}_{G}$ respectively represent the attention maps for feature D and feature G. $\gamma_{D}$ and $\gamma_{G}$ are trainable channel-wise scales and initialized with zeros for stabilizing training.

To better fuse the high-quality codebook prior feature into the degraded feature, we first generate an offset by concatenating the two output features. Then, we use the generated offset in the deformable convolution to distort the texture feature and effectively remove the blurry degradation, which can be formalized as follows:

where $\text{LKConv}(\cdot)$ and $\text{DeformConv}(\cdot)$ denote the $7\times 7$ convolution and the deformable convolution, respectively.

The training objective of VQCNIR comprises four components: (1) pixel reconstruction loss $\mathcal{L}_{pix}$ that minimizes the distance between the outputs and the ground truth; (2) code alignment loss $\mathcal{L}_{ca}$ enforces the codes of the night images to be aligned with the corresponding ground truth; (3) perceptual loss $\mathcal{L}_{per}$ which operates in the feature space, aims to enhance the perceptual quality of the restored images; and (4) adversarial loss $\mathcal{L}_{adv}$ for restoring realistic textures.

Specifically, we adopt the commonly-used L1 loss in the pixel domain as the reconstruction loss, represented by:

where the $x_{h}$ and $x_{n}$ denote high-quality ground truth and night image, respectively. To improve the matching performance of codes for night images with codes for high-quality images. We adopt the $L_{2}$ loss to measure the distance, which can be formulated as:

where $z^{e}$ and $z_{\textbf{q}}^{e}$ are the night image code and the ground truth code, respectively. The total training objective is the combination of the above losses:

where $\lambda_{pix}$, $\lambda_{ca}$, $\lambda_{per}$, and $\lambda_{adv}$ denote the scale factors of each loss function, respectively.

We train our VQCNIR network on the LOL-Blur dataset (2022), which consists of 170 sequences (10,200 pairs) of training data and 30 sequences (1,800 pairs) of test data. We use random rotations of 90, 180, 270, random flips, and random cropping to $256\times 256$ size for the augmented training data. We train our network using Adam (2014) optimizer with $\beta_{1}$=0.9, $\beta_{2}$=0.99 for a total of 500k iterations. The mini-batch size is set to 8. The initial learning rate is set to $1\times 10^{-4}$ and adopts the MultiStepLR to adjust the learning rate progressively. We empirically set $\lambda_{pix}$, $\lambda_{ca}$, $\lambda_{per}$, and $\lambda_{adv}$ to $\{1,1,1,0.1\}$. All experiments are performed on a PC equipped with Intel Core i7-13700K CPU, 32G RAM, and the Nvidia RTX 3090 GPU with CUDA 11.2.

In this section, we compare our proposed VQCNIR quantitatively and qualitatively with all the above methods on the LOL-Blur test set (2022). We use the two most widely evaluated metrics: PSNR and SSIM for a fair evaluation of all methods. In addition, we employ the LPIPS metric to evaluate the perceptual quality of the restored images.

To better illustrate the effectiveness of our method in the real scene, we compare our proposed VQCNIR with the above method quantitatively and qualitatively under the real Real-LOL-Blur dataset (2022). Since the real scene lacks a corresponding reference image to evaluate, three non-reference evaluation metrics were used for the evaluation: MUSIQ (2021), NRQM (2017), and NIQE (2012). The MUSIQ metric assesses mainly color contrast and sharpness, which is more appropriate for this task.

In this section, we have implemented a series of ablation experiments to better validate the effectiveness of each of our proposed modules. To verify the effectiveness of our proposed operations, a series of ablation experiments are presented and the results are shown in Table 3. Initially, we use the VQGAN as our baseline model. Table 3 shows that VQGAN does not effectively address low light and blur degradation, since VQGAN is a codebook prior learned from high-quality natural images and is unable to correctly match degraded features. By designing corresponding parallel decoders, the network can then effectively use high-quality priors to assist in the reconstruction of degraded features. However, the illumination inconsistency between the degraded features and the codebook prior can prevent accurate matching of the high-quality prior features, leading to the occurrence of artifacts. Furthermore, the degraded features are at some distance from high-quality features. Therefore, AIEM and DBCA can be used to effectively improve network performance and image quality. More results of the ablation experiments are provided in the supplementary material.