TextDiff: Mask-Guided Residual Diffusion Models for Scene Text Image Super-Resolution

Single image super-resolution (SISR) is a fundamental low-level task in computer vision, which goal is to generate HR images from LR inputs. As the seminal work of image super-resolution, SRCNN (Dong et al. 2014) achieves good performance through a simple convolutional network. Besides, in order to further improve the quality of super-resolution (SR) outputs, various methods (Liang et al. 2021; Saharia et al. 2022; Chen et al. 2023) have been proposed. Among them, methods based on the deep generative model, mainly including GAN-based (Wang et al. 2018; Soh et al. 2019; Wang et al. 2021) and diffusion-based methods (Saharia et al. 2022; Whang et al. 2022), have shown convincing image generation ability. However, the method for SISR tasks is not suitable for STISR. The key is that the method of SISR does not consider the structural characteristics of scene text.

Different from SISR, STISR aims to improve the quality of the image while paying attention to the recovery of the text structure. TSRN proposes a real scene text SR dataset, and uses CNN-BiLSTM layers to perceive the sequential information of the text. Following it, (Chen, Li, and Xue 2021) proposes a Transformer-Based Super-Resolution Network (TBSRN) to extract sequential information, designs a Position-Aware Module and a Content-Aware Module to highlight the position and the content of each character. In addition, TPGSR (Ma, Guo, and Zhang 2023) obtains the character probability sequence through a text recognition model and merges it with image feature by convolutions. Different from these, by combining the text mask and graphic recognition results of LR text images, DPMN (Zhu et al. 2023) proposes a plug-and-play module, which can improve the performance of existing models. While these methods have made significant advances, they do not yet fully address issues like text distortion and blurring, which can impair image readlibity. Therefore, our method aims to resolve these specific problems.

Diffusion models are latent variable generative frameworks in (Sohl-Dickstein et al. 2015). Recent work has demonstrated the significant potential of diffusion models in SISR. For example, (Li et al. 2022) exploits a Markov chain to convert HR images to latents in simple distribution and then generate SR predictions in the reverse process. However, to the best of our knowledge, diffusion models have not yet been used in STISR. Therefore, in this paper, we explore the performance of diffusion models on STISR for the first time.

The overall framework of TextDiff is illustrated in Figure 3. First, the LR image is input into the Text Enhancement Module (TEM). The TEM consists of two branches, $\textrm{B}_{T}$ and $\textrm{B}_{M}$. $\textrm{B}_{T}$ outputs a coarsely deblurred image, and $\textrm{B}_{M}$ learns to predict text masks. Then, guided by the mask, the Mask-Guided Residual Diffusion Module (MRD) learns the distribution of residuals between the ground-truth image and the coarsely deblurred image. Finally, this predicted residual is added to the coarsely deblurred image to obtain the final output.

As shown in Figure 3, the branch $\textrm{B}_{T}$ is a deblurring network incorporated with semantic information. The front part, called Semantic Block, adopts TP Generator (TPG) from TATT (Ma, Liang, and Zhang 2022) to extract semantic priors, which are then fused with feature maps through TP Interpreter (TPI) in TATT. The addition of this part can effectively integrate the semantic features of the text with the spatial distribution features of the text. The fused results are finally input into the Double Sequential Residual Block (DSRB). DSRB introduces wavelet transform on the basis of SRB (Wang et al. 2020) to realize simultaneous learning of spatial and frequency information. This aims to complement the spatial domain details with frequency domain information to generate an output with richer frequency details.

Additionally, the blurred regions in the scene text image mainly focus on the text. We need a text mask to separate the text from the background, and then use the mask to focus the network on the text area. So the branch $\textrm{B}_{M}$ is for mask prediction, implemented by a conventional convolutional network (see Figure 3). And the ground-truth masks are simply generated by calculating the average gray scale of the RGB images.

We define the Gradient Profile Loss as the pixel loss between the $\textrm{B}_{T}$ output and the ground-truth image to recover the approximate information of the global text structure in the LR image:

where $\nabla x_{gt}$ denotes the gradient field of HR images, and $\nabla x_{sr}$ denotes that of SR images.

Meanwhile, we define the loss for mask learning as dice loss (Wang et al. 2019). So the $\mathcal{L}_{Mask}$ is used to consider the contour similarity between the $\textrm{B}_{M}$ output $x_{m}$ and the ground-truth text mask $x_{gt_{m}}$. Dice loss can be calculated as the following Eq.2 and Eq.3:

$P_{x,y}$ and $G_{x,y}$ represent the pixel value (x, y) of the predicted masks and the ground-truths, respectively.

