Semi-supervised Class Imbalanced Deep Learning for Cardiac MRI Segmentation

Abstract

Despite great progress in semi-supervised learning (SSL) that leverages unlabeled data to improve the performance over fully supervised models, existing SSL approaches still fail to exhibit good results when faced with a severe class imbalance problem in medical image segmentation. In this work, we propose a novel Mean-teacher based class imbalanced learning framework for cardiac magnetic resonance imaging (MRI) segmentation, which can effectively conquer the problems of class imbalance and limited labeled data simultaneously. Specifically, in parallel to the traditional linear-based classifier, we additionally train a prototype-based classifier that makes dense predictions by matching test samples with a set of prototypes. The prototypes are iteratively updated by in-class features encoded in the entire sample set, which can better guide the model training by alleviating the class-wise bias exhibited in each individual sample. To reduce the noises in the pseudo labels, we propose a cascaded refining strategy by utilizing two multi-level tree filters that are built upon pairwise pixel similarity in terms of intensity values and semantic features. With the assistance of pixel affinities, soft pseudo labels are properly refined on-the-fly. Upon evaluation on ACDC and MMWHS, two cardiac MRI datasets with prominent class imbalance problem, the proposed method demonstrates the superiority compared to several state-of-the-art methods, especially in the case where few annotations are available (Code is available in https://github.com/IsYuchenYuan/SSCI).