SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lu Qi $^{1}$  Jason Kuen $^{2}$  Weidong Guo $^{3*}$  Jiumiang Gu $^{2}$   
Zhe Lin $^{2}$  Bo Du $^{4}$  Yu Xu $^{3}$  Ming-Hsuan Yang $^{1,5}$ $^{1}$ UC Merced  $^{2}$ Adobe Research  $^{3}$ QQ Brower Lab, Tencen  
 $^{4}$ Wuhan University  $^{5}$ Google Research
Despite the progress of image segmentation for accurate visual entity segmentation, completing the diverse requirements of image editing applications for different-level region-of-interest selections remains unsolved. In this paper, we propose a new task, All-Inclusive Multi-Level Segmentation (AIMS), which segments visual regions into three levels: part, entity, and relation (two entities with some semantic relationships). We also build a unified AIMS model through multi-dataset multi-task training to address the two major challenges of annotation inconsistency and task correlation. Specifically, we propose task complementarity, association, and prompt mask encoder for three-level predictions. Extensive experiments demonstrate the effectiveness and generalization capacity of our method compared to other state-of-the-art methods on a single dataset or the concurrent work on segment anything. We will make our code and training model publicly available.
Decomposing an image into semantically meaningful regions is a key goal of image segmentation. With the recent emergence of photorealistic and high-quality image generation technologies [1, 2], image segmentation has gained considerable attention for its ability to select visual entities of interest within images. Combining image segmentation with image generation has enabled a variety of image editing and manipulation applications, such as altering the style of a specific object or eliminating a distracting element from a scene. Compared to bounding box-based object detection, segmentation methods preserve the integrity of pixels outside the edited or manipulated region.
Panoptic segmentation [3], with its instance awareness and dense segmentation capabilities, is a promising solution for various editing and manipulation applications. However, two main limitations remain when using it: (1) it struggles to handle classes encountered 'in the wild' but absent in the training set [4, 5]; and (2) existing methods typically do not explore the more general possibilities of segmentation masks across multiple abstraction levels (e.g., a car's wheel, the whole car, both a person and car combined), with the exception of part-aware panoptic segmentation [6, 7] and panoptic scene graph [8]. These tasks also fail to generalize to classes beyond those in the training set, especially at the part-level task where annotations are sparse.
In this paper, we address the constraints inherent to different-level segmentation tasks by introducing AIMS (All-Inclusive Multi-Level Segmentation). AIMS represents a comprehensive framework designed to segment images across various hierarchical levels (part, entity, relation) using multiple query-driven image decoders. Drawing inspiration from recent advancements in Entity Segmentation [4, 5], which showcased remarkable generalization capabilities on unseen classes, AIMS performs class-agnostic segmentation across different hierarchical levels to enhance generalization.
Image (I)
Part-Level  $(S_{p})$
ID
Association
One-step Inference
Entity-Level  $(S_e)$
ID
Association
Relation-Level  $(S_{\tau})$
Prompt Inference
Relation-to-Entity
Figure 1: The overview of our All-Inclusive Multi-Level Segmentation (AIMS) task.
Entity-to-Part
Thanks to its hierarchical modeling and class-agnostic learning approach, AIMS provides more level-flexible and all-inclusive segmentation masks for users to select from, thereby enabling a wider range of image editing and manipulation applications. In Figure 1, the AIMS task has two modes of inference: one-step inference and prompt inference. In one-step inference, the model segments the image directly into three levels (part, entity, and relation level). Meanwhile, in prompt inference, the model separates the image further based on the provided mask. Notably, one-step inference necessitates an association of the results between any two neighboring levels.
Unlike existing part-aware panoptic or panoptic scene graph methods, AIMS model is not confined to training on a single dataset that has limited hierarchical levels. Instead, we adopt the strategy of multi-dataset training to build a unified AIMS model for all levels. To maintain consistency across multiple hierarchical levels sourced from different datasets, we introduce the Task Complementarity Module and the Association Module to AIMS model. These respectively facilitate feature interactions between different levels and establish explicit relationships between predictions at adjacent levels (e.g., a car's part mask should correspond with the overall car mask, and a car mask should belong to a mask with a car and driver.).
We train our AIMS model on existing segmentation datasets such as Pascal Panoptic Parts [6], COCO-PSG [8], PACO [9], and EntitySeg [5]. However, due to partial or incomplete annotations over the full image, with inconsistencies in unannotated areas across different datasets, sub-optimal supervision signals are generated during training. To address this issue, we propose a novel Mask Prompt Encoder for our AIMS model. This encoder segments an image based on a given mask prompt, which specifies the region slated for segmentation. We train AIMS to accommodate a diverse set of mask prompts, derived from various hierarchical levels and datasets while excluding unannotated regions. This encoder offers two distinct benefits: (1) it mitigates the inconsistency problem inherent in training with partial annotations from multiple datasets; (2) it exhibits strong generalizability to arbitrary mask prompts during inference, including those stemming from unseen classes and in-the-wild visual entities, as shown in Figure 5.
The main contributions of this work are as follows:
- We introduce the AIMS task, which segments images into three distinct levels: part, entity, and relation. AIMS is capable of performing a one-step inference for the entire image and a prompt inference for a specified mask area within the image, which has great potential for image editing and manipulation applications.
- We propose a unified AIMS model for this task, which incorporates multi-dataset and multi-task training. We have meticulously designed several modules, including the task complementary, association, and mask prompt encoder, to maintain the consistency of tasks or data annotations.
- Extensive experiments and in-depth analysis showcase the effectiveness of our proposed AIMS model across various segmentation tasks and settings. We compare the results from our AIMS model with the popular segmentation models (e.g., SAM[10]), demonstrating the fine-grained segmentation capabilities and generalization to both seen and unseen categories.
Figure 2: The outline of our AIMS framework encompasses an image encoder, a mask prompt encoder, and an image decoder. The outputs generated by AIMS consist of five distinct components, including segmentation results across three levels, and two association metrics.
Image Segmentation Significant advancements [11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21] have been made in image segmentation [22, 23, 3, 24, 25, 26] in recent years, with numerous methods and techniques emerging to handle various challenges through supervised learning. Notable works include FCN [11], DeepLab [12], PSPNet [13], Mask R-CNN [14], PANet [15], SOLO [16], CondInst [17], PanopticFPN [18], PanopticFCN [19], Max-Deeplab [20], and Mask2Former [21] for semantic, instance, and panoptic segmentation. Additionally, specialized tasks tailored to specific scenarios, such as part-aware panoptic segmentation [6, 27, 7, 28] and panoptic scene graphs [8, 29], have been developed. However, these tasks typically focus on single target and do not incorporate hierarchical concepts that are more in line with human behavior. This paper aims to address this problem by proposing a unified segmentation model that can organize pixels at three different levels: part, entity, and relation. We also introduce an interactive stage to decompose unseen regions into smaller entities or parts in an open-world setting.
Multi-Dataset Training Multi-dataset training has gained considerable attention in recent years, as it allows models to draw from multiple data sources to improve performance and generalization [30, 31, 32]. This approach has significantly benefited domain adaptation [33, 34, 35], dataset distillation [36, 37], curriculum learning [38, 39, 40, 41], data augmentation [42, 43, 44], federated learning [45, 46], and transfer learning [47, 48] in the context of image segmentation. This paper addresses the annotation inconsistency among multiple datasets for generalized hierarchical segmentation. Specifically, we tackle the issue of incomplete annotations in a single dataset across three levels using a prompt structure.
As shown in Figure 1, given an image  $\mathbf{I} \in \mathbb{R}^{H \times W \times 3}$ , the AIMS task segments an image into three segmentation maps  $\mathbf{S}_T \in \{0,1\}^{N_T \times H \times W}$  at part, entity, and relation levels and two association maps  $\mathbf{A}_{\mathbf{EP}} \in \{0,1\}^{N_e \times N_p}$  (entity and part levels) and  $\mathbf{A}_{\mathbf{RE}} \in \{0,1\}^{N_r \times N_e}$  (relation and entity levels), where  $H$  and  $W$  are the height and width of the image,  $N_T$  indicates  $N_T$  prediction results in class-agnostic binary format and  $T = \{\mathbf{p},\mathbf{e},\mathbf{r}\}$ . The two association maps  $\mathbf{A}_{\mathbf{EP}}$  and  $\mathbf{A}_{\mathbf{RE}}$  are required to build the association relationships between two level prediction results. We chose these three levels since they align with the human brain's hierarchical information processing system [49, 50, 51] and fulfill the majority of requirements for image editing tools [52]. The outputs of AIMS are akin to those of related tasks like part-aware panoptic segmentation and panoptic scene graphs, but the AIMS task incorporates more levels and operates in a class-agnostic manner.
Like the existing Transformer-based segmentation methods [21], our AIMS model employs an image encoder and decoder to extract robust image features and decode them into segmentation masks. However, it differs in that it addresses the challenges of multi-dataset multi-task training. As illustrated in Figure 3, we first establish a task complementarity module and an association module to build the correlation of each task, followed by a mask prompt encoder to tackle the partial annotation problem among different datasets and enable generalized inference.
Next, we introduce our base framework and then elaborate on our three proposed components.
Figure 3: Overview of image decoder of our All-Inclusive Multi-Level Segmentation framework. The terms 'E-P Correlation Module' and 'E-P Association Module' refer to the task correlation and association structure between the entity and part level, respectively. This explanation can also be applied to the 'R-E Correlation Module' and 'R-E Association Module'. The rightward arrows denote the outputs from the respective modules. 'CONCAT' or 'SPLIT' refers to the joining of two tensors or the division of a single tensor into two tensors, respectively, occurring along the dimension of the query number.
3.1 Base Framework
The base framework of AIMS shares the same image encoder  $\Theta_I$  and has three independent decoders for different-level predictions. The image encoder consists of a backbone and pixel decoder, extracting multi-scale image features  $\mathbf{F_s} \in \mathbb{R}^{\frac{H}{2^s} \times \frac{W}{2^s} \times 256}$  and  $s = \{2,3,4,5\}$ . We denote this image-encoding process as:
$$
\mathbf {F} _ {\mathbf {s}} = \Theta_ {I} (\mathbf {I}) \tag {1}
$$
For the three independent decoders  $\Phi$ , each decoder has 9 cascaded attention blocks, where each block includes a masked cross-attention, self-attention, and feedforward network. Each block of decoder iteratively refines query embeddings  $\mathbf{E}_{\mathbf{T}} \in \mathbb{R}^{N \times 256}$  by interacting with multi-scale features  $F$ , where  $\mathbf{E}_{\mathbf{T}}$  is initialized by a group of learnable parameters  $\mathbf{Q}_{\mathbf{T}}$ . This iterative decoding process is denoted as
$$
\mathbf {E} _ {\mathbf {T}} = \Phi_ {\mathbf {T}} \left(\mathbf {Q} _ {\mathbf {T}}, \mathbf {F} _ {\mathbf {s}}\right) = \prod_ {i = 1} ^ {9} \Phi_ {\mathbf {T}} ^ {i} \left(\mathbf {E} _ {\mathbf {T}} ^ {\mathrm {i}}, \mathbf {F} _ {\mathbf {s}}\right) = \prod_ {i = 1} ^ {9} R _ {\mathrm {f f} _ {\mathbf {T}}} ^ {i} \left(R _ {\mathrm {S A t t} _ {\mathbf {T}}} ^ {i} \left(R _ {\mathrm {X A t t} _ {\mathbf {T}}} ^ {i} \left(\mathbf {E} _ {\mathbf {T}} ^ {\mathrm {i}}, \mathbf {F} _ {\mathbf {s}}\right)\right)\right) \tag {2}
$$
where  $\Phi_{\mathbf{T}}$  and  $\Phi_{\mathbf{T}}^{i}$  indicate the decoder of  $\mathbf{T}$  type and its attention blocks.  $i$  indicates the  $i^{\mathrm{th}}$  attention block. For each attention block, it includes cross-attention  $R_{\mathrm{XAtt_T}}^i$ , self-attention  $R_{\mathrm{SAtt_T}}^i$ , and feedforward network  $R_{\mathrm{ff_T}}^i$ . At last,  $\mathbf{E}_{\mathbf{T}}$  is used for part-ness/entity-ness/relation-ness prediction and pixel-level mask prediction using the low-level image feature  $\mathbf{F_2}$  (derived from the image encoder). The prediction process can be formulated as:
$$
\mathbf {U} _ {\mathbf {T}} ^ {\mathrm {c}}, \mathbf {U} _ {\mathbf {T}} ^ {\mathrm {m}} = \operatorname {P r e d H e a d} _ {\mathbf {T}} \left(\mathbf {E} _ {\mathbf {T}}, \mathbf {F} _ {\mathbf {2}}\right) \tag {3}
$$
where  $\mathbf{U}_{\mathbf{T}}^{\mathrm{c}}$  and  $\mathbf{U}_{\mathbf{T}}^{\mathrm{m}}$  denote the part-ness/entity-ness/relation-ness prediction and pixel-level mask outputs.
During training, we employ the losses  $\mathcal{L}_{\mathbf{T}}$  for different level predictions, with respect to the corresponding ground truth  $\mathbf{G}_{\mathbf{T}}$ . The overall training loss for our model is defined as:
$$
\mathcal {L} _ {\text {b a s e}} = \sum_ {\mathbf {T} \in \{\mathbf {p}, \mathbf {e}, \mathbf {r} \}} \mathcal {L} _ {\mathbf {T}} ^ {\mathrm {c e}} \left(\mathbf {U} _ {\mathbf {T}} ^ {\mathrm {e}}, \mathbf {G} _ {\mathbf {T}} ^ {\mathrm {e}}\right) + \mathcal {L} _ {\mathbf {T}} ^ {\mathrm {b c e}} \left(\mathbf {U} _ {\mathbf {T}} ^ {\mathrm {m}}, \mathbf {G} _ {\mathbf {T}} ^ {\mathrm {m}}\right) + \mathcal {L} _ {\mathbf {T}} ^ {\mathrm {d i c e}} \left(\mathbf {U} _ {\mathbf {T}} ^ {\mathrm {m}}, \mathbf {G} _ {\mathbf {T}} ^ {\mathrm {m}}\right) \tag {4}
$$
where  $\mathcal{L}_{\mathbf{T}}^{\mathrm{ce}}$  denote binary cross-entropy loss for part-/entity-/relation-ness prediction. Similarly,  $\mathcal{L}_{\mathbf{T}}^{\mathrm{bce}}$  and  $\mathcal{L}_{\mathbf{T}}^{\mathrm{dice}}$  denote the binary cross-entropy and dice loss for different level mask prediction.
Level	COCO	PPP	COCO-PSG	PACO	EntitySeg
Relation	○	○	✓	○	○
Entity	✓	✓	✓	✓	✓
Part	○	✓	○	✓	○
(a)
Information	COCO	PPP	COCO-PSG	PACO	EntitySeg
Annotation Area	89.1%	91.7%	90.4%	8.4%	99.9%
Category Numbers	133	20	133	75	634
Image Numbers	115k	9.7k	46k	35k	32k
Table 1: The statistics of the datasets we used for training including COCO [53], EntitySeg [5], PascalVOC Part (PPP) [6], PACO [9], and COCO-PSG [8]. The (a) indicates whether the dataset has the annotations for some level. The (b) satisfy the information of each dataset.
3.2 Task Complementarity Module
In our base method, we utilize three independent decoders for predictions at different hierarchy levels, which merely share the image encoder. To better exploit cross-task knowledge, we design the task complementarity module (TCM) that fuses task information within each attention block of the decoders. To provide a clearer explanation, we use the task complementarity module between the entity and part tasks as an example, and it works identically for the relation and entity tasks. As depicted in Figure 3, we concatenate the embeddings of two tasks following the self-attention of the attention block between two level decoders. Subsequently, we feed the concatenated embedding into a self-attention  $G_{\mathrm{SAtt}}^{i}$  and feed-forward network  $G_{\mathrm{SAtt}}^{i}$  and then separate processed embeddings back into their original dimensions:
$$
\mathbf {E} _ {\mathbf {e} _ {\mathbf {s} ^ {\prime}}} ^ {\mathbf {i}}, \mathbf {E} _ {\mathbf {p} _ {\mathbf {s} ^ {\prime}}} ^ {\mathbf {i}} = G _ {\text {s p l i t}} ^ {i} \left(G _ {\text {f f}} ^ {i} \left(G _ {\text {S A t t}} ^ {i} \left(G _ {\text {C o n c a t}} ^ {i} \left(\mathbf {E} _ {\mathbf {e} _ {\mathbf {s}}} ^ {\mathbf {i}}, \mathbf {E} _ {\mathbf {p} _ {\mathbf {s}}} ^ {\mathbf {i}}\right)\right)\right)\right) \tag {5}
$$
where the  $\mathbf{E}_{\mathrm{es}}^{\mathrm{i}}$  and  $\mathbf{E}_{\mathrm{ps}}^{\mathrm{i}}$  denote the embedding processed by self-attention in the attention block. After that, we acquire the final embedding by performing a straightforward addition between the fused and original embeddings:
$$
\mathbf {E} _ {\mathbf {e} _ {\mathrm {s}} ^ {\prime \prime}} ^ {\mathrm {i}} = \mathbf {E} _ {\mathbf {e} _ {\mathrm {s}} ^ {\prime}} ^ {\mathrm {i}} + \mathbf {E} _ {\mathbf {e} _ {\mathrm {s}}} ^ {\mathrm {i}}, \quad \mathbf {E} _ {\mathbf {p} _ {\mathrm {s}} ^ {\prime \prime}} ^ {\mathrm {i}} = \mathbf {E} _ {\mathbf {p} _ {\mathrm {s}} ^ {\prime}} ^ {\mathrm {i}} + \mathbf {E} _ {\mathbf {p} _ {\mathrm {s}}} ^ {\mathrm {i}} \tag {6}
$$
Ultimately, we feed the enhanced embeddings at each level into the original feedforward network of the corresponding decoder:
$$
\mathbf {E} _ {\mathbf {T}} ^ {\mathbf {i} + \mathbf {1}} = R _ {\mathrm {f f} \mathbf {T}} ^ {i} \left(\mathbf {E} _ {\mathbf {T} _ {\mathbf {s} ^ {\prime \prime}}} ^ {\mathbf {i}}\right). \tag {7}
$$
3.3 Association Module
Aside from generating predictions at various levels, the AIMS model also strives to establish the association between low-level and higher-level results. For example, the part masks of a car should be explicitly associated with the entire car's mask. In image editing and manipulation, this can provide a more intuitive mask selection experience. While it is possible to associate the masks across different levels by computing mask intersection, such association knowledge cannot be instilled into the model. Instead, we introduce an association module to AIMS to learn cross-level associations by applying an association loss to the similarity between embeddings. Besides having the association graph as an outcome, this also helps the model be aware of the cross-level association relationships and thus benefiting the individual-level task performance. The association loss  $\mathcal{L}_{\mathrm{asso}}$  is defined as:
$$
\mathcal {L} _ {\mathrm {a s s o}} = \mathcal {L} _ {\mathrm {b c e}} \left(\operatorname {F C} \left(\mathbf {E} _ {\mathbf {e}}\right) \cdot \left(\operatorname {F C} \left(\mathbf {E} _ {\mathbf {p}}\right)\right) ^ {T}, \mathbf {G} _ {\mathbf {e p}}\right) + \mathcal {L} _ {\mathrm {b c e}} \left(\operatorname {F C} \left(\mathbf {E} _ {\mathbf {r}}\right) \cdot \left(\operatorname {F C} \left(\mathbf {E} _ {\mathbf {e}}\right)\right) ^ {T}, \mathbf {G} _ {\mathbf {r e}}\right) \tag {8}
$$
where FC indicates a fully-connected layer and  $\cdot$  denotes dot product. The  $\mathbf{G}_{\mathrm{ep}} \in \{\mathbf{0}, \mathbf{1}\}^{\mathbf{N}_{\mathrm{e}} \times \mathbf{N}_{\mathrm{p}}}$  and  $\mathbf{G}_{\mathrm{re}} \in \{\mathbf{0}, \mathbf{1}\}^{\mathbf{N}_{\mathrm{r}} \times \mathbf{N}_{\mathrm{e}}}$  are the binary association ground truth of entity-part and relation-entity levels. Finally, the overall training loss for AIMS is
$$
\mathcal {L} = \mathcal {L} _ {\text {b a s e}} + \mathcal {L} _ {\text {a s s o}} \tag {9}
$$
3.4 Mask Prompt Encoder
As aforementioned, we train AIMS with multiple existing segmentation datasets that each covers one or more hierarchy levels, whose details are provided in Table 1. Due to the use of predefined classes, most of the datasets suffer from the problem of partial annotations and have unannotated areas on the images. Furthermore, the unannotated areas appear and happen inconsistently among the multiple datasets (e.g.,  $91.6\%$  of PACO is unannotated, while it is  $0.01\%$  for EntitySeg). This leads to undesirable supervision signals that confuse the model because something that is treated as background in one dataset/level may qualify as foreground in another dataset/level.
Figure 4: Illustration of the task prompt types with different combinations between mask prompts and decoders. For each sub-figure, the right part is our target of the network.
We design the Mask Prompt Encoder (MPE), which allows our method to segment images conditioned on specific mask-controlled task prompts. As depicted in Figure 4, our approach integrates four types of mask prompts: full image, single entity, dual entities and multiple entities (where the upper limit is the entire annotated area). By employing diverse mask prompts along with multi-level mask decoders, AIMS can accomplish different segmentation goals. The mask prompts, covering either the full or partial image area, drive the model to segment the entire image or target area at part, entity, or relation levels using the corresponding decoders. This approach enables us to provide the model with a trustworthy supervision signal for the region to be segmented, even when dealing with datasets that lack full image annotations. Moreover, during inference, the mask prompt encoder can be applied to unrestricted mask prompts. To some extent, the model can subdivide the masks of unseen classes and/or in-the-wild visual entities into reasonable sub-masks, even if they are absent in training data.
In particular, the mask encoder accepts a binary mask  $\mathbf{M}_{\mathbf{P}}$ , which designates the area to be segmented. Our mask encoder utilizes a series of cascaded convolution layers to downscale  $\mathbf{M}_{\mathbf{P}}$ , resulting in  $\mathbf{M}_2 \in \mathbb{R}^{\frac{H}{4} \times \frac{W}{4} \times C_m}$ .
$$
\mathbf {M} _ {2} = \operatorname {M a s k E n c o d e r} \left(\mathbf {M} _ {\mathbf {P}}\right) \tag {10}
$$
where the detailed structure of the MaskEncoder consists of three operation blocks including CONV2+LN+GELU, CONV2+LN+GELU and CONV1. The 'CONV2', 'CONV1', 'LN', and 'GELU' represent a convolution layer with stride and a kernel size of 2 or 1, layer normalization, and Gaussian error linear units, respectively. Subsequently, we directly downscale  $\mathbf{M}_2$  by  $2\times$ ,  $4\times$ , and  $8\times$  to obtain  $\mathbf{M}_{\mathbf{s}}$ , where  $s\in 3,4,5$ , using bilinear interpolation.
Ultimately, we acquire multi-scale fusion features  $\mathbf{F}'s$  by summing the encoded image and mask features  $\mathbf{F}s$  and  $\mathbf{M}_s$ , respectively. For AIMS with MPE, we replace  $\mathbf{F}_s$  in Eq. 2 with  $\mathbf{F}_s'$ :
$$
\mathbf {F} _ {s} ^ {\prime} = \mathbf {F} _ {s} + \mathbf {M} _ {s} \tag {11}
$$
Given MPE, we delve deeper into the sampling strategy implemented across various mask prompts and subtasks. To effectively balance contributions from each subtask and prompt type, we adopt a uniform sampling approach capable of considering multiple factors, such as subtask difficulty and the distribution of different prompt types within the training data. Specifically, in each training iteration, we have five subtasks: segmenting the entire annotated area (either the full image or multiple entities) based on three hierarchy levels, subdividing a single entity to parts, and subdividing two merged entities with relationships to two independent entities. Additionally, each subtask has its own data pool for sampling, taking into account the statistical annotation information of each dataset. For more details, please refer to our supplementary file.
We construct our training set by aggregating images from five segmentation datasets, including COCO [53], EntitySeg [5], PascalVOC Part (PPP) [6], PACO [9], and COCO-PSG [8]. Given that