After the LR image is processed by the TEM, $x_{sr}$ and a text mask are obtained. As shown in Figure 3, $x_{sr}$ has restored the general outline and color of the text and other visual features, but compared with the ground-truth image, the edges of the text are still blurred and partially distorted. The residual $x_{res}$ between the $x_{sr}$ and $x_{gt}$ (i.e., $x_{res}=x_{gt}-x_{sr}$) delineates the text outline, as depicted in Figure 3. Thus, we leverage the diffusion model to refine the text contour under the guidance of the text mask to make it more accurate and natural. Experiments show that this residual modeling effectively alleviates the limitations of the regression model in learning text contours, enabling the generation of more perceptually pleasing text.

Our MRD consists of a diffusion process of progressively adding Gaussian noise and a reverse process of learning residual distributions for denoising.

Specifically, the diffusion process starts from the residual image $x_{res}$, also denoted as $x_{0}$. Then it repeatedly adds Gaussian noise according to the transition kernel $q(x_{t}\mid x_{t-1})$. And at the maximum time step t, we obtain $x_{T}$ which is pure Gaussian noise:

The noise schedule $\alpha_{t}$ is a pre-chosen hyperparameter that controls the variance of noise added at each step. Setting $\alpha_{t}\in(0,1)$ for all $t=1,...,T$ , $\alpha_{0}=1$, $\bar{\alpha}_{t}={\textstyle\prod_{i=0}^{t}}\alpha_{i}$, the diffusion process allows sampling $x_{t}$ at an arbitrary timestep t in closed form:

The reverse process converts noise $x_{T}$ back into data distribution $x_{0}$ conditioned on $x_{sr}$ and $x_{m}$. We adopt a deterministic manner to conduct the process:

where $\mu_{t}(x_{t},x_{0})$ is calculated as:

Furthermore, with $x_{0}$ as an unknown parameter, the reverse diffusion step can be implemented by substituting the estimate $f_{\theta}$ in place of $x_{0}$.

In Eq.10, $x_{m}$ plays an important role. Looking the U-Net structure in Figure 3, we introduced gated dilated convolution (Yang, Xiong, and Wu 2023). When processing the input image depth features and text mask, the gating mechanism learns to generate a gating vector that places more attention on the text regions, while dilated convolution enlarges the receptive field. The combination of both allows effective fusion of multi-scale features and enhances the restoration effect for text structure.

Moreover, $f_{\theta}$ predicts $x_{0}$ instead of noise. The rationale behind this includes: first, cascading both $x_{sr}$ and $x_{m}$ as conditional inputs allows $f_{\theta}$ to emphasize textual regions when predicting $x_{0}$; second, predicting noise and $x_{0}$ can be converted via Eq.6; third, predicting $x_{0}$ can better exploit $x_{sr}$ and $x_{m}$ as guiding conditions compared to predicting noise, which only depends on $x_{t}$. Meanwhile, predicting noise is more likely to cause diversity that compromises textual structure.

In practice, $f_{\theta}$ ensures that the learned conditional distribution $p_{\theta}(x_{t-1}\mid x_{t},x_{sr},x_{m})$ approximates the true reverse diffusion step $q(x_{t-1}\mid x_{t},x_{0})$ as closely as possible. Hence, the training objective is:

Besides, in order to better restore the contour edges of text, we extract textual structural edge information from the $f_{\theta}(x_{t},t,x_{sr},x_{m})$ and $x_{0}$ using a Laplacian kernel $f_{Edge}$ to compute the edge loss. This encourages the model to reconstruct sharper and more coherent textual contours:

The overall loss function for training the network is comprised of three parts: the loss from the TEM, the loss from the MRD, and a joint loss $\mathcal{L}_{joint}$ combining both modules.

where $\mathcal{L}_{TP}$ is a text prior loss for TPG (Ma, Liang, and Zhang 2022), $\lambda_{1}$ and $\lambda_{2}$ are hyperparameters. Specifically, the joint loss of the two modules is formulated as:

where $R_{i}$ denotes the feature map output from the CRNN (Shi, Bai, and Yao 2016) after the i-th activation layer. Since CRNN is an optical character recognition model capable of perceiving textual patterns, we utilize its publicly available pre-trained weights to extract features without additional training on TextZoom. Incorporating this CRNN-based loss allows the network to better restore textual content.

Thus, the overall loss function is

where $\lambda$ is a hyperparameter.

The training and inference procedures are presented in Algorithm 1 and Algorithm 2, respectively.

