Blind Face Restoration for Under-Display Camera via Dictionary Guided Transformer

Existing methods usually use the UDC degradation model in Eq. (1) to synthesize UDC-degraded images. However, there exists a big gap between this approximate degradation model and the real UDC degradation. For example, the parameters in degradation (intensity scale factor $\alpha$ and PSF $k$) are different in UDC configurations, which requires extensive human effort to obtain these parameters in advance [10, 5, 9]. Thus, we design an end-to-end UDC degradation modeling network called UDC-DMNet to directly learn the UDC degradation process with the existing real-world paired data.

As shown in Fig. 2, our UDC-DMNet models the UDC degradation process as two stages [8] i.e., color filtering and brightness attenuation modeling, and diffraction modeling. Correspondingly, our UDC-DMNet is a two-stage network, where the first stage is used for color filtering and brightness attenuation modeling, and the second stage is designed for diffraction modeling.

Color Filtering and Brightness Attenuation Modeling. When the light is normally incident on the OLED display, part of the light will be absorbed by the display, because the thin-film layers of an OLED have different transmission rates for different wavelengths. Thus, a clean image first exhibits color filtering and brightness attenuation in the UDC imaging process and the degradation is highly related to the channels of the image. To model this degradation process, as shown in Fig. 2, we first use one $3\times 3$ convolutional layer in the first stage of our UDC-DMNet to extract the initial features. Then, those features are fed into an encoder-decoder architecture with four downsampling and upsampling operations to derive multi-scale information. In each scale, we use a Residual Channel Attention Block (RCAB) [47, 48] to adaptively learn the process of brightness attenuation. In particular, RCAB consists of three convolution layers and one channel attention layer. With the help of RCAB, the first stage of our UDC-DMNet can fully simulate the brightness and color of the UDC image.

Diffraction Modeling. In the second stage of UDC imaging, the light will be incident on the pixel definition layer (PDL). PDL that has a periodic window pattern acts as a planar diffraction grating [49, 50]. And images formed under PDL will be blurry. Previous works [10, 5, 9] usually estimate the PSF to simulate this process. However, PSF is highly complicated, and the measured PSF is usually not universal. Thus, in this work, we also choose encoder-decoder architecture to model this diffraction process. To better cover the PSF of various scales, we use a convolution with the stride of 2 to downsample and use a transposed convolution to upsample the feature maps. Specifically, in each scale, we adopt a residual block with spatial attention [48] named Residual Spatial Attention Block (RSAB) as the basic layer to build the second stage of our UDC-DMNet. With the help of RSAB, our UDC-DMNet can well simulate the blur caused by spatial changes in the second stage.

Moreover, considering that the degradation of the first stage may affect the diffraction of the second stage and a long series of convolution operations are likely to cause blurring of edge features. Thus, as shown in Fig. 2, we propose a Cross-Stage Spatial Feature Transform (CSSFT) to connect two stages of our UDC-DMNet. CSSFT can help simplify the information flow among stages to address the above problem [51, 52]. Specifically, CSSFT can be formulated as follows,

where $Conv1(\cdot)$ and $Conv2(\cdot)$ are both $1\times 1$ convolution. $F_{e}$ and $F_{d}$ are the same scale features from the encoder and decoder in the first stage, respectively. ${\alpha}_{e}$, ${\alpha}_{d}$ are the parameters of the spatial feature transform, which can be defined as

where $F_{in}$ is the output of RSAB, and $F_{out}$ represents the output of the spatial feature transform.

In addition, as shown in Fig. 2, we introduce a Degradation Fusion Block (DFB) between the first and second stages of our UDC-DMNet to fuse color-filtering degraded features with the input image, which makes the second stage of UDC-DMNet better focus on diffraction modeling.

We first select the Peak Signal-to-Noise Ratio (PSNR) loss as the main supervision [53]. However, the PSNR loss function merely measures the pixel-wise distance and cannot properly describe the similarity of the degradation of two images. The generated degraded image will be over-smooth. Therefore, we further select perceptual loss [54] and adversarial loss [55] to learn the degradation patterns in training. We adopt three losses to train UDC-DMNet: PSNR loss $\mathcal{L}_{\text{psnr}}$, perceptual loss $\mathcal{L}_{\text{per}}$ and adversarial loss $\mathcal{L}_{\text{adv}}$. The PSNR loss is defined as:

