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: The paper presents a convolutional neural network method for panoptic (multi-class, multi-instance) image segmentation. Semantic segmentation is addressed in the standard fashion with pre-trained model weights (convolutional kernels). Novelty is introduced in the instance segmentation component where the parameters of the convolutional kernels responsible for producing activation maps, one for each instance, are adapted at test-time to fit the given image. Overall this is a neat idea and the experimental results are compelling. However, I found the paper difficult to understand on first readings and have several suggestions for improvement of the presentation.