Weakly supervised learning for multi-class medical image segmentation via feature decomposition

Abstract

Weakly supervised learning with image-level labels, releasing deep learning from highly labor-intensive pixel-wise annotation, has gained great attention for medical image segmentation. However, existing weakly supervised methods are mainly designed for single-class segmentation while leaving multi-class medical image segmentation rarely-explored. Different from natural images, label symbiosis, together with location adjacency, are much more common in medical images, making it more challenging for multi-class segmentation. In this paper, we propose a novel weakly supervised learning method for multi-class medical image segmentation with image-level labels. In terms of the multi-class classification backbone, a multi-level classification network encoding multi-scale features is proposed to produce binary predictions, together with the corresponding CAMs, of each class separately. To address the above issues (i.e., label symbiosis and location adjacency), a feature decomposition module based on semantic affinity is first proposed to learn both class-independent and class-dependent features by maximizing the inter-class feature distance. Through a cross-guidance loss to jointly utilize the above features, label symbiosis is largely alleviated. In terms of location adjacency, a mutually exclusive loss is constructed to minimize the overlap among regions corresponding to different classes. Experimental results on three datasets demonstrate the superior performance of the proposed weakly-supervised framework for both single-class and multi-class medical image segmentation. We believe the analysis in this paper would shed new light on future work for multi-class medical image segmentation. The source code of this paper is publicly available at https://github.com/HustAlexander/MCWSS.