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: This paper tries to combine DETR and Mask Encoding techniques for end-to-end instance segmentation: DETR could directly separate different objects and Mask Encoding could encode the arbitrary instance mask into a low-dimensional and fix representation. To verify the effectiveness of each component, the authors conducted extensive ablations for each component. The final performance is inspiring, espically using Trasnformer backbone.