where $\hat{\mathbf{X}}$ and $\mathbf{I}$ refer to the generated degraded image and ground-truth degraded image respectively. Then the perceptual loss is used to help our model produce visually pleasing results. We adopt a pre-trained VGG19 [56] to extract the perceptual features from the $Conv5\_4$ layer of VGG19, and then use the $L_{1}$ loss function to calculate the difference in the feature space between the generated degraded image and their corresponding ground-truth degraded image. Specifically, the perceptual loss is as follows:

We use adversarial loss on the output of the generator of UDC-DMNet to distinguish real and fake degraded images. This is helpful to make the UDC-DMNet model the degradation process better and narrow the gap between the generated distribution and the real-world distribution. Specifically, we adopt the discriminator $\mathcal{D}$ of PatchGAN [57] for adversarial training:

The final loss function $\mathcal{L}$ to train our proposed UDC-DMNet is shown as follows:

where ${\lambda}_{1}$ and ${\lambda}_{2}$ are hyper-parameters used to balance these three losses. In our experiments, they are both set to $1$.

We train our UDC-DMNet on the P-OLED and T-OLED datasets [23], respectively. Both datasets have $300$ pairs of degraded UDC images captured using a UDC device in the real world and the corresponding ground truth images. Each image is of resolution $1024\times 2048$. After the model training, we leverage the FFHQ dataset [19], which consists of $70,000$ high-quality face images as inputs to the trained UDC-DMNet to synthesize UDC face datasets, FFHQ-P and FFHQ-T, which will be specifically employed for the training of our proposed UDC face image restoration model illustrated in the following. For evaluation of UDC face image restoration, we use high-quality face images in CelebA-Test dataset [21] as inputs of the trained UDC-DMNet to synthesize CelebA-Test-P and CelebA-Test-T.

As shown in Fig. 3, DGFormer consists of four modules: shallow feature extractor, UDC restoration module, dictionary-guided restoration module, and image reconstruction module. Specifically, give a degraded UDC face image $X\in\mathbb{R}^{H\times W\times 3}$, the shallow feature extractor, consisting of a simple $3\times 3$ convolution, first extracts features $F_{0}\in\mathbb{R}^{H\times W\times C}$ from the input image $X$. Then, the feature $F_{0}$ is fed into the subsequent UDC restoration module to achieve coarse-grained UDC degradation removal. The UDC restoration module is composed of several specially designed transformer blocks. After that, the output feature of the UDC restoration module is forwarded into the dictionary-guided restoration module for fine-grained restoration. The dictionary-guided restoration module is a symmetric encoder-decoder structure, which is mainly built by our proposed dictionary-guided transformer blocks. Finally, the fine-grained feature is reconstructed to a clear image by the image reconstruction module. The UDC restoration module, dictionary-guided restoration module, and image reconstruction module are core modules of our DGFormer, thus we detail them in the following.

The UDC image usually has diverse and complicated degradation (e.g., low light, blur, noise etc.), thus directly restoring clear images from UDC images is difficult. In addition, the complex degradation in UDC input images causes difficulty in extracting face priors, which influences the performance. Thus, we propose a UDC Restoration Module as the successor of the shallow extractor to achieve coarse-grained UDC degradation removal. Previous work [4] claims that the UDC restoration network should consider global and local information in the network design because removing low light and color filtering degradation in UDC images requires global information in each color channel, and local information is more favorable for noise and blur removal. Thus, as shown in Fig. 3, we propose a new transformer block as the basic unit to build our UDC Restoration Module. The proposed transformer block considers both global and local information for feature processing. Specifically, in the proposed transformer block, we first use a pixel-wise convolution [58] and a depth-wise convolution [59] to generate query, key, and value which are helpful to enrich the local context. Then, we perform self-attention across channels [60] to generate an attention map encoding the global context implicitly. Finally, we resort to the gating mechanism in the feed-forward network named Gating Feed Forward Network (GFFN) to further improve the expressive capability of networks. Specifically, the gating mechanism is formulated as the element-wise product of two parallel paths. In the first path, we adopt a pixel-wise convolution and a depth-wise convolution to strengthen the locality. In the second path, we adopt a pixel-wise convolution followed by a Sigmoid function as a gating signal. The UCD restoration module can be formulated as

