SOLQ: Segmenting Objects by Learning Queries

In this paper , we propose an end-to-end framework for instance segmentation . Based on the recently introduced DETR [ 1 ] , our method , termed SOLQ , segments objects by learning unified queries . In SOLQ , each query represents one object and has multiple representations : class , location and mask . The object queries learned perform classification , box regression and mask encoding simultaneously in an unified vector form . During training phase , the mask vectors encoded are supervised by the compression coding of raw spatial masks . In inference time , mask vectors produced can be directly transformed to spatial masks by the inverse process of compression coding . Experimental results show that SOLQ can achieve state-of-the-art performance , surpassing most of existing approaches . Moreover , the joint learning of unified query representation can greatly improve the detection performance of DETR . We hope our SOLQ can serve as a strong baseline for the Transformer-based instance segmentation . Code is available at https : //github . com/megvii-research/SOLQ . 1 Introduction . Instance segmentation , serving as one of visual detection tasks , not only locates instances of different categories but also generates pixel-level mask for each instance . State-of-the-art instance segmentation methods [ 2 , 3 , 4 , 5 ] follow the two-stage paradigm , which first performs object detection and then segments the masks within detected boxes by RoIAlign [ 2 ] . Those methods are relatively easy to be optimized thanks to the deployment of mature object detectors [ 6 , 7 ] . However , the segmentation branch heavily relies on the detection branch , making it hard to achieve better joint learning of multiple tasks . Some recent works [ 8 , 9 , 10 , 11 ] build instance segmentation frameworks on top of anchorfree object detectors [ 12 , 13 ] to remove the ROI-Cropping operation , reducing the effect of feature misalignment . For example , CondInst [ 9 ] based on FCOS [ 12 ] employs dynamic convolutions [ 14 ] to perform instance segmentation . YOLACT [ 8 ] models the instance segmentation as a combination of prototypes weighted by learned mask coefficients for each anchor . However , the weights of dynamic convolutions or the mask coefficients are still generated by instance proposals . Regardless of the bounding boxes , SOLO [ 15 ] introduces the notion of `` instance categories '' and segments objects by locations . SOLOv2 [ 16 ] further solves the problem of inefficient mask representation by introducing dynamic convolution and so-called matrix Non-Maximum Suppression ( NMS ) . Though SOLO directly outputs instance masks based on locations , hand-crafted post-processes , like NMS , are still required to remove duplicated predictions . Also , the performance of small objects is far from satisfactory due to the imbalance of location samples between the large and small objects . Building an end-to-end instance segmentation framework is still a remaining problem . Our work is inspired by DETR [ 1 ] , which first proposes the end-to-end solution for object detection . In DETR , object detection is regarded as a set prediction problem and objects are represented by ∗Equal contribution . This work is supported by The National Key Research and Development Program of China ( No.2017YFA0700800 ) and Beijing Academy of Artificial Intelligence ( BAAI ) . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . learnable query embeddings . How to encode the spatial binary mask into such an end-to-end system is an opening question . As shown in Fig . 1 ( a ) , DETR is further extended to panoptic segmentation by directly reshaping the learnable embeddings into spatial domain and building a FPN-style [ 17 ] network to produce the final mask predictions . However , both the Transformer encoder and decoder fail to model the spatial information well . Therefore , it is inappropriate to generate the spatial mask based on such query embeddings . Besides , the spatial mask labels used for supervision are of large resolutions , which results in high computation cost and makes it separated to learn the detection and mask branches2 . So we need to find one mask representation that satisfies the following conditions : 1 ) The representation can naturally convert object mask from spatial domain to embedding domain ; 2 ) The process that encodes the spatial masks into embeddings should be reversible ; 3 ) Mask embeddings encoded can keep the principle components of spatial mask . We turn to methods in literature for help and surprisingly find that classical compression coding methods ( e.g . Sparse Coding [ 18 ] ) just satisfy these conditions mentioned above . The mask embeddings can be simply generated from the learnable queries . In training phase , ground-truth spatial mask of each instance can be projected into low-dimensional mask embedding by compression coding and the mask embeddings are used to supervise the learning of predicted mask embeddings . In inference phase , binary spatial mask can be reconstructed from predicted mask embedding by the inverse process of compression coding . With these analysis , we explore how to better encode the spatial mask into the end-to-end object detectors in this paper . Based on DETR , our proposed method , termed SOLQ , segments objects by learning queries . In SOLQ , we formulate the instance segmentation as the joint learning of unified query representation ( UQR ) . The UQR learned can be used to perform parallel predictions for three sub-tasks ( classification , localization and segmentation ) simultaneously and all predictions are obtained in a regression manner ( see Fig . 1 ( b ) ) . In this way , SOLQ directly outputs instance masks together with corresponding class confidences and box coordinates . The learning of UQR can be divided into two parts : generating instance-aware query embeddings and joint supervision of multi-task learning . Specifically , all candidate instances are initialized with several learnable queries , which interact with extracted image features in Transformer decoder to produce the instance-aware query embeddings . The instance-aware query embeddings are further input to three branches of sub-tasks , which contains several linear projection layers , to generate three sub-task vectors . For the classification and regression branches , we follow the same supervisions as in DETR [ 1 ] . For the mask branch , we conduct implicit supervision with the help of mask compression coding mentioned above . To summarize , our contributions are : • We propose an end-to-end framework for instance segmentation based on DETR . SOLQ formulates the instance segmentation as the joint learning of UQR . In UQR , the mask representation can be converted from spatial domain into embedding domain , which is consistent with the learnable query embeddings in DETR . • Experiments show that SOLQ with ResNet101 achieves 40.9 % mask AP and 48.7 % box AP on the challenging MS COCO dataset [ 19 ] without bells and whistles , outperforming SOLOv2 by 1.2 % mask AP and 6.1 % box AP . It is worthy noting that SOLQ can improve 2.0 % in box AP compared to DETR thanks to the joint learning of UQR . 2 Related Work . Instance Segmentation Instance segmentation is a classic but challenging computer vision task . It is required to output each object instance in image with instance-level category label and localization , along with pixel-level mask simultaneously . Currently , there are mainly three categories of instance segmentation methods : top-down , bottom-up and directly-predict methods . Top-down approaches [ 20 , 2 , 8 , 21 , 22 , 23 , 3 , 5 , 24 ] follow the detect-then-segment pipeline . They first generate bounding boxes by object detectors and segment the masks by ROIAlign [ 2 ] or dynamic convolutions [ 9 ] . Bottom-up methods [ 25 , 26 , 27 , 28 ] learn per-pixel embeddings via semantic label and then cluster them into instance groups . For directly-predict methods , PolarMask [ 29 ] employs polar coordinates to represent mask contours . Latest SOLO [ 15 ] and SOLOv2 [ 16 ] directly segment the objects by locations without dependence on bounding boxes or embedding learning . QueryInst [ 24 ] and ISTR [ 30 ] extend the Sparse RCNN [ 31 ] to perform end-to-end instance segmentation . In this paper , we 2Train the object detector first and then freeze the weights of object detector to train segmentation branch . explore an end-to-end instance segmentation solution by learning an unified query representation without any post-processing procedures , like Non-Maximal Suppression ( NMS ) . Transformer in Vision Transformer [ 32 ] introduces the self-attention mechanism to model longrange dependencies , and has been widely applied in natural language processing ( NLP ) . Recently , several works attempted to involve the Transformer architecture into various computer vision tasks and showed promising performances . The non-local block [ 33 ] is first proposed to enhance video recognition by aggregating spatial information . After that , CCNet [ 34 ] further extends the selfattention via sparse attention in semantic segmentation . DETR [ 1 ] and Deformable DETR [ 35 ] adopt learnable queries and Transformer architecture together with bipartite matching to perform object detection in end-to-end fashion , without any hand-crafted process such as NMS . IPT [ 36 ] proposes a transformer-based pretrained network for low-level image processing . ViT series [ 37 , 38 , 39 , 40 ] take an image as a sequence of patches and achieve the cross-patch interactions by Transformer architecture in image classification . Compression Coding in Vision Consider the advantage of low-dimension representation and less computation cost , some recent works have attempted to introduce compression coding methods , like Sparse Coding [ 41 ] , Principal Component Analysis ( PCA ) [ 42 ] and Discrete Cosine Transform ( DCT ) [ 43 ] , into computer vision field . For image classification , [ 44 ] takes the DCT coefficients obtained from RGB images as the inputs of convolutional neural networks ( CNNs ) to reduce the communication bandwidth between CPU and GPU . DCT-Mask [ 45 ] , based on Mask R-CNN , employs DCT supervision to produce high-quality mask representation . Analogously , [ 46 ] performed semantic segmentation on the DCT representation and fed the rearranged DCT coefficients to CNNs . MEInst [ 21 ] and ISTR [ 30 ] encode binary masks into fixed-dimensional mask vectors produced by PCA . 3 Method . 3.1 Reviewing DETR . Recently , DETR [ 1 ] succeed in object detection . It formulates object detection as a set prediction problem and introduces object queries , a set of learnable embeddings , to represent objects . In DETR , each object query predicts an object for a given input image and set prediction loss is adopted to achieve one-to-one matching between the predicted and ground-truth objects in training phase . Further , the transformer encoder-decoder architecture is employed to model the relation between query embeddings and instances for better one-to-one set prediction . Set Prediction Loss Let y = ( c , b ) and ŷ = ( ĉ , b̂ ) denote the ground truths y and the set of predictions ŷ , respectively . c , ĉ ∈ RJ×S are the corresponding class labels and predicted class scores , where J and S are the object number and class number . b , b̂ ∈ RJ×4 are the corresponding ground-truth and predicted box coordinates . ω ∈ ΩJ is the assignment between the ground truths and predictions . Then the optimal one-to-one assignment ω∗ can be calculated by bipartite matching [ 47 ] as : ω∗ = argmin ω∈ΩJ Ldet ( y , ω ( ŷ ) ) ( 1 ) where the bipartite matching loss for object detection Ldet can be summarised as : Ldet ( y , ŷ ) = λcls · Lcls ( c , ĉ ) + λL1 · LL1 ( b , b̂ ) + λgiou · Lgiou ( b , b̂ ) ( 2 ) Here Lcls denotes the focal loss [ 48 ] for classifications , LL1 and Lgiou are L1 loss and generalized IoU loss [ 49 ] for box coordinates , respectively . λcls , λL1 and λgiou are corresponding coefficients . Extension on Segmentation As shown in Fig . 1 ( a ) , DETR is further generalized to panoptic segmentation task by adding a multi-head attention ( MHA ) and FPN-style CNN after the Transformer decoder . Features from the encoder and learned query embeddings from the decoder are reshaped to spatial domain and then interact in MHA . The produced features are then gradually upsampled to the image size by the FPN-Style CNN to obtain spatial masks . It ’ s easy to adapt this framework to perform instance segmentation by cropping instance masks within detected bounding boxes . Object Representation In DETR , objects are represented as a set of object queries . Object queries are initialized by the learnable query embeddings and then interact with the image features in the transformer decoder to update their representation . Finally , object query , as an unified representation , is directly used to classify and localize objects . However , the representation of spatial mask is built by compulsorily reshaping the query embeddings to spatial domain in instance segmentation task . Such design leads to different representation forms compared to detection branch . Besides , the two-stage training process makes DETR fail to enjoy the benefit from multi-task learning .

Summary: The paper proposes and end-to-end instance segmentation approach based on recently proposed Transformer-based object detector DETR. In DETR instance segmentation is done by upsampling of ResNet features from multiple stages, which is computationally expensive. Instead, the authors propose to predict instance masks together with bounding box and class labels, by compressing the masks. This results in a simple and efficient detection and instance segmentation model. Moreover, AP on both tasks improves significantly over DETR and Deformable DETR, achieving very impressive results on par with state-of-the-art, on the challenging COCO dataset. The paper also provides a detailed ablation study of the proposed model.