SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Xiuye Gu
Yin Cui*†
Jonathan Huang*
Abdullah Rashwan*
Xuan Yang*
Xingyi Zhou*
Golnaz Ghiasi
Weicheng Kuo
Huizhong Chen
Liang-Chieh Chen
Observing the close relationship among panoptic, semantic and instance segmentation tasks, we propose to train a universal multi-dataset multi-task segmentation model: DaTaSeg. We use a shared representation (mask proposals with class predictions) for all tasks. To tackle task discrepancy, we adopt different merge operations and post-processing for different tasks. We also leverage weak-supervision, allowing our segmentation model to benefit from cheaper bounding box annotations. To share knowledge across datasets, we use text embeddings from the same semantic embedding space as classifiers and share all network parameters among datasets. We train DaTaSeg on ADE semantic, COCO panoptic, and Objects365 detection datasets. DaTaSeg improves performance on all datasets, especially small-scale datasets, achieving 54.0 mIoU on ADE semantic and 53.5 PQ on COCO panoptic. DaTaSeg also enables weakly-supervised knowledge transfer on ADE panoptic and Objects365 instance segmentation. Experiments show DaTaSeg scales with the number of training datasets and enables open-vocabulary segmentation through direct transfer. In addition, we annotate an Objects365 instance segmentation set of 1,000 images and release it as a public evaluation benchmark on https://laoreja.github.io/dataset/.
Image segmentation is a core computer vision task with wide applications in photo editing, medical imaging, autonomous driving, and beyond. To suit different needs, various forms of segmentation tasks have arisen, the most popular ones being panoptic [29], semantic [23], and instance [19] segmentation. Prior works generally use specific model architectures tailored to each individual task [28, 45, 21, 6]. However, these segmentation tasks are closely related, as they can all be regarded as grouping pixels and assigning a semantic label to each group. In this paper, we address the following question: Can we leverage a diverse collection of segmentation datasets to co-train a single model for all segmentation tasks? A successful solution to this problem would leverage knowledge sharing among datasets, boosting model performance across the board, especially on smaller datasets.
Existing works on unified segmentation models either focus on a single architecture to handle multiple tasks [10, 8, 74, 25], but with separate weights for different datasets; or a single set of weights for
*Equal contribution. Correspondence to: Xiuye Gu <xiuyegu@google.com>.
$^{\dagger}$ Work done while at Google. Now at NVIDIA.
$\ddagger$ Work done while at Google. Now at ByteDance.
multiple datasets [27, 33, 77], but on the same task. By contrast, in this work, we aim to train a single model on multiple datasets for multiple tasks. Our key idea of unifying these segmentation tasks is to use a universal intermediate mask representation: a set of mask proposals (i.e., grouped pixels) with class labels [63, 10]. Different segmentation tasks can be realized by applying different merge and post-processing operations on this unified representation. This allows us to train our network on the same output space for different tasks. Furthermore, using this representation, we can exploit weak bounding box supervision for segmentation, which is far cheaper to collect than mask annotations.
To encourage knowledge sharing and transfer among the multiple segmentation sources, our network architecture shares the same set of model weights across all datasets and tasks. In addition, we utilize text embeddings as the category classifier, which maps class labels from different datasets into a shared semantic embedding space. This design further enhances knowledge sharing among categories with similar meanings in different datasets, e.g., 'windowpane' and 'window-other'. Our approach can be contrasted with an alternative of using dataset-specific model components — we show in our experiments that our simpler, unified approach leads to improved performance and enables open-vocabulary segmentation via simply switching the text embedding classifier.
Putting these techniques together, we propose DaTaSeg as shown in Fig. 1, a universal segmentation model, together with a cotraining recipe for the multi-task and multi-dataset setting. We train DaTaSeg on the ADE20k semantic [75], COCO panoptic [29], and Objects365 detection [56] datasets. We show that DaTaSeg improves performance on all datasets comparing with training separately, and significantly benefits relatively small-scale datasets
Figure 1: Considering the similarities across various segmentation tasks, and recognizing the potential for enhancing segmentation performance by harnessing data from multiple sources, we propose to train a universal segmentation model, DaTaSeg, on multiple datasets to perform multiple segmentation tasks.
(ADE20k semantic), outperforming a model trained only on ADE20k semantic by  $+6.1$  and  $+5.1$  mIoU with ResNet50 [22] and ViTDet-B [39] backbones. The multi-dataset multi-task setting also allows us to seamlessly perform weakly-supervised segmentation by transferring knowledge from other fully supervised source datasets, which we demonstrate on ADE20k panoptic and Objects365 instance segmentation. DaTaSeg also directly transfers to other datasets not seen during training. It outperforms open-vocabulary panoptic segmentation methods on the Cityscapes dataset [11] and performs comparably with open-vocabulary semantic segmentation works on the Pascal Context dataset [49].
To summarize our contributions, we present DaTaSeg, a single universal segmentation model on multiple segmentation tasks and datasets. DaTaSeg leverages knowledge from multiple sources to boost performance on all datasets. It seamlessly enables weakly-supervised segmentation, directly transfers to other datasets and is capable of open-vocabulary segmentation. As an additional contribution, we have annotated a subset of the Objects365 validation set with groundtruth instance masks and release it as an evaluation benchmark for instance segmentation.
Muti-dataset training: Training on multiple datasets has become popular for developing robust computer vision models [64, 78, 33, 47, 61, 67]. For object detection, Wang et al. [64] train an object detector on 11 datasets in different domains and show improved robustness. UniDet [78] trains a unified detector on 4 large-scale detection datasets with an automatic label merging algorithm. DetectionHub [47] employs text embeddings to accommodate different vocabularies from multiple datasets, and features dataset-specific designs. For segmentation, MSeg [33] manually merges the vocabularies of 7 semantic segmentation datasets, and trains a unified model on all of them; however, the manual efforts are expensive and hard to scale. LMSeg [77] dynamically aligns segment queries with category embeddings. UniSeg [27] explores the label relation and conflicts between multiple
datasets in a learned way. Most of these works merge datasets using a unified vocabulary on a single task, while our work focuses on a more challenging problem: we merge datasets with different vocabularies and different tasks.
Unified segmentation model: Panoptic segmentation [29] unifies semantic and instance segmentation. Prior works [28, 66, 69, 7, 52] exploit separate modules for semantic segmentation [45, 5] and instance segmentation [19, 21], followed by another fusion module. Recently, the mask transformer framework proposed by MaX-DeepLab [63] directly predicts masks with class predictions, allowing end-to-end panoptic segmentation. MaskFormer [10] and K-Net [74] adopt a single transformer-based model for different segmentation tasks. Mask2Former [8] improves upon MaskFormer by proposing masked attention, while kMaX-DeepLab [71] develops k-means cross-attention. OneFormer [25] extends Mask2Former with a multi-task train-once design. All these works still train separate weights on different datasets. By contrast, we aim at a single unified model that can perform well across multiple segmentation tasks and datasets.
Weakly-supervised segmentation: Facing the issue of expensive segmentation mask annotations, many works have proposed to learn segmentation masks from cheaper forms of supervision [50, 12, 42]. In particular, box-supervised approaches are most related to our work, including BoxInst [59], Box2Mask [38], Cut-and-Paste [54] and MAL [34] for instance segmentation, Box2Seg [32] and BCM [58] for semantic segmentation, and DiscoBox [35] and SimpleDoesIt [26] for both semantic and instance segmentation. Despite the popularity of box-driven semantic and instance segmentation, fewer attempts have been made at weakly-supervised panoptic segmentation. Li et al. [37] employ a pretrained model to generate pseudo-ground-truth in advance, in addition to weak supervision from bounding boxes and image level labels. Shen et al. [57] add a handcrafted branch on top of semantic and instance segmentation models to generate panoptic proposals. Unlike these methods that require customized components, our approach can realize knowledge sharing among datasets with different forms of annotations in a more systematic and data-centric fashion, and seamlessly enable weakly-supervised segmentation.
3.1 A universal segmentation representation
Segmentation aims to group pixels of the same concept together. We propose to use an intermediate and universal representation, namely mask proposals, for all segmentation tasks. A mask proposal is a binary foreground mask with a class prediction. We use this representation for its versatility: one mask proposal can represent a single instance, any number of instances, a region, or even a part of an object. They can overlap and can be combined to form higher-level segmentation outputs — an instance ("thing") or a region of the same semantic class ("stuff"). Thus they are well suited for the different flavors of segmentation considered in this paper. We note that a similar concept has been used in existing works [63, 10, 8, 71], for specific segmentation tasks on a single dataset. We show this representation is especially beneficial in our multi-dataset setting, as different datasets define "thing" and "stuff" categories differently. For example, 'table' is a thing category in ADE20k panoptic dataset, but there is a 'table-merged' stuff category in COCO panoptic. In our framework, both "thing" and "stuff" categories are represented using this same representation, and we treat them differently in the next step.
3.2 Merging predictions for specific tasks
Notations: We introduce the notations we use in the equations.  $\widehat{(\cdot)}$  denotes predictions. Let  $(\hat{z}_j,\widehat{M_j})$  denote the  $j$ -th mask proposal  $\widehat{M}_j$  with its class prediction  $\hat{z}_j$ , where  $\widehat{M_j}\in \mathbb{R}^{H\times W}$  is the (presigmoid) logits for the mask proposal and  $\hat{z}_{j,c_k}\in \mathbb{R}$  is the logit for class  $c_{k}$ . We assume a fixed set of  $N$  mask proposals.  $\mathbb{C}_{d,thing}$  and  $\mathbb{C}_{d,stuff}$  are the set of thing and stuff categories in dataset  $d$ .
Merge operation (MERGE): Given the varying formats of groundtruth annotations for different segmentation tasks, we propose a MERGE operation to merge the mask proposals predicting the same category  $c_k$ . We adopt a simple element-wise max operation to merge the mask proposal logits. In particular, the merged mask proposal logits at the  $(p,q)$  location for class  $c_k$  is computed as:
$$
\mathcal {M} \left(c _ {k}\right) _ {p, q} = \max  _ {j} \left(\mathbb {1} [ \operatorname {a r g m a x} \left(\hat {z} _ {j}\right) = c _ {k} ] \cdot \widehat {M} _ {j, p, q}\right). \tag {1}
$$
Figure 2: The universal representation, MERGE operations, and losses for different segmentation tasks. We use a universal representation for all tasks: a set of mask proposals with class predictions. Then we adopt distinct MERGE operations based on the segmentation task. For panoptic segmentation, we merge proposals predicting the same "stuff" category. For semantic segmentation, both "thing" and "stuff" categories undergo merging. In instance segmentation, we do not perform MERGE and there are no "stuff" categories. During training, we use Hungarian matching, and apply different losses based on the supervision types.
We choose element-wise max so that applying MERGE on raw mask logits is equivalent to applying MERGE on the soft mask predictions post-sigmoid. We also choose to merge at the level of raw logits for numerical stability. If no proposal predicts class  $c_k$ , then the corresponding merged mask  $\mathcal{M}(c_k)_{p,q}$  is set to -100 to ensure it is close to 0 after sigmoid.
In addition to merging masks, we merge their corresponding class predictions by simply averaging the class prediction logits of the proposals predicting class  $c_k$ :
$$
\mathcal {Z} \left(c _ {k}\right) = \frac {\sum_ {j} \left(\mathbb {1} \left[ \operatorname {a r g m a x} \left(\hat {z} _ {j}\right) = c _ {k} \right] \cdot \hat {z} _ {j}\right)}{\sum_ {j} \left(\mathbb {1} \left[ \operatorname {a r g m a x} \left(\hat {z} _ {j}\right) = c _ {k} \right]\right) + \epsilon}, \tag {2}
$$
where  $\epsilon$  is a small number to prevent division by zero.
How we apply the above merge operations depends on the given task, as we now describe:
Panoptic segmentation: To make the predicted mask proposals have the same format as the groundtruth, we apply MERGE to all "stuff" categories  $c_k \in \mathbb{C}_{d,stoff}$ .
Semantic segmentation: In contrast with panoptic segmentation, the semantic segmentation groundtruth for each "thing" category is also a single mask (i.e., "thing" categories are treated equivalently as "stuff" categories). We thus apply MERGE to all predicted mask proposals to cover all thing and stuff categories  $\mathbb{C}_{d,\text{stuff}} \cup \mathbb{C}_{d,\text{thing}}$ .
Instance segmentation: MERGE is not needed in this task, since the desired outputs are separate masks for separate instances. There are no "stuff" categories in the vocabulary  $(\mathbb{C}_{d,stuff} = \emptyset)$ , and therefore no stuff proposals.
Training: During training, we apply the corresponding MERGE operations based on the segmentation task. We then use one-to-one Hungarian matching [31] to match the merged outputs with the groundtruth, and calculate the training losses. Detailed training losses are provided in Sec. 3.5.
Inference: At inference time, we apply the same MERGE operations to predicted mask proposals based on the given task. We post-process the mask proposals to obtain non-overlapping outputs for panoptic and semantic segmentation. For instance segmentation, we simply take the top mask proposals as final outputs, since overlaps are allowed. Please refer to supplementary for more details.
We illustrate how the MERGE operation works for different segmentation tasks in Fig. 2.
3.3 Weakly-supervised instance segmentation
Bounding boxes are much cheaper to annotate than instance masks, and thus detection datasets can have a much larger scale than instance segmentation datasets. In order to train on larger detection datasets with bounding box annotations only, as well as to demonstrate the versatility of our framework to handle various tasks, we propose to perform weak-bounding-box-supervised instance segmentation.
Figure 3: Overview of our universal multi-dataset multi-task segmentation model (DaTaSeg). We feed  $N$  learnable segment queries  $\mathcal{Q}_s$  into a transformer decoder, which cross-attends to multi-scale image features  $\mathcal{I}_{multi}$  from the pixel decoder. The outputs serve as the universal representation for all tasks:  $N$  mask proposals  $\widehat{M}_j$  accompanied by  $N$  class predictions  $\hat{z}_j$ . To promote knowledge sharing, we share all network parameters across different datasets and tasks. The only dataset-specific design is a classifier consisting of frozen text embeddings of categories specific to each dataset. Additionally, we employ a shared learnable background classifier.
We adopt a simple projection loss  $\mathcal{L}_{proj}$  from [60], which measures consistency of vertical and horizontal projections of the predicted masks  $\widehat{M}_j$  against groundtruth boxes  $b_{j}$ :
$$
\mathcal {L} _ {p r o j} (\widehat {M _ {j}}, b _ {j}) = \mathcal {L} _ {d i c e} \left(\operatorname {P r o j} _ {x} (S (\widehat {M _ {j}})), \operatorname {P r o j} _ {x} (b _ {j})\right) + \mathcal {L} _ {d i c e} \left(\operatorname {P r o j} _ {y} (S (\widehat {M _ {j}})), \operatorname {P r o j} _ {y} (b _ {j})\right), \tag {3}
$$
where  $\mathcal{L}_{\text{dice}}$  is the dice loss [48];  $b_{j}$  is the groundtruth bounding box matched to  $j$ -th mask proposal  $\widehat{M_j}$ ;  $\operatorname{Proj}_{x/y}$  denotes the projection operation along the  $x$  or  $y$  axis, which can be implemented by a max operation along the axis; and  $S(\cdot)$  denotes the sigmoid function.
By itself, a box consistency loss such as Eqn. 3 is insufficient as a supervision signal for segmentation (e.g., Eqn. 3 is equally satisfied by predicting the bounding box of an object instead of its mask). Thus other works have resorted to additional, often more complex, loss terms (such as the pairwise affinity loss from [60]). However, by training on multiple datasets and multiple segmentation tasks, these handcrafted losses are not necessary, as our model can transfer knowledge gained from other fully-supervised segmentation tasks on other datasets.
3.4 Network architecture with knowledge sharing
Network architecture: We now describe the network architecture (similar to [8]) that predicts mask proposal and class prediction pairs  $(\hat{z}_j, \widehat{M}_j)$ . The input image first goes through a backbone, and we use a pixel decoder to 1) output multi-scale image features  $\mathcal{I}_{multi}$ , and 2) output a high-resolution feature map  $\mathcal{F}$  that fuses information from the multi-scale image features.  $N$  mask proposal pairs are then generated from  $N$  learnable segment queries  $\mathcal{Q}_s$ : The segment queries are fed into a transformer decoder, which cross attends to the multi-scale image features  $\mathcal{I}_{multi}$  from the pixel decoder. The  $N$  decoder outputs are then passed to class embedding and mask embedding heads (both MLPs) to obtain  $N$  (mask embedding, class embedding) pairs. To extract final mask proposals and class predictions, we apply a dot product between the high-resolution feature map  $\mathcal{F}$  and the mask embeddings to obtain  $N$  mask proposal logits  $\widehat{M}_j$ . And to obtain class predictions, we compute dot products of the class embeddings against a per-category classifier  $w_k$  for each category  $c_k$ . Fig. 3 shows an overview.
Knowledge sharing: We bake knowledge sharing into our model in two ways. First, all network parameters are shared across all datasets, allowing knowledge sharing among different datasets and knowledge transfer among different tasks. Additionally, to share knowledge gathered from training samples of similar categories across datasets, e.g., 'windowpane' in ADE20k and 'window-other' in COCO, we propose to map all category names into a shared semantic embedding space. We feed the category names into a frozen pre-trained text encoder to set the per-category classifier  $w_{k}$ . i.e., we let the fixed text embeddings serve as classifiers. This is similar to open-vocabulary
segmentation [16, 70, 72], but with a different purpose: our focus here is on knowledge sharing among different datasets. This automatic approach scales better than manually unifying label spaces in different datasets [33].
3.5 Co-training strategy
We adopt a simple co-training strategy: at each iteration, we randomly sample one dataset, then sample the entire batch for that iteration from the selected dataset. This can be contrasted with sampling from multiple datasets in one iteration. The main advantage for our strategy is that it is simple to implement, and allows more freedom to use different settings for different datasets, including applying distinct losses to various tasks. To account for various dataset scales, we use per-dataset sampling ratios, which randomly samples each dataset at a dataset-specific ratio during training, similar to [41].
For fully-supervised tasks, we employ direct mask supervision, which can be obtained from panoptic/semantic/instance segmentation groundtruth. To calculate the Hungarian matching cost and the training loss, we use a combination of focal binary cross-entropy loss  $\mathcal{L}_{focal}$  [44], and dice loss  $\mathcal{L}_{dice}$  [48], following [10]. Regarding weak bounding box supervision, we adopt the projection loss  $\mathcal{L}_{proj}$  introduced in Sec. 3.3 for the matching cost and training loss. In both cases of supervision, we use the negative of the class prediction probability  $p$  as the classification matching cost [4], and use cross-entropy loss  $\mathcal{L}_{ce}$  for classification training. The total training loss is:
$$
\mathcal {L} _ {d} = \lambda_ {c e, d} \mathcal {L} _ {c e} + \lambda_ {f o c a l, d} \mathcal {L} _ {f o c a l} + \lambda_ {d i c e, d} \mathcal {L} _ {d i c e} + \lambda_ {p r o j, d} \mathcal {L} _ {p r o j}. \tag {4}
$$
The Hungarian matching cost is defined similarly:
$$
\mathcal {C} _ {d} = - \mu_ {c e, d} \cdot p + \mu_ {f o c a l, d} \mathcal {L} _ {f o c a l} + \lambda_ {d i c e, d} \mathcal {L} _ {d i c e} + \mu_ {p r o j, d} \mathcal {L} _ {p r o j}. \tag {5}
$$
Here,  $\lambda_{*,d}$  and  $\mu_{*,d}$  are the weights for dataset  $d$ .
Datasets and metrics: We train and evaluate DaTaSeg on COCO panoptic [29] and ADE20k semantic [75] using mask supervision, as well as Objects365-v2 [56] detection datasets using bounding box weak supervision. COCO panoptic is the most popular panoptic segmentation benchmark with 118,287 training images and 5,000 validation images. COCO has 80 thing categories and 53 stuff categories. ADE20k semantic is one of the most widely used semantic segmentation benchmarks with 150 categories, 20,210 training images, and 2,000 validation images.
For evaluation, besides the training datasets, we also evaluate on ADE20k panoptic, which uses the same validation images as ADE20k semantic but with panoptic annotations. The original 150 categories are divided into 100 thing categories and 50 stuff categories. Finally, we train on the Objects365-v2 detection (bounding boxes only) dataset, which has 365 categories and 1,662,292 training images. To evaluate the weakly-supervised instance segmentation results, we manually label a subset of 1,000 images from the Objects365 validation set. Our annotation pipeline involves 20 well-trained human annotators, following the protocol in [2], without using any automatic assistant tools. Please refer to Appendix B for more details. We release this Objects365 instance segmentation evaluation set as a public benchmark.
To compare DaTaSeg with the state-of-the-art methods, we in addition evaluate on popular open-vocabulary segmentation benchmarks. Following OpenSeg [16], we evaluate on PASCAL Context datasets [49] with 5k val images. We use its two versions: 59 classes (PC-59) and 459 classes (PC-459). We also evaluate on the 500-image val set of Cityscapes panoptic dataset [11], which focuses on urban street scenes with 11 stuff categories and 8 thing categories.
We report results in standard evaluation metrics: panoptic quality (PQ), mean intersection-overunion (mIoU), and mean average precision (AP) for panoptic, semantic, and instance segmentation, respectively.
Implementation details: We experiment with ResNet [22] and ViTDet [39] backbones. For ResNet, we use an ImageNet-1K [55] pretrained checkpoint, with a backbone learning rate multiplier of 0.1 [4]; the pixel decoder consists of a transformer encoder [62] and a modified FPN [43] following [10],