where $Transformer(\cdot)$ is the UDC Restoration module that contains several transformer blocks, $F_{0}$ represents output features of the shallow feature extractor, and $F_{1}$ is output features.

To achieve fine-grained UDC face restoration, we exploit the face dictionary as extra facial prior knowledge in the proposed dictionary-guided restoration module to further help restore face details of the face images. The dictionary-guided restoration module is designed under an encoder-decoder structure with skip connections, which can fully exploit both multi-scale features and face dictionary prior for effective face restoration. In our dictionary-guided restoration module, we propose a novel Dictionary-Guided Transformer Block (DGTB) as the basic block of the encoder-decoder structure to restore the facial details of the image. As shown in Fig. 3, the DGTB consists of a Dictionary-Guide Muti-head Attention (DGMA), a GFFN in the UDC Restoration Module, and several transformer blocks. DGMA has facial component dictionaries and queries these dictionaries through Dictionary Transform Block (DTB). By storing the high-quality textures of the eyes, nose, and mouth in the facial component dictionary in advance, we can significantly improve the model’s ability to restore face details.

DGMA is our core component, which is shown in Fig. 5. The facial component dictionary $Dic$ in each DGMA can be expressed as follows

where $y$ is the number of key values in the dictionary. To make sure the dictionary can be end-to-end optimized by the final loss, we set the $(key,value)$ in the dictionary as learnable parameters. The size of the $key$ and $value$ are related to the scale where DGMA is located. Taking the eye dictionary as an example, in the DGMA with an input feature size of $(\frac{H}{s}\times\frac{W}{s}\times sC)$, the size of each key and value is $(\frac{A^{eye}}{s}\times\frac{A^{eye}}{s})$, where $A^{eye}$ is the size of the detected eye landmark box. Supposing the input feature of DGMA is $F_{n}$, DGMA first takes the obtained $F_{n}$ into DTB. DTB first obtains facial component features $F^{d}_{n}$ such as eyes and mouth from $F_{n}$ through RoIAlign operation and detected facial landmarks. Then query features $q^{d}_{n}$ is obtained through a series of convolution operations on features $F^{d}_{n}$. Then DTB queries the component dictionary in a similar way as the attention mechanism. After that, DTB fuses the original component features and the component features queried in the dictionary, then copies and pastes fused component features $\hat{F^{d}_{n}}$ into the original features $F_{n}$ through reverse RoIAlign operation to obtain features ${F_{n}^{dic}}$. Then, instead of using self-attention directly to features ${F_{n}^{dic}}$, our DGMA takes input features $F_{n}$ as queries $Q=W_{p}W_{d}BN(F_{n})$, while the keys $K=W_{p}W_{d}BN(F_{n}^{dic})$ and values $V=W_{p}W_{d}BN(F_{n}^{dic})$ are generated from output feature of DTB $F_{n}^{dic}$, where $W_{p}$ is a pixel-wise convolution and $W_{d}$ is a depth-wise convolution. Finally, we adopt the attention mechanism to conduct face restoration. Overall, the DGMA process is defined as

where $\beta$ is a learnable scaling parameter. With the guidance of the face dictionary prior, our DGMA can fuse the information from the features $F_{n}$ and its corresponding high-quality face texture features $F_{n}^{dic}$. They can respectively provide identity information and high-quality facial details for face restoration. While normal self-attention usually can only consider the feature of a single information source.

For the DGTB of each scale, after using a DGMA and a GFFN, we choose to use several transformer blocks mentioned in UDCRM to refine the features of this scale instead of continuing to use DGMA. Because we find that the use of the DGMA query dictionary repeatedly can cause model performance to decline.

Following the Dictionary Guided Restoration Module, deep features $F_{d}\in\mathbb{R}^{H\times W\times C}$ are put into the image reconstruction module. In this module, we use four consecutive transformer blocks and a $3\times 3$ convolution to transform $F_{d}$ into the final clear face image. To simplify our DGFormer, we use the same transformer block that is used in the UDC restoration module.