We conduct training and evaluation on the TextZoom (Wang et al. 2020) which is collected in real-world scenarios. This dataset consists of 17,367 LR-HR image pairs in the training set. The test set is divided into easy, medium, and hard subsets comprising 1,619, 1,411, and 1,343 LR-HR pairs respectively, based on the camera focal length. The size of LR images is $16\times 64$, while the size of HR images is $32\times 128$.

Our model is implemented with PyTorch 2.0 deep learning library and all the experiments are conducted on one RTX 4090 GPU. AdamW (Loshchilov and Hutter 2017) is utilized as the optimizer with a learning rate of $1\times 10^{-4}$, and the batch size is set to 16. The total time steps $T$ are set to 200. The number of training iterations is one million. We use a linear increase in $\beta_{1:T}$ from $1\times 10^{-6}$ to $1\times 10^{-2}$, $\alpha_{t}=1-\beta_{t}$. Moreover, the two stages of the proposed network are trained jointly, where the weight coefficient $\lambda$ for the joint loss is set to 5. The weights for $\mathcal{L}_{GP}$ and $\mathcal{L}_{Mask}$ are set to 0.5 and 3 in Eq.13, respectively. Additional model configuration information is given in the supplementary material.

Following the common practice in TSRN, to measure the downstream task performance, we calculate the text recognition accuracy for the text recognition task. Additionally, We evaluate the quality of the SR images using metrics including the no-reference image quality assessment method MANIQA (Yang et al. 2022), the full-reference method LPIPS (Zhang et al. 2018), PSNR and SSIM.

We evaluate our TextDiff and compare it with existing super-resolution models, including SRCNN (Dong et al. 2015), SRResNet (Ledig et al. 2017), TSRN (Wang et al. 2020), TBSRN (Chen, Li, and Xue 2021), PCAN (Zhao et al. 2021), TPGSR (Ma, Guo, and Zhang 2023), DocDiff (Yang et al. 2023), TATT (Ma, Liang, and Zhang 2022), C3-STISR (Zhao et al. 2022) and DPMN (Zhu et al. 2023).

As shown in Table 1, we can clearly observe that our TextDiff achieves a significant improvement in accuracy compared to existing methods. On average, our method improves recognition accuracy by 2.3%, 2.2% and 0.5% on ASTER, MORAN and CRNN, respectively. Furthermore, TextDiff outperforms all existing methods in recognition accuracy with only 5-step sampling (see the quantitative results taking other sampling steps in the supplementary material.). The consistent accuracy gains on multiple models validate the effectiveness of our TextDiff. For comprehensiveness, we also compare TextDiff with SOTA diffusion-based document enhancement method (DocDiff), and the results show the superiority of our method for STISR.

Additionally, as shown in Figure 4, most methods have insufficient capability in recovering textual structures, whereas our method can effectively alleviate this problem. We also give quantitative results for image quality evaluation, as shown in Table 2. Our TextDiff achieves the best MANIQA and LPIPS metrics, while also achieving competitive PSNR and SSIM. Notably, we obtain the LPIPS of 0.0822, a 12% reduction compared to C3-STISR and a 10% reduction compared to TATT.

Finally, we list the failure cases of TextDiff and our future work in the supplementary material.

In this section, we perform ablation studies to analyze the contribution of different motivations and model components. All the evaluations are validated on TextZoom. The quantitative and qualitative results of the ablation experiments are provided in Table 3 and in the supplementary material, respectively.

We explore cascading our MRD module with other mainstream approaches to further validate its efficacy and extensibility. Specifically, we achieve joint inference without any joint training, simply by direct inference. The inference results are shown in Table 4 and Figure 1 (see additional cascade inference results in the supplementary material.). The results sufficiently demonstrate that the effectiveness of MRD comes not only from the algorithm itself, but also from its ability to complement other types of methods. For instance, from the cascaded results of MRD with TSRN, it is evident that MRD effectively complements the incomplete font structure obtained from TSRN, leading to results that are more aligned with human perception. As a result, the recognition accuracy of TSRN increases by 2%, without any additional training. Overall, embedding MRD as a post-processing module into existing pipelines can provide further performance boost.

We mainly discuss the applicability of PSNR and SSIM metrics for the STISR. As shown in Figure 5, we find that for phrases such as “hat&sturdy”, the PSNR and SSIM of images obtained by using existing methods are relatively high, but the recovered fonts are distorted or blurred. In contrast, the image recovered by TextDiff can perfectly restore the font structure, but the PSNR and SSIM are low. Whether it is from the perspective of human perception or model recognition results, our method is relatively good. This also confirms the quantitative results in Table 2. In summary, as discussed in (Wang et al. 2020) and (Yang et al. 2023), it can be concluded that PSNR and SSIM metrics only partially align with human perception when applied to text images.