[Paper Review] K-Net: Towards Unified Image Segmentation

Wenwei Zhang, Jiangmiao Pang+2 more|arXiv (Cornell University)|Jun 28, 2021

Advanced Neural Network Applications66 references197 citations

TL;DR

K-Net unifies semantic, instance, and panoptic segmentation with learnable group kernels that are dynamically updated per image, trained end-to-end via bipartite matching, achieving state-of-the-art single-model results while being NMS- and box-free.

ABSTRACT

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 state-of-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/.

Motivation & Objective

• Motivate unifying semantic, instance, and panoptic segmentation under a common kernel-based framework.
• Propose a dynamic, group-aware kernel update mechanism to handle varying object instances.
• Train instance kernels end-to-end using bipartite matching without bounding boxes or NMS.
• Show that a unified kernel-based approach improves panoptic, instance, and semantic segmentation on standard benchmarks.

Proposed method

• Represent segmentation targets as a fixed set of N kernels, each predicting one mask for either a potential instance or a semantic class.
• Introduce a kernel update head that form- and content-adaptively updates kernels using group features assembled from current predictions.
• Apply adaptive feature and kernel interaction with gating to refine kernels and masks iteratively for S rounds.
• Use multi-head attention among kernels to allow contextual interaction and derive final masks from updated kernels.
• Train instance kernels with a mask-driven, end-to-end Hungarian assignment (no boxes, no NMS).
• Apply the same kernel framework to panoptic or semantic segmentation by combining instance and semantic kernels as appropriate.

Experimental results

Research questions

• RQ1 Can semantic, instance, and panoptic segmentation be solved under a single kernel-based framework?
• RQ2 How can kernels be made adaptive to content to reliably separate instances of varying shapes and scales?
• RQ3 Is end-to-end training with bipartite (Hungarian) assignment feasible for instance kernels without relying on boxes or NMS?
• RQ4 What performance gains arise on COCO for panoptic/instance and on ADE20K for semantic segmentation when using K-Net?

Key findings

• achieves state-of-the-art single-model panoptic segmentation on COCO val with 47.1 PQ (R50-FPN) and 54.6 PQ on Swin-L, surpassing prior methods.
• achieves 54.3 mIoU on ADE20K val for semantic segmentation when integrated with baseline architectures, indicating strong cross-task benefits.
• instance segmentation results are competitive with Cascade Mask R-CNN while being box-free and NMS-free, and show significantly faster inference (e.g., K-Net-N256 at 19.8 FPS vs 10.3 FPS for Cascade on COCO).
• the kernel update head with group-feature assembly and adaptive gating substantially improves AP (e.g., Table 4 shows AP rising from 18.2 baseline to 34.7 with full components).
• 100 instance kernels are sufficient for strong COCO performance, with performance saturating beyond a certain number of rounds (~3) of kernel updates.
• K-Net can outperform several box-based and kernel-based baselines with fewer training epochs (36), and provides robust performance with modest compute.