SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Quang Nguyen<sup>1,2*</sup> Truong Vu<sup>1*</sup> Anh Tran<sup>1</sup> Khoi Nguyen<sup>1</sup>
$^{1}$ VinAI Research,  $^{2}$ Ho Chi Minh City University of Technology, VNU-HCM
Preparing training data for deep vision models is a labor-intensive task. To address this, generative models have emerged as an effective solution for generating synthetic data. While current generative models produce image-level category labels, we propose a novel method for generating pixel-level semantic segmentation labels using the text-to-image generative model Stable Diffusion (SD). By utilizing the text prompts, cross-attention, and self-attention of SD, we introduce three new techniques: class-prompt appending, class-prompt cross-attention, and self-attention exponentiation. These techniques enable us to generate segmentation maps corresponding to synthetic images. These maps serve as pseudo-labels for training semantic segmenters, eliminating the need for labor-intensive pixel-wise annotation. To account for the imperfections in our pseudo-labels, we incorporate uncertainty regions into the segmentation, allowing us to disregard loss from those regions. We conduct evaluations on two datasets, PASCAL VOC and MSCOCO, and our approach significantly outperforms concurrent work. Our benchmarks and code will be released at https://github.com/VinAIResearch/Dataset-Diffusion.
Semantic segmentation is a fundamental task in computer vision. Its objective is to assign semantic labels to each pixel in an image, making it crucial for applications such as autonomous driving, scene comprehension, and object recognition. However, one of the primary challenges in semantic segmentation is the high cost associated with manual annotation. Annotating large-scale datasets with pixel-level labels is labor-intensive, time-consuming, and requires substantial human effort.
To address this challenge, an alternative strategy involves leveraging generative models to synthesize datasets with pixel-level labels. Past research efforts have utilized Generative Adversarial Networks (GANs) to effectively generate synthetic datasets for semantic segmentation, thereby mitigating the reliance on manual annotation [1-3]. However, GAN models primarily concentrate on object-centric images and have yet to capture the intricate complexities present in real-world scenes.
On the other hand, text-to-image diffusion models have emerged as a promising technique for generating highly realistic images from textual descriptions [4-7]. These models possess unique characteristics that make them well-suited for the generation of semantic segmentation datasets. Firstly, the text prompts used as input to these models can serve as valuable guidance since they explicitly specify the objects to be generated. Secondly, the application of cross and self-attention maps in the image generation process endows these models with informative spatial cues, enabling precise extraction of object positions within the generated images.
By leveraging these characteristics of text-to-image diffusion models, the concurrent works DiffuMask [8] and DiffusionSeg [9] effectively generate pairs of synthetic images and corresponding
*First two authors contribute equally. The work is done during Quang Nguyen's internship at VinAI Research.
Figure 1: Overview of our Dataset Diffusion for synthetic dataset generation. (Left) Given the target classes, our framework generates high-fidelity images with their corresponding pixel-level semantic segmentations. These segmentations serve as pseudo-labels for training a semantic segmenter. (Right) The trained semantic segmenter is able to predict the semantic segmentation of a test image.
segmentation masks. DiffuMask achieves this by utilizing straightforward text prompts, such as "a photo of a [class name] [background description]", to generate image and segmentation mask pairs. Meanwhile, DiffusionSeg focuses on creating synthetic datasets that address the challenge of object discovery, which involves identifying salient objects within an image. While these approaches successfully produce images paired with their corresponding segmentation masks, they are currently limited to generating a single object segmentation mask per image.
In this paper, we present Dataset Diffusion, a novel framework for synthesizing high-quality semantic segmentation datasets, as shown in Fig. 1. Our approach focuses on generating realistic images depicting scenes with multiple objects, along with precise segmentation masks. We introduce two techniques: class-prompt appending, which encourages diverse object classes in the generated images, and class-prompt cross-attention, enabling more precise attention to each object within the scene. We also introduce self-attention exponentiation, a simple refinement method using self-attention maps to enhance segmentation quality. Finally, we employ the generated data to train a semantic segmenter using uncertainty-aware segmentation loss and self-training.
To evaluate the quality of the synthesized datasets, we introduce two benchmark datasets: synth-VOC and synth-COCO. These benchmarks utilize two well-established semantic segmentation datasets, namely PASCAL VOC [10] and COCO [11], to standardize the text prompt inputs and ground-truth segmentation evaluation. On the synth-VOC benchmark, Dataset Diffusion achieves an impressive mIoU of 64.8, outperforming DiffuMask [8] by a substantial margin. On the synth-COCO benchmark, the DeepLabV3 model trained on our synthesized dataset achieves noteworthy results of 34.2 in mIoU compared to the model trained on real images with full supervision.
In summary, the contributions of our work are as follows:
- We present a framework that effectively employs a state-of-the-art text-to-image diffusion model to generate synthetic datasets with pixel-level annotations.
- We introduce a simple and effective text prompt design that facilitates the generation of complex and realistic images, closely resembling real-world scenes.
- We propose a straightforward method that utilizes self and cross-attention maps to achieve highly accurate segmentation, thereby improving the quality and reliability of the synthesized datasets.
- We introduce synth-VOC and synth-COCO benchmarks for evaluating the performance of semantic segmentation dataset synthesis.
In the following, Sec. 2 reviews prior work, Sec. 3 describes our proposed framework, and Sec. 4 presents our experimental results. Finally, Sec. 5 concludes with some remarks and discussions.
Semantic segmentation is a critical computer vision task that involves classifying each pixel in an image to a specific class label. Popular semantic segmentation approaches include the fully convolutional network (FCN) [12] and its successors, such as DeepLab [13], DeepLabV2 [14], DeepLabv3 [15], DeepLabv3+ [16], UNet [17], SegNet [18], PSPNet [19], and HRNet [20]. Recently,
transformer-based approaches like SETR [21], Segmenter [22], SegFormer [23], and Mask2Former [24] have gained attention for their superior performance over convolution-based approaches. In our framework, we focus on generating synthetic datasets that can be used with any semantic segmenter, so we use DeepLabv3 and Mask2Former as they are commonly used.
Text-to-image diffusion models have revolutionized image generation research, moving beyond simple class-conditioned to more complex text-conditioned image generation. Examples include GLIDE [25], Imagen [6], Stable Diffusion (SD) [5], Dall-E [4], eDiff-I [7], and Muse [26]. These models can generate images with multiple objects interacting with each other, more closely resembling real-world images rather than the single object-centric images generated by prior generative models. Our Dataset Diffusion marks a milestone in synthetic dataset generation literature, moving from image-level annotation to pixel-level annotation. We utilize Stable Diffusion [5] in our framework, as it is the only open-sourced pretrained text-to-image diffusion model available at the time of writing.
Diffusion models for segmentation. Diffusion models have proven effective for semantic, instance, and panoptic segmentation tasks. These models either use input images to condition the maskdenoising process [27-33], or employ pretrained diffusion models as feature extractors [34-37]. However, they still require ground-truth (GT) segmentation for training. In contrast, our framework utilizes only a pretrained SD to generate semantic segmentation without GT labels.
Generative Adversarial Networks (GANs) for synthetic segmentation datasets. GANs have been employed in the generation of synthetic segmentation datasets, as demonstrated in previous works such as [1, 3, 38, 39]. However, these approaches primarily focus on object-centric images, where a single mask is segmented for the salient object or specific parts of common objects like faces, cars, or horses, as exemplified in [2]. In contrast, our framework is designed to generate semantic segmentations for more complex images, where multiple objects interact with each other at the scene level. Furthermore, while some techniques [38, 39] support foreground/background subtraction, and others [1, 3] still require human annotations, our objective is to generate semantic segmentations for multiple object classes in each image without the need for human involvement.
Diffusion models for synthetic data generation have been used to improve the performance of image classification [40, 41], domain adaptation for classification [42, 43], and zero/few-shot learning [44-47]. However, these methods produce only image-level annotations as augmentation datasets. In contrast, our framework produces pixel-level annotations, which is considerably more challenging.
Recently, there have been concurrent works [8, 9] that utilize Stable Diffusion (SD) for generating object segmentation without any annotations. However, they focus on segmenting a single object in an image rather than multiple objects. Their text-prompt inputs to SD are simple, usually “a photo of a [class name]”. The semantic segmenter trained on these annotations can segment multiple objects to some extent. Our framework, on the other hand, employs more complex text prompts where multiple objects can coexist and interact, making it more suitable for the semantic segmentation task in real-world images.
3 Dataset Diffusion
Problem setting: Our objective is to generate a synthetic dataset  $\mathcal{D} = (I_i, S_i)_{i=1}^N$ , consisting of high-fidelity images  $\mathcal{I}$  and pixel-level semantic masks  $\mathcal{S}$ . These images and masks capture both the semantic and location information of the target classes  $\mathcal{C} = \{c_1, c_2, \dots, c_K\}$ , where  $K$  represents the number of classes. The purpose of constructing this dataset is to train a semantic segmenter  $\Phi$  without relying on human annotation.
In our approach, we follow a three-step process. Firstly, we prepare relevant text prompts  $\mathcal{P}$  containing the target classes (Sec. 3.1). Secondly, using Stable Diffusion (SD) as our model, we generate images  $\mathcal{I}_i\in \mathbb{R}^{H\times W\times 3}$  and their corresponding semantic segmentations  $S_{i}\in \{0,\dots ,K\}^{H\times W}$ , where 0 represents the background class (Sec. 3.2). These images and segmentations form the synthetic dataset  $\mathcal{D}$ . Lastly, we train a semantic segmenter  $\Phi$  on  $\mathcal{D}$  and evaluate its performance on the test set of standard semantic segmentation datasets (Sec. 3.3). It is worth noting that our approach primarily focuses on segmenting common objects in everyday scenes, where the SD model excels, rather than specialized domains like medical or aerial images. The overall framework is depicted in Fig. 2.
Figure 2: Three stages of Dataset Diffusion. In the first stage, the target classes are provided, and text prompts are generated using language models such as ChatGPT [48]. Real captions (for COCO) or image-based captions (for VOC) can also be used for prompt generation to ensure standard evaluation. The text prompts are then augmented with the target class labels to avoid missing objects. In the second stage, given the augmented text prompt, a frozen Stable Diffusion [5] is employed to generate an image and its self- and cross-attention maps. The cross-attention map for each target class is refined using the self-attention map to match the object's shape. Finally, the generated images and corresponding semantic segmentations are used to train a semantic segmenter with uncertainty-aware loss and the self-training technique.
3.1 Preparing Text Prompts for Stable Diffusion
To prepare prompts containing a given list of classes for SD, one option is to utilize a large language model (LLM) such as ChatGPT [48] to generate the sentences, similar to the method described in [9]. This approach can be valuable in real-world applications.
However, for evaluating the quality of the synthetic dataset, we need to rely on standard datasets for semantic segmentation like PASCAL VOC [10] or COCO [11] to create standardized benchmarks. In this regard, we propose using the provided or generated captions of the training images in these datasets as the text prompts for SD. This is solely for the purpose of standard benchmarking where the text prompts are fixed, and we do not utilize real images or image-label associations in our synthetic dataset generation. We call these new benchmarks as synth-VOC and synth-COCO.
When using the COCO dataset, we can rely on the provided captions to describe the training images. However, in the case of the PASCAL VOC dataset, which lacks captions, we employ a state-of-the-art image captioner like BLIP [49] to generate captions for each image. However, we encountered several issues with the provided or generated captions. Firstly, the text prompts may not use the exact terms as the target class names  $\mathcal{C}$  provided in the dataset. For instance, terms like "man" and "woman" may be used instead of "person", or "bike" instead of "bicycle", resulting in a mismatch with the target classes. Secondly, many captions do not contain all the classes that are actually present in the images (as illustrated in Fig. 3). This leads to a shortage of text prompts for certain classes, affecting the generation process for those particular classes.
To address the issues, we propose a method that leverages the class labels provided by the datasets. We append the provided (or generated) captions  $\mathcal{P}_i$  with the class labels, creating new text prompts  $\mathcal{P}'_i$  that explicitly incorporate all the target classes  $\mathcal{C}_i = [c_1; \ldots; c_M]$ , where  $M$  is the number of classes in image  $i$ . This is achieved through the text appending operation or class-prompt appending technique:  $\mathcal{P}'_i = [\mathcal{P}_i; \mathcal{C}_i]$ . For example, in the case of the left image in Fig. 3, the final text prompt would be "a photograph of a kitchen inside a house; bottle microwave sink refrigerator". This ensures that the new text prompts encompass all the target classes, addressing the issue of mismatched or missing class names in the captions.
3.2 Generating Segmentation from Self and Cross-attention Maps
We build our segmentation generator on Stable Diffusion (SD) by leveraging its self and cross-attention layers. Given a text prompt  $\mathcal{P}'$  first encoded by a text encoder into text embedding  $e\in \mathbb{R}^{\Lambda \times d_e}$  with the text length  $\Lambda$  and the number of dimensions  $d_{e}$ , SD seeks to output the final latent state  $z_0\in \mathbb{R}^{H\times W\times d_z}$ , where  $H,W,d_{z}$  are height, width, and number of channels of  $z_0$ , reflecting the content encoded in  $e$  from the initial latent state  $z_{T}\sim \mathcal{N}(0,I)$  after  $T$  denoising steps.
Caption: A photograph of a kitchen inside a house.  
Provided classes: bottle, microwave, sink, refrigerator
Caption: A bike leaning against a sign in Scotland.  
Provided classes: bicycle, backpack, bottle
Caption: A man riding a dirt bike in a forest.  
Provided classes: person, motorcycle
Figure 3: Common issues of using provided (or generated) captions. Red classes are often missing from the captions, resulting in a lack of text prompts for those classes. Blue classes may have different terms used in the captions, causing a discrepancy between the target class names and the text prompts.
Generated image
Figure 4: Given a text prompt "A bike is parked in a room; bicycle", we obtain the generated image, cross-attention map, enhanced cross-attention map by the self-attention with  $\tau = \{1,2,4\}$  described in the Eq. (3), and mask with uncertainty value (white region) by Eq. (4) and Eq. (5).
Cross-attention
$\tau = 1$
$\tau = 2$
$\tau = 4$
Mask + Uncertainty
At each denoising step  $t$ , a UNet architecture with  $L$  layers of self and cross-attention is used to transform  $z_{t}$  to  $z_{t - 1}$ . In particular, at layer  $l$  and time step  $t$ , the self-attention layer captures the pairwise similarity between positions within a latent state  $z_{t}^{l}$  in order to enhance the local feature with the global context in  $z_{t}^{l + 1}$ . In the meantime, the cross-attention layer models the relationship between each position of the latent state  $z_{t}^{l}$  and each token of the text embedding  $e$  so that  $z_{t}^{l + 1}$  can express more of the content encoded in  $e$ .
Formally, the self-attention map  $\mathcal{A}_S^{l,t} \in [0,1]^{HW \times HW}$  and cross-attention map  $\mathcal{A}_C^{l,t} \in [0,1]^{HW \times \Lambda}$  at layer  $l$  and time step  $t$  are computed as follows:
$$
\mathcal {A} _ {S} ^ {l, t} = \operatorname {S o f t m a x} \left(\frac {Q _ {z} K _ {z} ^ {\top}}{\sqrt {d _ {l}}}\right), \quad \mathcal {A} _ {C} ^ {l, t} = \operatorname {S o f t m a x} \left(\frac {Q _ {z} K _ {e} ^ {\top}}{\sqrt {d _ {l}}}\right), \tag {1}
$$
where  $Q_{z}, K_{z}, K_{e}$  are the query of  $z$ , key of  $z$ , and key of  $e$ , respectively, obtained by linear projections and taken as inputs to the attention mechanisms, and  $d_{l}$  is # features at layer  $l$ .
Since we only want to obtain the cross-attention map of the class labels  $C_i$  of image  $i$  for semantic segmentation, we introduce class-prompt cross-attention that is similar to cross-attention in Eq. (1) but produced by only taking the softmax over the class name part  $C_i$  rather than entire of the text prompt  $\mathcal{P}'_i$ . In practice, we form a new text prompt  $\hat{\mathcal{P}}_i = C_i$  just for the purpose of extracting the cross-attention maps while the original text prompt  $\mathcal{P}'_i$  for generating images keeps unchanged. After this, we obtain  $\mathcal{A}_C^{l,t} \in [0,1]^{HW \times M}$ , where  $M$  is the number of classes in the image.
With the observation that using different ranges of timesteps only affects the final result marginally, (provided in Supp.), we average these cross and self-attention maps over layers and timesteps:
$$
\mathcal {A} _ {S} = \frac {1}{L \times T} \sum_ {l = 1} ^ {L} \sum_ {t = 0} ^ {T} \mathcal {A} _ {S} ^ {l, t}, \quad \mathcal {A} _ {C} = \frac {1}{L \times T} \sum_ {l = 1} ^ {L} \sum_ {t = 0} ^ {T} \mathcal {A} _ {C} ^ {l, t}, \tag {2}
$$
Although the cross-attention maps  $\mathcal{A}_C$  already exhibit the location of the target classes in the image, they are still coarse-grained and noisy, as illustrated in Fig. 4. Thus, we propose to use the self-attention map  $\mathcal{A}_S$  (as illustrated in Fig. 6 - Left) to enhance  $\mathcal{A}_C$  for a more precise object location. This is because the self-attention maps capturing the pairwise correlations among positions within the latent  $z_{t}$  can help propagate the initial cross-attention maps to the highly similar positions, e.g., non-salient parts of the object, thereby enhancing their quality. Therefore, we propose self-attention exponentiation where the self-attention map  $\mathcal{A}_S$  is powered to  $\tau$  before multiplying to the cross-attention map  $\mathcal{A}_C$  as:
$$
\mathcal {A} _ {C} ^ {*} = \left(\mathcal {A} _ {S}\right) ^ {\tau} \cdot \mathcal {A} _ {C}, \quad \mathcal {A} _ {C} ^ {*} \in [ 0, 1 ] ^ {H W \times M}. \tag {3}
$$
Next, we aim to identify two matrices:  $\mathcal{V} \in [0,1]^{H \times W}$  representing the objectness value at each location (the higher the objectness, the more likely that location contains an object), and  $S \in \{1,\dots,M\}^{H \times W}$  indicating which objects in the class labels  $C_i$  that each location could be. To obtain those, we perform the pixel-wise arg max and max operator (over the category  $M$  dimension):
$$
\mathcal {S} = \underset {m} {\arg \max } \mathcal {A} _ {C} ^ {*; m}, \quad \mathcal {V} = \underset {m} {\max } \mathcal {A} _ {C} ^ {*; m}. \tag {4}
$$
At a location  $x$  in the map  $\mathcal{V}$ , if its value is less than a threshold, one can set its label to the background class 0. However, we find that using a fixed threshold does not work for all images. Instead, we use a lower threshold  $\alpha$  for certain background decisions and a higher threshold  $\beta$  for certain foreground decisions. Any value that falls inside the range  $(\alpha, \beta)$  expresses an uncertain mask prediction with value  $U = 255$ . That is, the final mask  $\bar{S}$  is illustrated in the last image of Fig. 4 and calculated as:
$$
\bar {S} _ {x} = \left\{ \begin{array}{l l} 0 & \text {i f} \mathcal {V} _ {x} \leq \alpha , \\ U & \text {i f} \alpha <   \mathcal {V} _ {x} <   \beta , \\ S _ {x} & \text {o t h e r w i s e .} \end{array} \right. \tag {5}
$$
3.3 Training Semantic Segmenter on Generated Segmentation
Given the synthetic images  $\mathcal{I}$  and semantic segmentation masks  $\bar{S}$ , we train a semantic segmenter  $\Phi$  with an uncertainty-aware cross-entropy loss. Specifically, for pixels marked as uncertain, we ignore the loss from those as:  $\mathcal{L} = \sum_{x} \mathbb{1}(\bar{S}_x \neq U) \mathcal{L}_{\mathrm{CE}}(\hat{S}_x, \bar{S}_x)$ , where  $\mathbb{1}$  is the indication function,  $\mathcal{L}_{\mathrm{CE}}$  is the cross entropy loss, and  $\hat{S} = \Phi(\mathcal{I})$  is the predicted segmentation from the generated image  $\mathcal{I}$ .
We further enhance the segmentation mask  $\bar{S}$  by the self-training technique [50]. That is, after being trained with  $\bar{S}$ , the segmenter  $\Phi$  makes its own prediction on  $\mathcal{L}$  as pseudo labels  $S^{*}$  without uncertainty value  $U$ . Finally, the final semantic segmenter  $\Phi^{*}$  is the segmenter  $\Phi$  trained again on  $S^{*}$ .
Datasets: We evaluate our Dataset Diffusion on two datasets: PASCAL VOC 2012 [10] and COCO 2017 [11]. The PASCAL VOC 2012 dataset has 20 object classes and 1 background class. For standard semantic segmentation evaluation, this dataset is usually augmented with the SBD dataset [51] to have a total of 12,046 training, 1,449 validation, and 1,456 test images. We additionally augment the training images with captions generated from BLIP [49]. The COCO 2017 dataset contains 80 object classes and 1 background class with 118,288 training and  $5K$  validation images, along with provided captions for each image. It is worth noting that we only use the image-level class annotation to form the text prompts as described in Sec. 3.1.
We introduce the set of our prepared text prompts along with the validation set of each dataset as synth-VOC and synth-COCO – the two benchmarks for evaluation of semantic segmentation dataset synthesis. To create a balance synthetic dataset among classes, we generate  $2k$  images per object class for PASCAL VOC, resulting in a total of  $40k$  image-mask pairs and about  $1k$  images per object class for COCO, resulting in a total of  $80k$  image-mask pairs. If the number of text prompts associated with a certain class is insufficient, we use more random seeds to generate more images.
Evaluation metric: We evaluate the performance of Dataset Diffusion using the mean Intersection over Union (mIoU) metric. The mIoU(\%) score measures the overlap between the predicted segmentation masks and the ground truth masks for each class and takes the average across all classes.