SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Method	Train Data	Backbone	Inference(PPP)	Inference(COCO-PSG)
AP	AP50	AP75	APE	AP50	AP75	AREP	AP	AP50	AP75	APE	AP50	AP75	ARRE
Baseline (Sec. 3.1)	PPP [6]	Swin-Tiny	24.3	54.3	18.1	48.6	70.0	49.8	65.8	-	-	-	25.7	53.1	34.1	-
COCO-PSG [8]	-	-	-	44.5	65.6	46.1	-	39.6	59.7	41.3	41.7	67.3	42.4	51.8
All	24.5	55.4	18.2	53.4	75.4	55.8	69.7	38.9	59.0	40.7	40.4	65.0	41.2	50.9
AIMS	26.5	59.0	20.2	56.1	78.5	58.5	72.3	40.5	60.6	42.3	42.1	67.2	43.3	53.1
Swin-Large	29.3	63.0	23.1	60.0	81.6	63.1	75.6	44.2	65.1	45.5	46.8	70.4	48.7	56.9
Table 2: Comparison among single dataset training, baseline, and our proposed method AIMS on multi-dataset training.
each dataset was originally created for a distinct task, we split each dataset into training and evaluation, and subsequently merge them to form unified training and evaluation sets. Overall, our training and evaluation sets have about 236.7K and 2.1K images. Please refer to our supplementary file for details on the split criteria. With the exception of COCO and EntitySeg, which are solely for entity-level prediction with full or partial image mask prompts, PPP dataset is utilized for entity/part prediction with full or partial image mask prompts and the entity-to-part task. COCO-PSG is employed for relation/entity prediction with full or partial image mask prompts and a relation-to-entities task. We keep only the relationship annotations in COCO-PSG that involve physical contact, to be consistent with typical part-whole relationships. PACO dataset is exclusively used for the entity-to-part task.
Training settings. We train our model for 36,000 iterations using a base learning rate of 0.0001 and weights pre-trained on COCO-Entity [3] with the exception of images contained in our validation set. The longer edge size of the images is set to 1,333 pixels, while the shorter edge size is randomly sampled between 640 and 800 pixels, with a stride of 32 pixels. The learning rate is decayed by a factor of 0.1 after 28,000 and 33,000 iterations, respectively. During each training iteration, we sample the data and tasks as introduced in the sampling strategy of Section 3.4, with a batch size of 64 on 8 A100 GPUs.
Federated evaluation. We build our evaluation set by collecting images from PPP [6] and COCO-PSG [8] dataset, where each image only has two level annotations. Similar to LVIS [46], our evaluation images have non-exhaustive annotations. To better approach this, we organize the evaluation set into two subsets: one for entity-part and another for relation-entity. For each subset, we evaluate our model by the instance segmentation metric mean average precision (mAP) for each level and recall AR@100 for cross-level association predictions. For brief illustration,  $\mathrm{AP}^p$ ,  $\mathrm{AP}^e$ , and  $\mathrm{AP}^r$  are mAP for part-, entity- and relation- predictions. The  $\mathrm{AR}^{ep} @ 100$  and  $\mathrm{AR}^{re} @ 100$  are metrics for entity-part and relation-entity association performance.
In the following, we first ablate our experiments with a Swin-Tiny [54] backbone and then report the quantitative performance and visualization results of our best model with the Swin-Large [54] backbone to other state-of-the-art methods.
4.1 Ablation Study
Multi-dataset training. Table 2 demonstrates the overall performance improvement by introducing multiple datasets. In the  $1^{\text{st}}$  and  $2^{\text{nd}}$  rows, we directly apply the base framework (baseline) with single dataset training and separate decoders. From the inference results, we observe that the model trained on the PPP dataset yields lower performance at the entity level when evaluated on the COCO-PSG dataset. This is primarily due to the fact that the COCO-PSG dataset contains 4 times more images than the PPP dataset, which does not provide sufficient entity-level annotations for the model. This issue can be alleviated through multi-dataset training, as depicted in the third row.
When we train our baseline method with all datasets, there are slight improvements or even degradation in part-level prediction for the PPP dataset, and relation and entity-level prediction for the COCO-PSG dataset. This can be attributed to annotation inconsistencies and incompleteness among multiple datasets, which results in ambiguous training signals for the model. By exploiting task complementarity (with Task Complementarity Module) and mitigating annotation issues (with MPE), our proposed AIMS framework can improve performance on both evaluation datasets.
Improvement trajectory of the proposed method. In Table 3(a), we show the component-wise improvements over the baseline. The  $1^{\text{st}}$  and last rows represent the performance of our baseline (three separate decoders for three levels) and the full AIMS framework. The middle three rows display the performance enhancement over the baseline method when integrated with different modules. It is evident that both MPE and TCM contribute to improvements in performance across all three
MPE	TCM	AM	PPP	COCO-PSG
APP	APE	AREP	APR	APE	ARRE
○	○	○	24.5	53.4	69.7	38.9	40.4	50.9
✓	○	○	25.7	54.8	71.0	39.5	41.2	52.0
○	✓	○	25.9	55.0	71.2	39.9	41.6	52.3
○	○	✓	24.3	53.1	70.6	39.0	40.5	52.0
✓	✓	✓	26.5	56.1	72.3	40.5	42.1	53.1
(a)
FUM	PAM	REM	EPM	PPP	COCO-PSG
APP	APE	AREP	APR	APE	ARRE
○	○	○	○	24.5	53.2	69.6	38.9	40.3	51.0
✓	○	○	○	24.5	53.4	69.7	38.9	40.4	50.9
○	✓	○	○	22.1	51.8	65.8	35.8	37.7	47.6
✓	✓	○	○	25.6	54.9	70.4	39.6	41.3	52.1
○	○	✓	✓	20.4	49.7	62.3	32.5	35.6	44.8
✓	✓	✓	✓	26.5	56.1	72.3	40.5	42.1	53.1
(b)
Table 3: Ablation studies. (a) Improvement trajectory of the proposed modules. The abbreviations 'MPE', 'TCM', and 'AM' represent the mask prompt encoder, task complementarity module, and association module, respectively. (b) Ablation study concerning the utilization of mask prompt types. The terms 'FUM', 'PAM', 'REM', and 'EPM' denote the mask types for the full image, partial image with multiple entities, dual entities' mask, and single entity mask, respectively. The checkmark symbol  $(\checkmark)$  and the circle symbol  $(\circ)$  indicate whether the ablated module is used or not.
Structure	PPP	COCO-PSG
APP	APE	AREP	APR	APE	ARRE
None	24.5	53.4	69.7	38.9	40.4	50.9
(CONV2B×2)+(CONV1)	26.5	56.1	72.3	40.5	42.1	53.1
(CONV2B×2)+(CONV1×2)	26.4	56.1	72.3	40.5	42.1	53.1
(CONV2B×2)	25.9	54.7	71.1	39.6	41.6	52.5
(a)
SA	FF	ADD	PPP	COCO-PSG
APP	APE	AREP	APR	APE	ARRE
○	○	○	24.5	53.4	69.7	38.9	40.4	50.9
✓	○	○	25.5	54.8	70.7	39.1	40.9	51.7
✓	✓	○	25.7	54.3	70.8	39.7	41.2	52.3
✓	✓	✓	26.5	56.1	72.3	40.5	42.1	53.1
(b)
Table 4: Ablation studies. (a) The structure design of prompt mask encoder. The 'CONV2B' and 'CONV1' indicates convolution block operation with kernel and stride size 2 followed by layer normalization and GELU and convolution with kernel and stride size 1. The  $\times 2$  means the two sequential structures. (b) The structure design of the task correlation module. The terms 'SA', 'FF', and 'ADD' denote self-attention, feed-forward, and addition respectively.
levels (mAP and association recall) on the two evaluation subsets, highlighting the effectiveness of these proposed modules on all levels. With the addition of the association module, a more significant improvement in association performance can be achieved while maintaining a similar level of segmentation performance compared to the baseline. This indicates that our proposed association module does not negatively impact the segmentation performance. Furthermore, all three modules can mutually benefit each other, giving the best performance as shown in the last row.
Mask Prompt Encoder In Table 3(b), we demonstrate the impact of using different mask prompt types during training. Note that only full-image mask prompts are used during inference. The  $1^{\text{st}}$  and  $2^{\text{nd}}$  rows show that full-image mask prompts do not bring any impact. In the third row, solely using a mask prompt based on the partial annotation area of the image results in a significant performance decline due to the mask prompt inconsistency between the training and inference stages, because inference only uses full-image prompts. However, when both full and partial image prompts are used, as shown in the  $4^{\text{th}}$  row, better performance is achieved by simultaneously mitigating the incomplete annotation problem and maintaining training-inference consistency. As shown by the last row, the introduction of dual-entity or single-entity mask prompts further improves performance, by pushing the model to learn better how to split masks between two immediate adjacent levels (relation-to-entities & entity-to-parts). This is different than (but complements) using full- or partial-image mask prompts which impose large gaps between the source and target levels (e.g., image-to-parts).
In Table 4(a), we examine the influence of various mask encoder architectures The  $1^{\mathrm{st}}$  row represents our baseline without a prompt mask encoder. As demonstrated in the  $2^{\mathrm{nd}}$  and  $4^{\mathrm{th}}$  rows, we find that adding a 'CONV2' and 'CONV1' block with convolution kernel and stride sizes of 2 and 1 is sufficient for the mask encoder architecture. Incorporating more 'CONV1' blocks does not lead to a significant improvement in the final performance, as shown in the third row. This is because a single 'CONV1' layer is enough to encode mask information after the feature that has been already processed by two 'CONV2' blocks, as evident in the last two rows.
Task Complementarity Module. Table 4(b) presents the ablation study on the design of TCM. It is evident that the self-attention block is capable of fusing information from each level to enhance each respective task, as demonstrated in the  $2^{\text{nd}}$  row. Moreover, adding the original level-specific embeddings back to fused embeddings is crucial for improving the final performance. This is because the level-specific embeddings carries unique information specific to each task.
Figure 5: Qualitative comparison of SAM [10] and our method in interactive segmentation. The prompt mask indicates an entity mask that is provided to the network for further separation to the next low level. For the SAM model, we use hyper-parameters for more fine-grained segmentation, which will be introduced in our supplementary material. The images in the first row and the last two rows are from PACO [9] (in-domain) and ImageNet [56] (out-of-domain).
Method	PPP	COCO-PSG
PQ	PartPQ	PWQ	R/mR@20	R/mR@50	R/mR@100
PanopticPartFormer++ [7]	68.0	67.0	68.3	-	-	-
PSGTR [8]	-	-	-	28.2/15.4	32.1/20.3	35.3/21.5
Ours	72.9	70.8	72.3	36.7/22.9	39.4/25.6	42.3/28.2
(a)
Method	Train Data	VG-150 ARR	PACO AREAP	ADE20K AREAP
EntitySeg [5]	33K	-	87.2	82.7
SAM [10]	11M	-	86.2	82.3
Ours	237K	67.4	87.1	83.0
Table 5: Ablation studies. (a) Comparison with state-of-the-art (SOTA) methods in panoptic part segmentation and panoptic scene graph tasks. (b) Performance comparison of segmentation generalization for images in the real world. The  $\mathrm{AR}^R$ ,  $\mathrm{AR}^E$ , and  $\mathrm{AR}^P$  represent recall@100 at relation, entity, and part levels, respectively.
4.2 Comparison to state-of-the-art methods
In the following comparison, we evaluate our model against other state-of-the-art methods. To ensure a fair comparison, all methods employ Swin-Large backbone, with the exception of SAM, which uses a ViT-Huge backbone [55].
Category-aware segmentation. We transfer our AIMS method to other class-aware segmentation tasks, including panoptic part segmentation and panoptic scene graph on PPP and COCO-PSG datasets, respectively. To equip our method with class-aware capability, we directly use our AIMS model as pre-trained weights to fine-tune the corresponding class-aware training dataset. In Table 5(a), we observe that our AIMS model outperforms the corresponding state-of-the-art methods on both tasks. This is because our AIMS model can utilize more training data from various datasets by employing our proposed prompt mask encoder to address annotation inconsistencies among them. Moreover, our AIMS model benefits from the two-stage process of class-agnostic pretraining and class-aware fine-tuning [38], which has been demonstrated to improve the performance of the second stage class-aware fine-tuning when class-agnostic pretraining is performed with more data.
Interactive segmentation. In Table 5(b), we compare AIMS to another popular interactive segmentation method, SAM [10], on the PACO and ADE20K [57] datasets, where ADE20K is for zero-shot evaluation at entity and part levels. Neither SAM nor AIMS uses ADE20K for training, and ADE20K's classes are fairly dissimilar to PascalVOC's [58]. The  $1^{\text{st}}$  row represents our baseline method with purely class-agnostic training without interactive segmentation, which outperforms SAM at the pure entity level even though the SAM model is trained on 11 million images. This indicates that the zero-shot generalization ability is attributed to class-agnostic rather than interactive training. On the other hand, SAM's interactive segmentation can divide masks into more fine-grained segments, which our AIMS is also capable of. On both evaluation datasets, AIMS remarkably outperforms SAM, using merely a 273K training set that is far smaller than SAM's 11M training set. Our proposed MPE and TCM allow our model to leverage more desirable training signals and multi-level task information in a data-efficient manner. Furthermore, our AIMS model can make relation-level predictions, which are higher than part and entity levels and absent in the SAM model.
Figure 5 provides a visual comparison between the SAM [10] model and our model in a real-world scenario. Like AIMS, the SAM model also supports image segmentation at three levels (object,
Figure 6: The illustration of the flexibility of AIMS tasks to tackle the subjective annotation issues in existing datasets.
part, and sub-part level). For an equitable comparison, we use the same entity mask prompt for further delineation. As shown in Figure 5, both models capably segment parts of the entity with distinct approaches. The SAM model leans towards identifying parts based on color and texture clues, whereas our AIMS model relies more on semantic understanding.
The flexibility of our AIMS task. Figure 7 illustrates how our proposed AIMS task provides the flexibility for segmenting anything. This is to address the subjective annotation issues in existing datasets. For instance, in the image (a), the throw pillows and the sofa are predicted as a single entity, which aligns with the ground truth annotation. Nonetheless, in certain scenarios, a user might wish to edit the throw pillows independently. With our AIMS method, utilizing a full image prompt for part-level prediction does not yield these separated masks. However, by selecting this entity mask as an entity prompt for part-level prediction, we can successfully differentiate the three throw pillows.
Additionally, in image (b), a user may want to segment the windows into several components. However, original ground truth annotations typically consider the two windows as a single entity. To meet this requirement, we initially utilize a full image prompt and an entity encoder to identify the mask of the whole windows. Following this, we apply this two-entity mask prompt and relation Decoder to split them into two independent window masks. Ultimately, we can select one window for further segmentation into two parts: the window edge and curtains.
This paper presents AIMS, a unified segmentation model that parses images at three levels: part, entity, and relation. To address the absence of a dataset with annotations for all three levels, we construct a unified dataset by aggregating multiple existing datasets, each providing annotations for one or two levels. Our base method uses three separate decoders for different-level predictions. To handle the annotation inconsistency issue, we propose Mask Prompt Encoder, which offers a more accurate task signal to the model, and Task Complementarity Module to enhance each level's prediction. Extensive experiments demonstrate the effectiveness of our method in both closed-world and zero-shot settings. Our work can serve as a springboard for further research on this new segmentation problem.
This research is based upon work supported by the Office of the Director of National Intelligence (ODNI), Intelligence Advanced Research Projects Activity (IARPA), via IARPA R&D Contract No. 2022-21102100001. The views and conclusions contained herein are those of the authors and should not be interpreted as necessarily representing the official policies or endorsements, either expressed or implied, of the ODNI, IARPA, or the U.S. Government. The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright annotation thereon.
[46] Agrim Gupta, Piotr Dollar, and Ross Girshick. Lvis: A dataset for large vocabulary instance segmentation. In CVPR, 2019. 3, 7
[47] Davood Karimi, Simon K Warfield, and Ali Gholipour. Transfer learning in medical image segmentation: New insights from analysis of the dynamics of model parameters and learned representations. Artificial intelligence in medicine, 2021. 3
[48] Ya-Yen Tsai, Bidan Huang, Yao Guo, and Guang-Zhong Yang. Transfer learning for surgical task segmentation. In ICRA, 2019. 3
[49] David Marr. Vision: A computational investigation into the human representation and processing of visual information. San Francisco: W.H. Freeman. Print., 1982. 3
[50] Leslie G Ungerleider and James V Haxby. 'what'and 'where'in the human brain. Current opinion in neurobiology, 4(2):157-165, 1994. 3
[51] Danielle S Bassett and Michael S Gazzaniga. Understanding complexity in the human brain. Trends in cognitive sciences, 15(5):200-209, 2011. 3
[52] Jing Shi, Wei Xiong, Zhe Lin, and Hyun Joon Jung. Instantbooth: Personalized text-to-image generation without test-time finetuning. In CVPR, 2023. 3
[53] Tsung-Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dólar, and C Lawrence Zitnick. Microsoft coco: Common objects in context. In ECCV, 2014. 5, 6, 14
[54] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In ICCV, 2021. 7
[55] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In arXiv, 2020. 9
[56] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. In NeurIPS, 2012. 9
[57] Bolei Zhou, Hang Zhao, Xavier Puig, Sanja Fidler, Adela Barriuso, and Antonio Torralba. Scene parsing through ade20k dataset. In CVPR, 2017. 9
[58] Mengde Xu, Zheng Zhang, Fangyun Wei, Han Hu, and Xiang Bai. Side adapter network for open-vocabulary semantic segmentation. In arXiv, 2023. 9
[59] Alina Kuznetsova, Hassan Rom, Neil Alldrin, Jasper Uijlings, Ivan Krasin, Jordi Pont-Tuset, Shahab Kamali, Stefan Popov, Matteo Malloci, Alexander Kolesnikov, et al. The open images dataset v4. IJCV, 2020. 16, 20
[60] Christoph Schuhmann, Richard Vencu, Romain Beaumont, Robert Kaczmarczyk, Clayton Mullis, Aarush Katta, Theo Coombes, Jenia Jitsev, and Aran Komatsuzaki. Laion-400m: Open dataset of clip-filtered 400 million image-text pairs. In NeurIPS Workshop, 2021. 16, 20
In this supplementary file, we provide more experimental details and empirical results to further demonstrate the benefits of our proposed All-Inclusive Multi-Level segmentation task (AIMS).
- Experimental details on dataset split, sampling strategy and SAM's [10] configuration.
- More ablation studies on sample strategy and dataset usage.
- More visualization results of our AIMS model with various inference modes.
- User study of AIMS and the concurrent work SAM [10]
The code and models to reproduce our experiments will be released.
Dataset Split As outlined in the main paper, all models are trained using five datasets: COCO [53], EntitySeg [5], PascalVOC Part (PPP)[6], PACO[9], and COCO-PSG [8]. Given that the original training and validation splits of these datasets are tailored for single tasks, we collate the images and reorganize them to suit our AIMS task. Initially, we select 1069 and 1000 validation images from PPP [6] (which covers the part and entity levels) and COCO-PSG [8] (which covers the entity and relation levels) respectively. Following this, we eliminate any duplicate images in the unified training set that are present in the validation images, resulting in a refined training set comprised of approximately 236.7K unique images.
Sampling Strategy In Table 7, we present the eight different sample types, labeled as Sample ID 1 through 8, for each iteration, assuming a batch size of 8. Additionally, images for each of these sample types are uniformly selected from the datasets mentioned in the 'Dataset' column.
SAM's configuration To obtain more fine-grained part-level predictions, we follow the hyperparameters found on SAM's GitHub page: points_per_side (2), pred_iou thresh (0.86), stability_score thresh (0.92), crop_n_layers (1), crop_n_points_downscale_factor (2), min_mask_region_area (100).
B Ablation Studies
Sample Strategy Table 7 presents an ablation study on various sample strategies by examining the impact of each sample type during each iteration. The first row serves as our baseline, adhering to our base framework with a split decoder for varying-level predictions, and it excludes the usage of a prompt mask encoder. The remaining rows illustrate the influence on performance when task prompts are employed during each training iteration.
For instance, the second row reveals that the exclusive usage of a full image task prompt brings a marginal performance improvement over the baseline. Since full-image mask prompts are always the same, injecting them into the original image features is identical to not using any mask prompts. Conversely, the third row shows that employing solely partial-image mask prompts can considerably deteriorate performance, which is inconsistent with the inference process that involves full-image mask prompts.
In the fourth row, we observe that using both full and partial mask prompts can further enhance performance by providing more accurate signals to the network. Additionally, the introduction of mask prompts at the entity-to-part and relation-to-entity levels consistently yields performance improvements, as depicted in subsequent rows.
Dataset Usage Table 8 displays an ablation study focusing on the utilization of various datasets. The first three rows present the model performance when trained on a single dataset. This results in degraded performance on cross-dataset evaluation due to inconsistencies in annotations across different datasets. As shown in the fourth row, combining the three datasets for training enhances the overall validation performance. This improvement is due to the consistency maintained between the training and validation splits. Incorporating PACO and Entityseg datasets at part and entity levels further improve the performance for the respective levels.
Sample ID	Mask Prompt	Decoder	Dataset
1	Entity	COCO, EntitySeg, PPP, COCO-PSG
2	Full Image	Part	PPP
3	Relation	COCO-PSG
4	Entity	COCO, EntitySeg, PPP, COCO-PSG, PACO
5	Partial Image	Part	PPP, PACO
6	Relation	COCO-PSG
7	One Entity	Part	PPP, PACO
8	Two Entities	Relation	COCO-PSG
Table 6: The illustration of eight sample types on each training iteration.
Sample ID	PPP (Inference)	COCO-PSG (Inference)
1	2	3	4	5	6	7	8	\(AP^P\)	\(AP^E\)	\(AR^{EP}\)	\(AP^R\)	\(AP^E\)	\(AR^{RE}\)
○	○	○	○	○	○	○	○	24.5	53.4	69.7	38.9	40.4	50.9
✓	✓	✓	○	○	○	○	○	25.0	53.5	69.8	39.2	40.8	51.3
○	○	○	✓	✓	✓	○	○	22.5	51.4	67.3	37.1	37.6	48.9
✓	✓	✓	✓	✓	✓	○	○	25.7	54.4	70.9	39.6	41.2	51.8
✓	✓	✓	✓	✓	✓	✓	○	26.4	55.9	72.1	39.6	41.3	51.8
✓	✓	✓	✓	✓	✓	○	✓	25.7	54.3	71.0	40.4	42.0	53.0
✓	✓	✓	✓	✓	✓	✓	✓	26.5	56.1	72.3	40.5	42.1	53.1