SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jingye Chen\*13, Yupan Huang\*23, Tengchao Lv\*3, Lei Cui\*3, Qifeng Chen\*1, Furu Wei\*3
$^{1}$ HKUST  $^{2}$ Sun Yat-sen University  $^{3}$ Microsoft Research
qwerty.chen@connect.ust.hk, huangyp28@mail2.sysu.edu.cn, cqf@ust.hk
{tengchaolv, lecu, fuwei}@microsoft.com
Diffusion models have gained increasing attention for their impressive generation abilities but currently struggle with rendering accurate and coherent text. To address this issue, we introduce TextDiffuser, focusing on generating images with visually appealing text that is coherent with backgrounds. TextDiffuser consists of two stages: first, a Transformer model generates the layout of keywords extracted from text prompts, and then diffusion models generate images conditioned on the text prompt and the generated layout. Additionally, we contribute the first large-scale text images dataset with OCR annotations, MARIO-10M, containing 10 million image-text pairs with text recognition, detection, and character-level segmentation annotations. We further collect the MARIO-Eval benchmark to serve as a comprehensive tool for evaluating text rendering quality. Through experiments and user studies, we show that TextDiffuser is flexible and controllable to create high-quality text images using text prompts alone or together with text template images, and conduct text inpainting to reconstruct incomplete images with text. The code, model, and dataset will be available at https://aka.ms/textdiffuser.
(b) Text-to-Image with Template
(c) Text Inpainting
Figure 1: TextDiffuser generates accurate and coherent text images from text prompts or together with template images, as well as conducting text inpainting to reconstruct incomplete images.
*Equal contribution during internship at Microsoft Research.
The field of image generation has seen tremendous progress with the advent of diffusion models [2, 15, 16, 18, 25, 67, 70, 72, 79, 93] and the availability of large-scale image-text paired datasets [17, 74, 75]. However, existing diffusion models still face challenges in generating visually pleasing text on images, and there is currently no specialized large-scale dataset for this purpose. The ability of AI models to generate accurate and coherent text on images is crucial, given the widespread use of text images in various forms (e.g., posters, book covers, memes, etc.) and the difficulty in creating high-quality text images, which typically require professional skills and numerous times of designers.
Traditional solutions to creating text images involve using image processing tools like Photoshop to add text onto images directly. However, these often result in unnatural artifacts due to the background's complex texture or lighting variations. Recent efforts have used diffusion models to overcome the limitations of traditional methods and enhance text rendering quality. For instance, Imagen [72], eDiff-I [2], and DeepFloyd [12] observe diffusion models generate text better with T5 series text encoders [61] than the CLIP text encoder [60]. Liu et al. employ character-aware text encoders to improve text rendering [46]. Despite some success, these models only focus on text encoders, lacking control over the generation process. A concurrent work, GlyphDraw [49], improves the controllability of models by conditioning on the location and structures of Chinese characters. However, GlyphDraw does not support multiple text bounding-box generation, which is not applicable to many text images such as posters and book covers.
In this paper, we propose TextDiffuser, a flexible and controllable framework based on diffusion models. The framework consists of two stages. In the first stage, we use a Layout Transformer to locate the coordinates of each keyword in text prompts and obtain character-level segmentation masks. In the second stage, we fine-tune the latent diffusion model by leveraging the generated segmentation masks as conditions for the diffusion process and text prompts. We introduce a character-aware loss in the latent space to further improve the quality of generated text regions. Figure 1 illustrates the application of TextDiffuser in generating accurate and coherent text images using text prompts alone or text template images. Additionally, TextDiffuser is capable of performing text inpainting $^2$  to reconstruct incomplete images with text. To train our model, we use OCR tools and design filtering strategies to obtain 10 million high-quality image-text pairs with OCR annotations (dubbed as MARIO-10M), each with recognition, detection, and character-level segmentation annotations. Extensive experiments and user studies demonstrate the superiority of the proposed TextDiffuser over existing methods on the constructed benchmark MARIO-Eval. The code, model and dataset will be publicly available to promote future research.
Text Rendering. Image generation has made significant progress with the advent of diffusion models [18, 25, 67, 72, 79, 63, 70, 2, 8, 13, 52, 48, 26, 80], achieving state-of-the-art results compared with previous GAN-based approaches [64, 100, 43, 58]. Despite rapid development, current methods still struggle with rendering accurate and coherent text. To mitigate this, Imagen [72], eDiff-I [2], and DeepFolyd [12] utilize a large-scale language model (large T5 [61]) to enhance the text-spelling knowledge. In [46], the authors noticed that existing text encoders are blind to token length and trained a character-aware variant to alleviate this problem. A concurrent work, GlyphDraw [49], focuses on generating high-quality images with Chinese texts with the guidance of text location and glyph images. Unlike this work, we utilize Transformer [81] to obtain the layouts of keywords, enabling the generation of texts in multiple lines. Besides, we use character-level segmentation masks as prior, which can be easily controlled (e.g., by providing a template image) to meet user needs.
Several papers have put forward benchmarks containing a few cases regarding text rendering for evaluation. For example, Imagen [72] introduces DrawBench containing 200 prompts, in which 21 prompts are related to visual text rendering (e.g., A storefront with 'Hello World' written on it). According to [46], the authors proposed DrawText comprising creative 175 prompts (e.g., letter 'c' made from cactus, high-quality photo). GlyphDraw [49] designs 218 prompts in Chinese and
2Different from text editing [88, 37, 32], the introduced text inpainting task aims to add or modify text guided by users, ensuring that the inpainted text has a reasonable style (i.e., no need to match the style of the original text during modification exactly) and is coherent with backgrounds.
English (e.g., Logo for a chain of grocery stores with the name 'Grocery'). Considering that existing benchmarks only contain a limited number of cases, we attempt to collect more prompts and combine them with existing prompts to establish a larger benchmark MARIO-Eval to facilitate comprehensive comparisons for future work.
Image inpainting. Image inpainting is the task of reconstructing missing areas in images naturally and coherently. Early research focused on leveraging low-level image structure and texture to address this task [3, 6, 5]. Later, deep learning architectures such as auto-encoder [55, 45], GAN [65, 98], VAE [103, 105], and auto-regressive Transformers [57, 83, 96] were applied to tackle this problem. Recently, diffusion models have been used to generate high-quality and diverse results for unconditional image inpainting [71, 48, 67, 11, 99], text-conditional image inpainting [52, 1] and image-conditional image inpainting [92]. Our work falls under the category of text-conditional image inpainting using diffusion models. In contrast to prior works that focused on completing images with natural backgrounds or objects, our method focuses on completing images with text-related rendering, also named text inpainting, by additional conditioning on a character-level segmentation mask.
Optical Character Recognition. Optical Character Recognition (OCR) is an important task that has been studied in academia for a long period [87, 7]. It has undergone a remarkable development in the last decade, contributing to many applications like autonomous driving [89, 73], car license plate recognition [101, 53], GPT models [76, 28], etc. Various datasets [31, 19, 90, 91] and downstream tasks are included within this field, such as text image recognition [77, 41, 95, 14], detection [106, 51, 42, 84], segmentation [90, 91, 107, 68], super-resolution [9, 85, 50, 104], as well as some generation tasks, including text image editing [88, 78, 94, 59, 38], document layout generation [56, 22, 20, 40, 33], font generation [30, 21, 36, 54], etc. Among them, the font generation task is most relevant to our task. Font generation aims to create high-quality, aesthetically pleasing fonts based on given character images. In contrast, our task is more challenging as it requires the generated text to be legible, visually appealing, and coherent with the background in various scenarios.
As illustrated in Figure 2, TextDiffuser consists of two stages: Layout Generation and Image Generation. We will detail the two stages and introduce the inference process next.
3.1 Stage1: Layout Generation
In this stage, the objective is to utilize bounding boxes to determine the layout of keywords (enclosed with quotes specified by user prompts). Inspired by Layout Transformer [20], we utilize the Transformer architecture to obtain the layout of keywords. Formally, we denote the tokenized prompt as  $\mathcal{P} = (p_0,p_1,\dots,p_{L - 1})$ , where  $L$  means the maximum length of tokens. Following LDM [67], we use CLIP [60] and two linear layers to encode the sequence as  $\mathrm{CLIP}(\mathcal{P})\in \mathbb{R}^{L\times d}$ , where  $d$  is the dimension of latent space. To distinguish the keywords against others, we design a keyword embedding  $\mathrm{Key}(\mathcal{P})\in \mathbb{R}^{L\times d}$  with two entries (i.e., keywords and non-keywords). Furthermore, we encode the width of keywords with an embedding layer  $\mathrm{Width}(\mathcal{P})\in \mathbb{R}^{L\times d}$ . Together with the learnable positional embedding  $\mathrm{Pos}(\mathcal{P})\in \mathbb{R}^{L\times d}$  introduced in [81], we construct the whole embedding as follows:
$$
\operatorname {E m b e d d i n g} (\mathcal {P}) = \operatorname {C L I P} (\mathcal {P}) + \operatorname {P o s} (\mathcal {P}) + \operatorname {K e y} (\mathcal {P}) + \operatorname {W i d t h} (\mathcal {P}). \tag {1}
$$
The embedding is further processed with Transformer-based  $l$ -layer encoder  $\Phi_E$  and decoder  $\Phi_D$  to get the bounding boxes  $\mathbf{B} \in \mathbb{R}^{K \times 4}$  of  $K$  key words autoregressively:
$$
\mathbf {B} = \Phi_ {D} \left(\Phi_ {E} (\operatorname {E m b e d d i n g} (\mathcal {P}))\right) = \left(\mathbf {b} _ {0}, \mathbf {b} _ {1}, \dots , \mathbf {b} _ {K - 1}\right). \tag {2}
$$
Specifically, we use positional embedding as the query for the Transformer decoder  $\Phi_D$ , ensuring that the  $n$ -th query corresponds to the  $n$ -th keyword in the prompt. The model is optimized with  $l1$  loss, also denoted as  $|\mathbf{B}_{GT} - \mathbf{B}|$  where  $\mathbf{B}_{GT}$  is the ground truth. Further, we can utilize some Python packages like Pillow to render the texts and meanwhile obtain the character-level segmentation mask  $\mathbf{C}$  with  $|\mathcal{A}|$  channels, where  $|\mathcal{A}|$  denote the size of alphabet  $\mathcal{A}$ . To this end, we obtain the layouts of keywords and the image generation process is introduced next.
Figure 2: TextDiffuser consists of two stages. In the first Layout Generation stage, a Transformer-based encoder-decoder model generates character-level segmentation masks that indicate the layout of keywords in images from text prompts. In the second Image Generation stage, a diffusion model generates images conditioned on noisy features, segmentation masks, feature masks, and masked features (from left to right) along with text prompts. The feature masks can cover the entire or part of the image, corresponding to whole-image and part-image generation. The diffusion model learns to denoise features progressively with a denoising and character-aware loss. Please note that the diffusion model operates in the latent space, but we use the image pixels for better visualization.
3.2 Stage2: Image Generation
In this stage, we aim to generate the image guided by the segmentation masks  $\mathbf{C}$  produced in the first stage. We use VAE [35] to encode the original image with shape  $H\times W$  into 4-D latent space features  $\mathbf{F}\in \mathbb{R}^{4\times H^{\prime}\times W^{\prime}}$ . Then we sample a time step  $T\sim \mathrm{Uniform}(0,T_{\max})$  and sample a Gaussian noise  $\epsilon \in \mathbb{R}^{4\times H^{\prime}\times W^{\prime}}$  to corrupt the original feature, yielding  $\hat{\mathbf{F}} = \sqrt{\bar{\alpha}_T}\mathbf{F} + \sqrt{1 - \bar{\alpha}_T}\epsilon$  where  $\bar{\alpha}_{T}$  is the coefficient of the diffusion process introduced in [25]. Also, we downsample the character-level segmentation mask  $\mathbf{C}$  with three convolution layers, yielding 8-D  $\hat{\mathbf{C}}\in \mathbb{R}^{8\times H^{\prime}\times W^{\prime}}$ . We also introduce two additional features, called 1-D feature mask  $\hat{\mathbf{M}}\in \mathbb{R}^{1\times H^{\prime}\times W^{\prime}}$  and 4-D masked feature  $\hat{\mathbf{F}}_M\in \mathbb{R}^{4\times H^{\prime}\times W^{\prime}}$ . In the process of whole-image generation,  $\hat{\mathbf{M}}$  is set to cover all regions of the feature and  $\hat{\mathbf{F}}_M$  is the feature of a fully masked image. In the process of part-image generation (also called text inpainting), the feature mask  $\hat{\mathbf{M}}$  represents the region where the user wants to generate, while the masked feature  $\hat{\mathbf{F}}_M$  indicates the region that the user wants to preserve. To simultaneously train two branches, we use a masking strategy where a sample is fully masked with a probability of  $\sigma$  and partially masked with a probability of  $1 - \sigma$ . We concatenate  $\hat{\mathbf{F}},\hat{\mathbf{C}},\hat{\mathbf{M}},\hat{\mathbf{F}}_M$  in the feature channel as a 17-D input and use denoising loss between the sampled noise  $\epsilon$  and the predicted noise  $\epsilon_{\theta}$ :
$$
l _ {\text {d e n o i s i n g}} = \left\| \epsilon - \epsilon_ {\theta} \left(\hat {\mathbf {F}}, \hat {\mathbf {C}}, \hat {\mathbf {M}}, \hat {\mathbf {F}} _ {M}, \mathcal {P}, T\right) \right\| _ {2} ^ {2}. \tag {3}
$$
Furthermore, we propose a character-aware loss to help the model focus more on text regions. In detail, we pre-train a U-Net [69] that can map latent features to character-level segmentation masks. During training, we fix its parameters and only use it to provide guidance by using a cross-entropy loss  $l_{char}$  with weight  $\lambda_{char}$  (See more details in Appendix A). Overall, the model is optimized with
$$
l = l _ {\text {d e n o i s i n g}} + \lambda_ {\text {c h a r}} * l _ {\text {c h a r}}. \tag {4}
$$
Finally, the output features are fed into the VAE decoder to obtain the images.
(a) MARIO-LAION
Figure 3: Illustrations of three subsets of MARIO-10M. See more details in Appendix C.
(b) MARIO-TMDB
(c) MARIO-OpenLibrary
3.3 Inference Stage
TextDiffuser provides a high degree of controllability and flexibility during inference in the following ways: (1) Generate images from user prompts. Notably, the user can modify the generated layout or edit the text to meet their personalized requirements; (2) The user can directly start from the second stage by providing a template image (e.g., a scene image, handwritten image, or printed image), and a segmentation model is pre-trained to obtain the character-level segmentation masks (Appendix B); (3) Users can modify the text regions of a given image using text inpainting. Moreover, this operation can be performed multiple times. These experimental results will be presented in the next section.
4 MARIO Dataset and Benchmark
As there is no large-scale dataset designed explicitly for text rendering, to mitigate this issue, we collect 10 million image-text pairs with OCR annotations to construct the MARIO-10M Dataset. We further collect the MARIO-Eval Benchmark from the subset of the MARIO-10M test set and other existing sources to serve as a comprehensive tool for evaluating text rendering quality.
4.1 MARIO-10M Dataset
The MARIO-10M is a collection of about 10 million high-quality and diverse image-text pairs from various data sources such as natural images, posters, and book covers. Figure 3 illustrates some examples from the dataset. We design automatic schemes and strict filtering rules to construct annotations and clean noisy data (more details in Appendix D and Appendix E). The dataset contains comprehensive OCR annotations for each image, including text detection, recognition, and character-level segmentation annotations. Specifically, we use DB [42] for detection, PARSeq [4] for recognition, and manually train a U-Net [69] for segmentation. We analyze the performance of OCR tools in Appendix F. The total size of MARIO-10M is 10,061,720, from which we randomly chose 10,000,000 samples as the training set and 61,720 as the testing set. MARIO-10M is collected from three data sources:
MARIO-LAION derives from the large-scale datasets LAION-400M [75]. After filtering, we obtained 9,194,613 high-quality text images with corresponding captions. This dataset comprises a broad range of text images, including advertisements, notes, posters, covers, memes, logos, etc.
MARIO-TMDB derives from The Movie Database (TMDB), which is a community-built database for movies and TV shows with high-quality posters. We filter 343,423 English posters using the TMDB API out of 759,859 collected samples. Since each image has no off-the-shelf captions, we use prompt templates to construct the captions according to movie titles.
MARIO-OpenLibrary derives from Open Library, which is an open, editable library catalog that creates a web page for each published book. We first collect 6,352,989 original-size Open Library covers in bulk. Then, we obtained 523,684 higher-quality images after filtering. Like MARIO-TMDB, we manually construct captions using titles due to the lack of off-the-shelf captions.
4.2 MARIO-Eval Benchmark
The MARIO-Eval benchmark serves as a comprehensive tool for evaluating text rendering quality collected from the subset of the MARIO-10M test set and other sources. It comprises 5,414 prompts in total, including 21 prompts from DrawBenchText [72], 175 prompts from DrawTextCreative [46], 218 prompts from ChineseDrawText [49] and 5,000 image-text pairs from a subset of the MARIO-10M test set. The 5,000 image-text pairs are divided into three sets of 4,000, 500, and 500 pairs, and are named LAIONEval4000, TMDBEval500, and OpenLibraryEval500 based on their respective data sources. We offer examples in Appendix G to provide a clearer understanding of MARIO-Eval.
Evaluation Criteria: We evaluate text rendering quality with MARIO-Eval from four aspects: (1) Fréchet Inception Distance (FID) [24] compares the distribution of generated images with the distribution of real images. (2) CLIPScore calculates the cosine similarity between the image and text representations from CLIP [29, 60, 23]. (3) OCR Evaluation utilizes existing OCR tools to detect and recognize text regions in the generated images. Accuracy, Precision, Recall, and F-measure are metrics to evaluate whether keywords appear in the generated images. (4) Human Evaluation is conducted by inviting human evaluators to rate the text rendering quality of generated images using questionnaires. More explanations are shown in Appendix H.
For the first stage, we utilize the pre-trained CLIP [60] to obtain the embedding of given prompts. The number of Transformer layers  $l$  is set to 2, and the dimension of latent space  $d$  is set to 512. The maximum length of tokens  $L$  is set to 77 following CLIP [60]. We leverage a commonly used font "Arial.ttf" and set the font size to 24 to obtain the width embedding and also use this font for rendering. The alphabet  $\mathcal{A}$  comprises 95 characters, including 26 uppercase letters, 26 lowercase letters, 10 digits, 32 punctuation marks, and a space character. After tokenization, only the first substring is marked as the keyword when several subtokens exist for a word.
For the second stage, we implement the diffusion process using Hugging Face Diffusers [82] and load the checkpoint "runwayml/stable-diffusion-v1-5". Notably, we only need to modify the input dimension of the input convolution layer (from 4 to 17), allowing our model to have a similar scale of parameters and computational time as the original model. In detail, the height  $H$  and  $W$  of input and output images are 512. For the diffusion process, the input is with spatial dimension  $H' = 64$  and  $W' = 64$ . We set the batch size to 768 and trained the model for two epochs, taking four days using 8 Tesla V100 GPUs with 32GB memory. We use the AdamW optimizer [47] and set the learning rate to 1e-5. Additionally, we utilize gradient checkpoint [10] and xformers [39] for computational efficiency. During training, we follow [25] to set the maximum time step  $T_{max}$  to 1,000, and the caption is dropped with a probability of  $10\%$  for classifier-free guidance [27]. When training the part-image generation branch, the detected text box is masked with a likelihood of  $50\%$ . We use 50 sampling steps during inference and classifier-free guidance with a scale of 7.5 following [67].
Number of Transformer layers and the effectiveness of width embedding. We conduct ablation studies on the number of Transformer layers and whether to use width embedding in the Layout Transformer. The results are shown in Table 1. All ablated models are trained on the training set of MARIO-10M and evaluated on its test set. Results show that adding width embedding improves performance, boosting IoU by  $2.1\%$ ,  $2.9\%$ , and  $0.3\%$  when the number of Transformer layers  $l$  is set to 1, 2, and 4, respectively. The optimal IoU is achieved using two Transformer layers and the width embedding is included. See more visualization results in Appendix I.
Character-level segmentation masks provide explicit guidance for generating characters. The character-level segmentation masks provide explicit guidance on the position and content of characters during the generation process of TextDiffuser. To validate the effectiveness of using character-level segmentation masks, we train ablated models without using the masks and show results in Appendix
Table 1: Ablation about Layout Transformer.
#Layer	Width(P)	IoU↑
1	-	0.268
✓	0.289
2	-	0.269
✓	0.298
4	-	0.294
✓	0.297
Table 2: Ablation on weight of character-aware loss.
λchar	Acc↑
0	0.396
0.001	0.486
0.01	0.494
0.1	0.420
1	0.400
Table 3: Ablation on two-branch training ratio  $\sigma$  .
ratio	Acc↑/Det-F↑/Spot-F↑
0	0.344 / 0.870 / 0.663
0.25	0.562 / 0.899 / 0.636
0.5	0.552 / 0.881 / 0.715
0.75	0.524 / 0.921 / 0.695
1	0.494 / 0.380 / 0.218
Table 4: The performance of text-to-image compared with existing methods. TextDiffuser performs the best regarding CLIPScore and OCR evaluation while achieving comparable performance on FID.
Metrics	StableDiffusion [67]	ControlNet [102]	DeepFloyd [12]	TextDiffuser
FID↓	51.295	51.485	34.902	38.758
CLIPScore↑	0.3015	0.3424	0.3267	0.3436
OCR(Accuracy)↑	0.0003	0.2390	0.0262	0.5609
OCR(Precision)↑	0.0173	0.5211	0.1450	0.7846
OCR(Recall)↑	0.0280	0.6707	0.2245	0.7802
OCR(F-measure)↑	0.0214	0.5865	0.1762	0.7824
J. The generated texts are inaccurate and not coherent with the background compared with texts generated with TextDiffuser, highlighting the importance of explicit guidance.