3D Graph-S2Net: Shape-Aware Self-ensembling Network for Semi-supervised Segmentation with Bilateral Graph Convolution

Abstract

Semi-supervised learning (SSL) algorithms have attracted much attentions in medical image segmentation due to challenge in acquiring pixel-wise annotations by using unlabeled data. However, most of existing SSLs neglected the geometric shape constraint in object, leading to unsatisfactory boundary and non-smooth of object. In this paper, we propose a shape-aware semi-supervised 3D medical image segmentation network, named 3D Graph-S2Net, which incorporates the flexible shape information and learns duality constraints between semantics and geometrics in the graph domain. Specifically, our method consists of two parts: a multi-task learning network (3D S2Net) and a graph-based cross-task module (3D BGCM). The 3D S2Net improves the existing self-ensembling model (i.e., Mean-Teacher model) by adding a signed distance map (SDM) prediction task, which encodes richer features of object shape and surface. Moreover, the 3D BGCM explores the co-occurrence relations between the semantics segmentation and SDM prediction task, so that the network learns stronger semantic and geometric correspondences from both labeled and unlabeled data. Experimental results on the Atrial Segmentation Challenge confirm that our 3D Graph-S2Net outperforms the state-of-the-arts in semi-supervised segmentation.