Mobile Image Restoration via Prior Quantization

Any digital imaging system suffers from optical aberration, and thus correcting aberration is necessary for accurate measurements [1]. Unfortunately, the optical aberration in image is affected by many factors, which is formulated as:

here $(h,w)$ indicates the coordinates of the pixel, $\lambda$ is the wavelength, $\mathcal{C}_{e}$ and $I_{e}$ denote the spectral response and the point spread function (PSF), $L_{e}$, $J_{e}$, and $N_{e}$ are the latent sharp image, the observed image, and the measurement noise, respectively. Note that the subscript $e$ indicates the measurement that represents the energy received by the sensor. Therefore, energy dispersion and field-of-view (FoV) clue of the optical system, spectral properties, and sensor noise together contribute to the optical aberration expression on digital images. In other words, these factors serve as the priors to facilitate the correction.

Although there are many algorithms for optical aberration correction, it still faces a few challenges for widespread application. One issue is that the existing methods generally have a limited application scope, e.g., the deconvolution is inefficient in handling spatial-varying kernels, and the deep learning method is trained for a specific device according to the data [2]. Another challenge is the quality of restoration, which is generally unsatisfactory due to the lack of sufficient information. An inherent defect of these unflexible methods is that they are incapable of mining the interaction between the digital pixel and other optical priors [3].

As mentioned above, the expression of optical aberration correlates with multiple factors. Designing a general method to utilize these priors for targeted correction is an issue worth discussing. Meanwhile, this baseline can help analyze the correlation between different priors and optical aberration correction, aiming to guide the co-design of the hardware configuration and the post-processing pipeline in the high-end imaging system.

In this letter, we propose a prior quantization model to correct the optical degradation influenced by multiple factors, where different priors corresponding to each factor are fed to the model. However, due to the content differences among the multimodal priors, they cannot directly integrate into the model for efficient post-processing. To this end, we encode the various priors into a high-dimensional latent space and characterize it by a learnable codebook. The learned code reduces the dimension of multimodal priors representation and models the global interrelations of auxiliary information. Then the model integrates the quantized prior representation into the image restoration branch by fusing it with features of different scales. Finally, we supervise the restoration in multiple scales to ensure that the cross-scale information used for fusion is accurate and valid. Instead of a black box, the model can adjust the input pixel-level prior according to the demands of the user, aiming to achieve a targeted development. Moreover, the learned codebook bridges the gap between different priors and restoration. Therefore, the proposed model has the potential to analyze the utilization of priors, which is meaningful for imaging system design.

As mentioned above, multiple factors contribute to the expression of optical aberration. Thus the model must possess the ability to perceive a complex combination and transform it into a simpler form. Directly engaging the image signal with individual priors is inefficient, where the model will spend substantial computing overhead on the auxiliary information. Therefore, we encode these priors into a high-dimension space and represent them with a learnable codebook. Since the aberration is relatively constant in the neighborhood, the optical prior of an image $x\in\mathcal{R}^{H\times W\times 3}$ can be represented by a spatial collection of codebook $z_{q}\in\mathcal{R}^{h\times w\times d_{z}}$, where $d_{z}$ is the dimension of the code. To effectively learn such a codebook, we propose to exploit the spatial extraction capability of convolution and incorporate with the neural discrete representation learning [4]. First, we use CNNs to encode the priors into a latent space with the same spatial resolution ($h\times w$). Then, we embed the encoded features into the spatial code $\hat{z}_{ij}\in\mathcal{R}^{d_{z}}$ and quantify each code onto its closest entry in the discrete codebook $\mathcal{Z}=\{z_{k}\}^{K}_{k=1}\subset\mathcal{R}^{n_{z}}$:

Third, we concatenate the quantized prior representation with the features in the restoration branch. And the subsequent ResBlock actively fused the image feature, where the prior after quantization provides clues of the spatial/channel importance to the reconstruction branch.

In most reconstruction models, the skip connections are only processed in one scale when the coarse-to-fine architecture is applied. Inspired by the dense connection between multi-scale features, we implement a multi-scale feature fusion (MFF) module to fuse the quantified representation with the features from other scales [5]. As shown in the right-bottom of Fig. 1, this module receives the outputs of different encoding scales ($ES_{i}$) as inputs. After adjusting the channels of each $ES^{out}_{i}$ by convolution, we use pixel shuffle to transform the encoded information into the same spatial resolution and then perform the fusion [6]. The output of the MFF is delivered to its corresponding decoding scales. In this way, each scale can perceive the encoded information of other scales, especially the lowest-scale features filtered by the quantified priors, resulting in improved restoration quality.

Here we discuss the loss function of our model. Due to the quantization operation in Eq. 2 is non-differentiable in backpropagation, the CNN encoder of priors cannot receive a gradient to optimize its parameters if the entire network is directly trained with end-to-end supervision. Fortunately, the strategy of the codebook alignment allows the gradient propagated from decoders to update these CNN encoders. Specifically, we keep the encoded priors approaching the vectors of the learnable codebook. In this way, the quantization progress like a gradient estimator, allowing the CNN encoder to estimate the backpropagation and update the parameters even when the gradient is truncated in training. Thus, we apply the codebook alignment loss to realize the optimization of encoders [7]:

here $\hat{z}_{ij}=\{E_{1}(p_{1}),E_{2}(p_{2}),\dots,E_{n}(p_{n})\}$, where $p_{i}$ and $E_{i}$ are the $i^{th}$ prior and encoder, respectively. $\{\cdot\}$ is the operation to concatenate the encoded features. $sg[\cdot]$ denotes the stop-gradient operation to ensure the loss function only guide the encoded priors and the codebook vectors approaching.

Because our network relies on the refinement in a coarse-to-fine manner, we supervise the image reconstruction on different scales. For the supervision of image content, we find that the L1 loss performs better on quantitative metrics for optical aberration correction. However, since the content loss only measures the pixel-level difference, the model does a good job restoring the low-frequency information (such as brightness and color, etc.) while performing poorly in restoring the textures of the scene. To prevent the limitation of pixel-level supervision, we adopt the supervision in the Fourier domain as an auxiliary loss [5]. Compared with the gradient constraints (e.g., total variation) or the feature similarity on the pre-trained model (e.g., perceptual loss), this item provides more information on different spatial frequencies. The loss $\mathcal{L}_{content}$ is formulated as follows:

here the output of the k-th scale is $I_{k}$ and its corresponding ground-truth is $\hat{I_{k}}$, where $K$ is the number of scales. The content loss in each scale is the average on the total elements, whose number is denoted as $t_{k}$. $\mathcal{F}$ is the fast Fourier transform (FFT) operation that performs on different scales and the $\lambda$ is experimentally set to 0.1. The overall loss function is the sum of $\mathcal{L}_{content}$ and $\mathcal{L}_{align}$.

In summary, we develop a deep learning model to utilize multiple priors for optical aberration correction. Even though the content of the priors varies a lot from each other, the proposed quantization strategy efficiently fuses these factors and guides the correction of optical aberration. Comprehensive experiments show that integrating all the related information does benefit the restoration of optical aberration, and the proposed model can generalize to different mobile devices without finetuning. Moreover, the learned model quantifies the correlation and significance of different priors for the correction procedure. Thus, the aberration correction in the post-processing system can be more efficient when the critical clues are obtained. In the future, we will put efforts into engaging the proposed model with different imaging devices, aiming to realize efficient correction and better generalization.