SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Weijia Wu $^{1,3}$  Yuzhong Zhao $^{2}$  Hao Chen $^{1}$  Yuchao Gu $^{3}$  Rui Zhao $^{3}$  Yefei He $^{1}$  Hong Zhou $^{1*}$  Mike Zheng Shou $^{3*}$  Chunhua Shen $^{1,4}$
<sup>1</sup>Zhejiang University, China <sup>2</sup>University of Chinese Academy of Sciences, China <sup>3</sup>Show Lab, National University of Singapore <sup>4</sup>Ant Group
Current deep networks are very data-hungry and benefit from training on large-scale datasets, which are often time-consuming to collect and annotate. By contrast, synthetic data can be generated infinitely using generative models such as DALL-E and diffusion models, with minimal effort and cost. In this paper, we present DatasetDM, a generic dataset generation model that can produce diverse synthetic images and the corresponding high-quality perception annotations (e.g., segmentation masks, and depth). Our method builds upon the pre-trained diffusion model and extends text-guided image synthesis to perception data generation. We show that the rich latent code of the diffusion model can be effectively decoded as accurate perception annotations using a decoder module. Training the decoder only needs less than  $1\%$  (around 100 images) manually labeled images, enabling the generation of an infinitely large annotated dataset. Then these synthetic data can be used for training various perception models for downstream tasks.
To showcase the power of the proposed approach, we generate datasets with rich dense pixel-wise labels for a wide range of downstream tasks, including semantic segmentation, instance segmentation, and depth estimation. Notably, it achieves (1) state-of-the-art results on semantic segmentation and instance segmentation; (2) significantly more robust on domain generalization than using the real data alone; and state-of-the-art results in zero-shot segmentation setting; and (3) flexibility for efficient application and novel task composition (e.g., image editing).
The project website is at: weijiawu.github.io/DatasetDM.
Modern deep-learning models for perception tasks often require a large amount of labeled data to achieve good performance. Unfortunately, collecting large-scale data and labeling the corresponding pixel-level annotations is a time-consuming and expensive process. For example, collecting images of urban driving scenes requires physical car infrastructure, and labeling a segmentation annotation for a single urban image in Cityscapes [10] can take up to 60 minutes. Moreover, in certain specialized domains, such as medical or human facial data, collecting relevant information can be challenging or even impossible, owing to privacy concerns or other factors. The above challenges can be a barrier to advancing artificial intelligence in computer vision.
To reduce costs, many previous researchers have primarily focused on weakly supervised [57] and unsupervised solutions [51] to address the problem. For instance, certain segmentation priors [1, 2, 32] use weak or inexpensive labels to train robust segmentation models. With the advancement of
*H. Zhou and M. Shou are the corresponding authors.
Figure 1: Synthetic Data from DatasetDM. The high-quality and infinitely numerous synthetic images with perception annotation can yield significant improvements for various downstream tasks.
generative models, such as DALL-E [40], Stable Diffusion [41], some researchers have begun to explore the potential of synthetic data, attempting to use it to assist in training models, or even replace real data for perception task. Most works focus on classification, face recognition [34, 49, 22], salient object detection [54] and segmentation tasks [31, 53, 58], with only a minority trying to address problems such as human pose estimation [18] or medical image analysis [21]. In the era of GANs, DatasetGAN [58] and BigDatasetGAN [29] are recognized as pioneering works that utilize the feature space of pre-trained GANs and design a shallow decoder for generating pixel-level labels in the context of segmentation tasks. Following the two works, Hands-off [55] extends this approach to multitasking scenarios, such as depth estimation. However, these methods still suffer from three major drawbacks: 1) Due to the limitations of the representation ability of early (up to 2020) GAN models, the quality of the synthesized data is often dissatisfactory, leading to an inferior performance on downstream tasks. 2) These methods primarily focus on independent downstream tasks and no one tries to explore a unified data synthesis paradigm with a generalized decoding framework. 3) The training cost is still relatively high, while these methods do not make full use of the visual knowledge contained within the latent codes of powerful text-to-image models.
By leveraging large-scale datasets of image-text pairs (e.g., LAION5B [45]), recent text-to-image diffusion models (e.g., Stable Diffusion [41]) present phenomenal power in generating diverse and high-fidelity images with rich texture, diverse content, and reasonable structures. The phenomenon suggests that large text-to-image diffusion models can implicitly learn valuable, and rich high-level and low-level visual representations from massive image-text pairs. It is natural ask: Can we leverage the knowledge learned by these models to generate perception annotations and extend the paradigm of text-to-image generation to text-to-data generation?
In this paper, built upon the powerful text-to-image diffusion model, we present DatasetDM, a generalized dataset generation model that can produce an unlimited number of synthetic images and the corresponding perception annotations, as shown in Fig. 1. The key to our approach is a unified perception decoder, namely P-Decoder, that decodes the latent code of the diffusion model to perception annotations. Inspired by align/instruct tuning from LLM [50], a method for inducing output following capability with minimal human-labeled, we only use less than  $1\%$  manually labeled images to train the decoder, enabling infinite annotated data generation. The generated data can subsequently be utilized to train any perception models for various downstream tasks, including but not limited to segmentation, depth estimation, and pose estimation. To maintain the robust image generation capabilities of the diffusion model, we freeze the model weights and use image inversion to extract the latent code, i.e., multi-scale features, which are then fed into the unified P-Decoder. The P-Decoder accommodates various downstream tasks within a unified transformer framework, with only minor variations as depicted in Fig. 3.
To summarize, our contributions are four-fold:
- We introduce DatasetDM: a versatile dataset generation model featuring a perception decoder capable of producing an unlimited quantity of high-fidelity synthetic images, along with various perception annotations, including depth, segmentation, and human pose estimation.
- Visual align/instruct tuning, a method for inducing output following capability with minimal human-labeled data. With less than  $1\%$  of the original dataset, i.e., around 100 images,
DatasetDM pushes the limits of text-to-image generation, pioneering a novel paradigm: text-to-data generation. This breakthrough is made possible by leveraging the rich latent code of the pre-trained diffusion model.
- Experiments demonstrate that the existing perception models trained on synthetic data generated by DatasetDM exhibit outstanding performance on six datasets across five different downstream tasks. For instance, the synthetic data delivers remarkable gains of  $13.3\%$  mIoU and  $12.1\%$  AP for semantic segmentation on VOC 2012 and instance segmentation on COCO 2017, respectively.
- Text-guided data generation allows for the generation of a diverse range of data, which has been shown to provide a more robust solution for domain generalization and long-tail data distribution. Moreover, DatasetDM offers a flexible approach for novel task composition, as exemplified by its ability to facilitate image editing (see Fig. 4).
Text-to-Image Generation. Several mainstream methods exist for the task, including Generative Adversarial Networks (GANs)[19], Variational Autoencoders (VAEs)[28], flow-based models [12], and Diffusion Probabilistic Models (DPMs) [47, 41, 25, 20]. Recently, as a likelihood-based method, diffusion models have gained significant attention with promising generation abilities. They match the underlying data distribution by learning to reverse a noising process. Thanks to the high-quality image synthesis and the stability of training, diffusion models are quickly emerging as a new frontier [41, 40, 26, 43] in the field of image synthesis.
Text-guided diffusion models are used for text-conditional image generation, where a text prompt  $\mathcal{P}$  guides the generation of content-related images  $\mathcal{I}$  from a random Gaussian noise  $z$ . Visual and textual embeddings are typically fused using cross-attention. Recent large text-to-image diffusion models, such as Stable Diffusion [41] of Stability AI, DALL-E2 [40] of OpenAI, and Imagen [43] of Google, have shown powerful performance in generating diverse and high-fidelity images with rich textures and reasonable structures. Their impressive synthesis capabilities suggest that these models can implicitly learn valuable representations with the different semantic granularity from large-scale image-text pairs. In this paper, we leverage these learned representations (latent codes), to extend the paradigm of text-to-image generation to text-to-data generation.
Synthetic Datasets for Perception Task. Previous studies in dataset synthesis, such as Virtual KITTI [17] and Virtual KITTI 2 [7], primarily rely on 3D computer simulations to address standard 2D computer vision problems, such as object detection [38], scene understanding [44], and optical flow estimation [6]. However, these methods are limited by the domain of 3D models and cannot be generalized to arbitrary scenes and open-set categories. For instance, Virtual KITTI is exclusively focused on autonomous driving scenes and supports only 20 commonly occurring categories, which cannot be extended to open-scene domains like the COCO benchmark [33].
In contrast, synthetic data generated using generation models (i.e., GAN [19, 35, 55] and Diffusion Model [47]) are more flexible and can support a wider range of tasks and open-world scenes for various downstream tasks, such as classification task [22], face recognition [22], salient object detection [54], semantic segmentation [29, 58, 3, 60], and human pose [18]. Inspired by the success of large-scale generative models, such as Stable Diffusion [41], trained on massive datasets like LAION5B [45], recent studies have begun to explore the potential of powerful pre-trained diffusion generative models. Based on DDPM [25], DatasetDDPM [3] design several CNN layers as the annotation decoder to generate semantic segmentation data by performing a multi-class task. Li et al. [31] utilized Stable Diffusion and Mask R-CNN pre-trained on the COCO dataset [33] to design and train a grounding module for generating images and semantic segmentation masks. DiffuMask [53] produces synthetic image and semantic mask annotation by exploiting the potential of the cross-attention map between text and image from the text-supervised pre-trained Stable Diffusion model. The above methods focus on semantic segmentation, which cannot handle tasks such as instance segmentation. In this paper, we take a further step by utilizing a generalized perception decoder to parse the rich latent space of the pre-trained diffusion model, enabling the generation of perception for a variety of downstream tasks, including depth, segmentation, and human pose.
Diffusion Model for Perception Task. Some recent works [59, 56, 52] has also attempted to directly employ diffusion models for perceptual tasks. VPD [59] explores the direct application of pre-trained
Stable Diffusion to design perception models. ODISE [56] unify pre-trained text-image diffusion and discriminative models to perform open-vocabulary panoptic segmentation. Different from these approaches, our focus lies in synthetic data augmentation for perception tasks, and design a unified transformer-based decoder to enhance more perception tasks from data aspect.
3.1 Formulation
Given a language prompt  $\mathcal{S}$ , text-guided diffusion models generate content-related images  $\mathcal{I} \in \mathcal{R}^{H \times W \times 3}$  from a random Gaussian noise  $z \sim \mathcal{N}(0, \mathbf{I})$ . The standard text-guided image denoising processing can be formulated as:  $\mathcal{I} = \Phi_{\mathrm{T2I}}(z, \mathbb{S})$ , where  $\Phi_{\mathrm{T2I}}(\cdot)$  refers to a pre-trained text-to-image diffusion model. In this paper, we adopt Stable Diffusion [41] as the base for the diffusion model, which consists of three components: a text encoder  $\tau_{\theta}(\cdot)$  for embedding prompt  $\mathcal{S}$ ; a pre-trained variational autoencoder (VAE) [14] that encodes  $\mathcal{E}(\cdot)$  and decodes  $\mathcal{D}(\cdot)$  latent code of images; and a time-conditional UNet  $(\epsilon_{\theta}(\cdot))$  [42] that gradually denoises the latent vectors. To fuse visual and textual embeddings, cross-attention layers are typically used in the UNet for each denoising step. The denoising process is modeled as a Markov chain:  $\boldsymbol{x}_{t-1} = f(\boldsymbol{x}_t, \epsilon_{\theta})$ , where  $\boldsymbol{x}_t$  denote latent code at timestep  $t$ , and  $t \in [1, T]$ . The latent noise at the final timestep  $T$ , denoted as  $\boldsymbol{x}_T$ , is equivalent to the random Gaussian noise  $z$ .  $f(\cdot)$  is the denoising function [25].
In this paper, we design a perception decoder that can effectively parse the latent space of the UNet  $\epsilon_{\theta}(\pmb{x}_t,t,\tau_\theta (\mathcal{S}))$ . By doing so, we extend the text-to-image generation approach to a text-to-data paradigm:
$$
\{\mathcal {I}, \mathcal {P} _ {1: k} \} = \Phi_ {\mathrm {T 2 D}} (\boldsymbol {z}, \mathcal {S}), \tag {1}
$$
where  $\mathcal{P}_{1:k}$  denotes the corresponding perception annotations, and  $k$  is the number of the supported downstream tasks. In fact, the paradigm can support any image-level perception task, such as semantic segmentation, instance segmentation, pose estimation, and depth estimation.
This paper introduces a novel paradigm called text-to-data generation, which extends text-guided diffusion models trained on large-scale image-text pairs. Our key insight is that using a small amount of real data (using less than  $1\%$  existing labeled dataset) and a generic perception decoder to interpret the diffusion latent spaces, results in the generation of infinite and diverse annotated data. Then the synthetic data can be used to train any existing perception methods and apply them to real images.
The proposed DatasetDM framework, presented in Fig. 2, comprises two stages. 1) The first stage—Training—involves using diffusion inversion ( $\S 3.3$ ) to obtain the latent code of the real image and extract the text-image representations ( $\S 3.3$ ). These representations and their corresponding annotations are then used to train the perception decoder ( $\S 3.4$ ). 2) The second stage—Inference ( $\S 3.5$ )—uses GPT-4 to enhance the diversity of data and generates abundant images, while the P-Decoder produces corresponding perception annotations such as masks and depth maps.
3.3 Hypercolumn Representation Extraction
The first step in the training stage of DatasetDM is to extract the hypercolumn representation of real images from the latent space of the diffusion model, as shown in Fig. 2(a). To achieve this, we employ the diffusion inversion technique, which involves adding a certain level of Gaussian noise to the real image and then extracting the features from the UNet during the denoising process.
Image Inversion for Diffusion Model. Give a real image  $\mathcal{X} \in \mathcal{R}^{H \times W \times 3}$  from the training set, the diffusion inversion (diffusion forward process) is a process that approximates the posterior distribution  $q(\pmb{x}_{1:T} | \pmb{x}_0)$ , where  $\pmb{x}_0 = \mathcal{E}(\mathcal{X})$ . This process is not trainable and is based on a fixed Markov chain that gradually adds Gaussian noise to the image, following a pre-defined noise schedule  $\beta_1, \ldots, \beta_T$  [25]:
$$
q \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x} _ {t - 1}\right) := \mathbb {N} \left(\boldsymbol {x} _ {t}; \sqrt {1 - \beta_ {t}} \boldsymbol {x} _ {t - 1}, \beta_ {t} \mathbf {I}\right), \tag {2}
$$
Figure 2: The overall framework of DatasetDM. DatasetDM consists of two main steps: 1) Training. Using diffusion inversion to extract the latent code from a small amount of data and then train the perception decoder. 2) Text-guided data generation. A large language model such as GPT-4 is utilized to prompt infinite and diverse data generation for various downstream tasks.
where  $t$  represents the  $t$ -th time step, and we set it to 1 for a single forward pass in our paper. A single forward pass for visual representation extraction usually provides two advantages, i.e., faster convergence and better performance [56].
Text-Image Representation Fusion. With the latent code  $\boldsymbol{x}_t$  of the real image and the corresponding language prompt  $S$ , we extract the multi-scale feature maps and cross attention maps from the UNet  $\epsilon_{\theta}$  as follows:
$$
\{\mathcal {F}, \mathcal {A} \} = \epsilon_ {\theta} (\boldsymbol {x} _ {t}, t, \tau_ {\theta} (\mathcal {S})), \tag {3}
$$
where  $\mathcal{S}$  for training set is simply defined using a template "a photo of a  $[\mathrm{CLS}]_1$ ,  $[\mathrm{CLS}]_2$ , ...". During the data generation phase, GPT-4 is used to provide diverse prompt languages. The variable  $\mathcal{F}$  denotes the multi-scale feature maps from four layers of the U-Net, corresponding to four different resolutions, i.e.,  $8 \times 8$ ,  $16 \times 16$ ,  $32 \times 32$ , and  $64 \times 64$ , as illustrated in Fig. 2. Additionally,  $\mathcal{A}$  represents the cross-attention maps of text-to-image from the 16 cross-attention layers in the U-Net, which implement the function  $\mathcal{A} = \text{softmax}\left(\frac{QK^T}{\sqrt{d}}\right)$ , where  $d$  is the latent projection dimension. We group the 16 cross-attention maps into 4 groups with the same resolutions, and compute their average within each group, which results in the average cross-attention maps  $\hat{\mathcal{A}}$ .
Prior works [53, 59, 24] have proved the effectiveness of class-discriminative and localization of cross-attention map between the visual embedding and the conditioning text features. Thus we concatenate the cross-attention maps  $\hat{\mathcal{A}}$  and the multi-scale feature maps  $\mathcal{F}$  to obtain the final extracted hyper-column representation, and further use a  $1\times 1$  convolution to fuse them:  $\hat{\mathcal{F}} = \mathrm{Conv}([\mathcal{F},\hat{\mathcal{A}}])$
3.4 Perception Decoder
The P-Decoder is utilized to translate the representation  $\hat{\mathcal{F}}$  into perception annotations, which are not limited to a specific type for each downstream task. To achieve this goal, inspired by previous works [8, 61], we devised a generalized architecture. This architecture is depicted in Fig. 3, with only minor variations (i.e., whether to startup some layers) for each downstream task. For example, the pixel decoder and transformer decoder are required for generic segmentation, while only the pixel decoder is necessary for depth and pose estimation.
Generic Image Segmentation. In Fig. 3-(a), we present the adaptation for semantic and instance segmentation tasks, which includes two components: the pixel decoder and the transformer decoder.
(a) Generic Segmentation
(b) Open-Vocabulary Segmentation
Figure 3: Various types of tasks with our proposed P-Decoder. The proposed decoder is a generalized architecture for the six supported tasks, with only minor variations required for different downstream applications, i.e., determining whether to activate certain layers.
(c) Depth/Pose Estimation
Similar to Mask2former [8], the transformer decoder comprises a stack of transformer layers with cross-attention and self-attention. This component refines the queries and renders the outputs. The pixel decoder is made up of multiple upsampling CNN layers, and it is used to obtain per-pixel embeddings. Given the representation  $\hat{\mathcal{F}}$  and  $N$  learnable queues  $\{Q_0, Q_1 \dots Q_T\}$  as input, the decoder outputs  $N$  binary masks  $\mathbf{O} = o_1, \dots, o_N \in \{0, 1\}^{N \times H \times W}$ , along with the corresponding category. This is achieved through simple matrix multiplication between the outputs of the transformer decoder and the pixel decoder. Following Mask2former [8], one query is responsible for a class or instance for semantic and instance segmentation, respectively. To optimize the mask prediction, we use the binary cross-entropy loss [9] and dice loss [37]. On the other hand, the cross-entropy loss is used for classification prediction.
Open-Vocabulary Segmentation. Based on the architecture of generic segmentation, a semantics-related query concept is proposed to support the open-vocabulary segmentation task, as shown in Fig 3-(b). During the training stage, given the category name  $\{C_0,C_1\dots C_K\}$  (e.g., dog, cat) on an image from the training set, we use the text encoder  $\tau_{\theta}(\cdot)$ , e.g., CLIP [39] to encode them into the embedding space. Then concatenating them with the queries to equip the semantics of class vocabulary to the query as follows:
$$
\hat {Q} _ {i} = \operatorname {M L P} \left(\left[ Q _ {i}, \tau_ {\theta} \left(C _ {j}\right) \right]\right), \tag {4}
$$
where  $Q_{i}$  and  $C_{j}$  is the  $i$ -th query embedding and  $j$ -th class. MLP refers to a learned MLP, used to fuse the class embedding and learnable query embedding. Thus DatasetDM can generate an open-vocabulary mask by incorporating a new class name, as illustrated in Fig. 5 (b).
Depth and Pose Estimation. For depth and pose estimation, the output format is predetermined, eliminating the need to differentiate between classes or instances. In this context, the pixel decoder is only required to predict a fixed number of maps  $\mathbf{O} \in \mathcal{R}^{M \times H \times W}$ . The value of  $M$  is set to either 1 or 17 (corresponding to 17 human key points), depending on whether the task is depth or human pose estimation. For human pose estimation, we use the mean squared error as the loss function, and the ground truth heatmaps are generated by applying a 2D Gaussian with a standard deviation of 1 pixel centered on each key point. As for depth estimation, we update the loss function from the classic scale-invariant error [30, 13].
3.5 Prompting Text-Guided Data Generation
In Fig. 2 (b), we present the inference pipeline for text-guided data generation. There are two main differences compared to the training phase: firstly, the prompts come from a large language model instead of a fixed template, and secondly, the denoising process is extended to  $T$  steps to obtain the synthetic images. Large language model, i.e., GPT-4, is adopted to enhance the diversity of generative data, while recent works [16, 23, 4] have proven their powerful understanding and adaptability for the real world. As shown in Fig. 4, instead of the template-based prompts from humans, we guide GPT-4 to produce diverse, and infinite prompts. For different downstream tasks and datasets, we give different guided prompts for GPT-4. For example, as for the urban scene of Cityscapes [10], the simple guided prompt is like 'Please provide 100 language descriptions of urban driving scenes for the Cityscapes benchmark, containing a minimum of 15 words each. These descriptions will serve as a guide for Stable Diffusion in generating images.' In this approach, we collected  $L$
Collaboration between the creative thinker, the painter, the annotator and editor
Human: Please provide multiple descriptions for the appearance of one person with a focus on their clothing, pants, and shoes. These descriptions will serve as a guide
for Stable Diffusion in generating corresponding images.