K-Net: Towards Unified Image Segmentation

Semantic , instance , and panoptic segmentations have been addressed using different and specialized frameworks despite their underlying connections . This paper presents a unified , simple , and effective framework for these essentially similar tasks . The framework , named K-Net , segments both instances and semantic categories consistently by a group of learnable kernels , where each kernel is responsible for generating a mask for either a potential instance or a stuff class . To remedy the difficulties of distinguishing various instances , we propose a kernel update strategy that enables each kernel dynamic and conditional on its meaningful group in the input image . K-Net can be trained in an end-to-end manner with bipartite matching , and its training and inference are naturally NMS-free and box-free . Without bells and whistles , K-Net surpasses all previous published stateof-the-art single-model results of panoptic segmentation on MS COCO test-dev split and semantic segmentation on ADE20K val split with 55.2 % PQ and 54.3 % mIoU , respectively . Its instance segmentation performance is also on par with Cascade Mask R-CNN on MS COCO with 60 % -90 % faster inference speeds . Code and models will be released at https : //github.com/ZwwWayne/K-Net/ . 1 Introduction . Image segmentation aims at finding groups of coherent pixels [ 48 ] . There are different notions in groups , such as semantic categories ( e.g. , car , dog , cat ) or instances ( e.g. , objects that coexist in the same image ) . Based on the different segmentation targets , the tasks are termed differently , i.e. , semantic and instance segmentation , respectively . There are also pioneer attempts [ 19 , 29 , 50 , 62 ] to joint the two segmentation tasks for more comprehensive scene understanding . Grouping pixels according to semantic categories can be formulated as a dense classification problem . As shown in Fig . 1- ( a ) , recent methods directly learn a set of convolutional kernels ( namely semantic kernels in this paper ) of pre-defined categories and use them to classify pixels [ 40 ] or regions [ 22 ] . Such a framework is elegant and straightforward . However , extending this notion to instance segmentation is non-trivial given the varying number of instances across images . Consequently , instance segmentation is tackled by more complicated frameworks with additional steps such as object detection [ 22 ] or embedding generation [ 44 ] . These methods rely on extra components , which must guarantee the accuracy of extra components to a reasonable extent , or demand complex postprocessing such as Non-Maximum Suppression ( NMS ) and pixel grouping . Recent approaches [ 34 , 49 , 55 ] generate kernels from dense feature grids and then select kernels for segmentation to simplify the frameworks . Nonetheless , since they build upon dense grids to enumerate and select kernels , these methods still rely on hand-crafted post-processing to eliminate masks or kernels of duplicated instances . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . In this paper , we make the first attempt to formulate a unified and effective framework to bridge the seemingly different image segmentation tasks ( semantic , instance , and panoptic ) through the notion of kernels . Our method is dubbed as K-Net ( ‘ K ’ stands for kernels ) . It begins with a set of convolutional kernels that are randomly initialized , and learns the kernels in accordance to the segmentation targets at hand , namely , semantic kernels for semantic categories and instance kernels for instance identities ( Fig . 1- ( b ) ) . A simple combination of semantic kernels and instance kernels allows panoptic segmentation naturally ( Fig . 1- ( c ) ) . In the forward pass , the kernels perform convolution on the image features to obtain the corresponding segmentation predictions . The versatility and simplicity of K-Net are made possible through two designs . First , we formulate K-Net so that it dynamically updates the kernels to make them conditional to their activations on the image . Such a content-aware mechanism is crucial to ensure that each kernel , especially an instance kernel , responds accurately to varying objects in an image . Through applying this adaptive kernel update strategy iteratively , K-Net significantly improves the discriminative ability of the kernels and boosts the final segmentation performance . It is noteworthy that this strategy universally applies to kernels for all the segmentation tasks . Second , inspired by recent advances in object detection [ 4 ] , we adopt the bipartite matching strategy [ 47 ] to assign learning targets for each kernel . This training approach is advantageous to conventional training strategies [ 37 , 46 ] as it builds a one-to-one mapping between kernels and instances in an image . It thus resolves the problem of dealing with a varying number of instances in an image . In addition , it is purely mask-driven without involving boxes . Hence , K-Net is naturally NMS-free and box-free , which is appealing to real-time applications . To show the effectiveness of the proposed unified framework on different segmentation tasks , we conduct extensive experiments on COCO dataset [ 38 ] for panoptic and instance segmentation , and ADE20K dataset [ 70 ] for semantic segmentation . Without bells and whistles , K-Net surpasses all previous state-of-the-art single-model results on panoptic ( 54.6 % PQ ) and semantic segmentation benchmarks ( 54.3 % mIoU ) and achieves competitive performance compared to the more expensive Cascade Mask R-CNN [ 3 ] . We further analyze the learned kernels and find that instance kernels incline to specialize on objects at specific locations of similar sizes . 2 Related Work . Semantic Segmentation . Contemporary semantic segmentation approaches typically build upon a fully convolutional network ( FCN ) [ 40 ] and treat the task as a dense classification problem . Based on this framework , many studies focus on enhancing the feature representation through dilated convolution [ 7–9 ] , pyramid pooling [ 57 , 67 ] , context representations [ 64 , 65 ] , and attention mechanisms [ 32 , 63 , 68 ] . Recently , SETR [ 69 ] reformulates the task as a sequence-to-sequence prediction task by using a vision transformer [ 17 ] . Despite the different model architectures , the approaches above share the common notion of making predictions via static semantic kernels . Differently , the proposed K-Net makes the kernels dynamic and conditional on their activations in the image . Instance Segmentation . There are two representative frameworks for instance segmentation – ‘ topdown ’ and ‘ bottom-up ’ approaches . ‘ Top-down ’ approaches [ 2 , 15 , 22 , 27 , 33 ] first detect accurate bounding boxes and generate a mask for each box . Mask R-CNN [ 22 ] simplifies this pipeline by directly adding a FCN [ 40 ] in Faster R-CNN [ 46 ] . Extensions of this framework add a mask scoring branch [ 24 ] or adopt a cascade structure [ 3 , 5 ] . ‘ Bottom-up ’ methods [ 1 , 30 , 43 , 44 ] first perform semantic segmentation then group pixels into different instances . These methods usually require a grouping process , and their performance often appears inferior to ‘ top-down ’ approaches in popular benchmarks [ 14 , 38 ] . Unlike all these works , K-Net performs segmentation and instance separation simultaneously by constraining each kernel to predict one mask at a time for one object . Therefore , K-Net needs neither bounding box detection nor grouping process . It focuses on refining kernels rather than refining bounding boxes , different from previous cascade methods [ 3 , 5 ] . Recent attempts [ 10 , 54 , 55 , 58 ] perform instance segmentation in one stage without involving detection nor embedding generation . These methods apply dense mask prediction using dense sliding windows [ 10 ] or dense grids [ 54 ] . Some studies explore polar [ 58 ] representation , contour [ 45 ] , and explicit shape representation [ 60 ] of instance masks . These methods all rely on NMS to eliminate duplicated instance masks , which hinders end-to-end training . The heuristic process is also unfavorable for real-time applications . Instance kernels in K-Net are trained in an end-to-end manner with bipartite matching and set prediction loss , thus , our methods does not need NMS . Panoptic Segmentation . Panoptic segmentation [ 29 ] combines instance and semantic segmentation to provide a richer understanding of the scene . Different strategies have been proposed to cope with the instance segmentation task . Mainstream frameworks add a semantic segmentation branch [ 28 , 31 , 55 , 59 ] on an instance segmentation framework or adopt different pixel grouping strategies [ 11 , 61 ] based on a semantic segmentation method . Recently , DETR [ 4 ] tries to simplify the framework by transformer [ 51 ] but need to predict boxes around both stuff and things classes in training for assigning learning targets . These methods either need object detection or embedding generation to separate instances , which does not reconcile the instance and semantic segmentation in a unified framework . By contrast , K-Net partitions an image into semantic regions by semantic kernels and object instances by instance kernels through a unified perspective of kernels . Concurrent to K-Net , some recent attempts [ 12 , 35 , 52 ] apply Transformer [ 51 ] for panoptic segmentation . MaskFormer reformulates semantic segmentation as a mask classification task , which is commonly adopted in instance-level segmentation . From an inverse perspective , K-Net tries to simplify instance and panoptic segmentation by letting a kernel to predict the mask of only one instance or a semantic category , which is the essential design in semantic segmentation . In contrast to K-Net that directly uses learned kernels to predict masks and progressively refines the masks and kernels , MaX-DeepLab [ 52 ] and MaskFormer [ 12 ] rely on queries and Transformer [ 51 ] to produce dynamic kernels for the final mask prediction . Dynamic Kernels . Convolution kernels are usually static , i.e. , agnostic to the inputs , and thus have limited representation ability . Previous works [ 16 , 18 , 25 , 26 , 71 ] explore different kinds of dynamic kernels to improve the flexibility and performance of models . Some semantic segmentation methods apply dynamic kernels to improve the model representation with enlarged receptive fields [ 56 ] or multi-scales contexts [ 20 ] . Differently , K-Net uses dynamic kernels to improve the discriminative capability of the segmentation kernels more so than the input features of kernels . Recent studies apply dynamic kernels to generate instance [ 49 , 55 ] or panoptic [ 34 ] segmentation predictions directly . Because these methods generate kernels from dense feature maps , enumerate kernels of each position , and filter out kernels of background regions , they either still rely on NMS [ 49 , 55 ] or need extra kernel fusion [ 34 ] to eliminate kernels or masks of duplicated objects . Instead of generated from dense grids , the kernels in K-Net are a set of learnable parameters updated by their corresponding contents in the image . K-Net does not need to handle duplicated kernels because its kernels learn to focus on different regions of the image in training , constrained by the bipartite matching strategy that builds a one-to-one mapping between the kernels and instances . 3 Methodology . We consider various segmentation tasks through a unified perspective of kernels . The proposed K-Net uses a set of kernels to assign each pixel to either a potential instance or a semantic class ( Sec . 3.1 ) . To enhance the discriminative capability of kernels , we contribute a way to update the static kernels by the contents in their partitioned pixel groups ( Sec . 3.2 ) . We adopt the bipartite matching strategy to train instance kernels in an end-to-end manner ( Sec . 3.3 ) . K-Net can be applied seamlessly to semantic , instance , and panoptic segmentation as described in Sec . 3.4 .

Summary: This paper tries to unify semantic and instance segmentation with a group of learnable kernels that is conditioned on image features. Although the idea of using dynamic kernels (aka dynamic convolutions) has been already explored in instance segmentation [44, 47] and panoptic segmentation [30], this paper proposes to use sparse dynamic kernels instead of using kernels that are generated on dense grid in [30, 44, 37]. This paper further borrows the idea of bipartite matching loss in DETR [4] to make their model NMS free. The experiments show competitive results compared to state-of-the-art methods on multiple segmentation tasks.