To train our DGFormer, we use L1 loss between the recovered images and the ground truth images.

In the ablation of UDC-DMNet, we analyze the effectiveness of its components. We use UDC-DMNet with different settings to synthesize different P-OLED datasets to train an identical restoration model, NAFNet [65]. For a fair comparison, the model is trained on different datasets by $200$ epochs with the same training parameters. After training, we evaluate the restoration model on real-world P-OLED UDC images. Table IV reports the performance of different settings.

We first perform ablation experiments on the loss function. When using only PSNR loss to train the network, the performance of the network declines significantly. The reason is that images generated by UDC-DMNet are too smooth, and UDC-DMNet fails to generate any noise degradation in the generated images. The restoration model trained with such synthetic data cannot remove UDC degradation well. Then ablation experiments are carried out on perceptual loss and adversarial loss respectively. Compared to using only PSNR loss, adding both losses can improve the utility of the synthetic dataset and make the restoration model trained by it perform better.

Secondly, we perform ablation experiments on each component in the UDC-DMNet structure. Specifically, we first remove RSAB and RCAB respectively, adopting a normal residual block [68] in each stage. Removing RSAB causes a drop of $0.99$ dB in terms of PSNR and removing RCAB causes a drop of $1.09$ dB in terms of PSNR. It demonstrates that both kinds of attention blocks are useful for simulating degradation. Then we remove the Cross-Stage Spatial Feature Transform (CSSFT), and the performance drops by about $0.45$ dB regarding PSNR. Then we remove the Degradation Fusion Block (DFB), and the performance drops by about $1.16$ dB in terms of PSNR. All the ablation studies demonstrate the effectiveness of our proposed UDC-DMNet for UDC image synthesis. Furthermore, we modify the network to single-stage to explore the effectiveness of multi-stage modeling of the degradation process. This leads to a significant drop in performance, indicating that multi-stage modeling affects the synthesis performance significantly. The qualitative results for different configurations are shown in Fig. 10. UDC-DMNET, without any of the concerned components, can hardly simulate real-world UDC degradation. As shown in the figure, this leads to poor performance of the restoration model trained on these synthetic datasets.

We further analyze and discuss the effectiveness of the internal modules in our DGFormer. Table V shows the results, where all models are trained on FFHQ-P and tested on CelebA-Test-P.

Firstly, we replace the transformer blocks with residual blocks [69, 68] for training, leading to a significant drop in performance. It implies that transformer blocks are important for the integration of global information and local information. Secondly, we remove the UDC Restoration Module (UDCRM) and Image Reconstruction Module (IRM) respectively, which both result in a performance drop, implying that both coarse-grained UDC degradation removal and image feature refinement are indispensable. Thirdly, we remove the Dictionary Transform Block (DTB) in each Dictionary-Guided Transformer Block (DGTB). In this way, we will not use the facial component dictionary in the network. This also causes performance drop, and the facial details are easily lost, demonstrating that dictionaries are helpful in the restoration of facial component details. Fourth, we replace DGMA with normal self-attention, which results in a performance drop. This shows that DGMA, which combines low-quality features and features enhanced by prior information, is more conducive to image restoration than self-attention using only a single feature as an information source. Finally, we replace the normal transformer block in DGTB with DGMA + GFFN, causing a performance drop, which shows that the repeated use of the dictionary has a negative impact on the restoration of details.

The visual results of varying configurations are presented in Fig. 11. Regarding ‘use Resblocks’ and ‘w/o UDCRM’ in the figure, the absence of a transformer structure and UDCRM cannot fully eliminate UDC degradation, leading to blurring and noise artifacts in the restored images. For ‘w/o DTB’ in the figure, the restoration model without the dictionary component struggles to effectively restore eye details. In addition, as shown in the ‘Replace transformer block’ of the figure, when we replace the transformer block in DGTB with DGMA + GFFN, the eyes of the restored image become blurrier. Furthermore, in comparison with ‘w/o IRM’ and ‘Replace DGMA’, images restored by full DGFormer showcase more comprehensive and detailed eye features. This illustrates the refinement effect of IRM and DGMA on image features.