Online pseudo labeling for polyp segmentation with momentum networks

TL;DR

This paper introduces an online pseudo labeling method with momentum networks for semi-supervised polyp segmentation, improving label quality during training and achieving high accuracy with limited labeled data.

Contribution

It proposes a novel online pseudo labeling strategy that updates the teacher model during training using a momentum network, enhancing pseudo label quality in semi-supervised learning.

Findings

Achieved 84.1% Dice Score on five datasets with only 20% labeled data.

Surpassed common semi-supervised methods by 3%.

Approached fully-supervised performance on some datasets.

Abstract

Semantic segmentation is an essential task in developing medical image diagnosis systems. However, building an annotated medical dataset is expensive. Thus, semi-supervised methods are significant in this circumstance. In semi-supervised learning, the quality of labels plays a crucial role in model performance. In this work, we present a new pseudo labeling strategy that enhances the quality of pseudo labels used for training student networks. We follow the multi-stage semi-supervised training approach, which trains a teacher model on a labeled dataset and then uses the trained teacher to render pseudo labels for student training. By doing so, the pseudo labels will be updated and more precise as training progress. The key difference between previous and our methods is that we update the teacher model during the student training process. So the quality of pseudo labels is improved during the student training process. We also propose a simple but effective strategy to enhance the quality of pseudo labels using a momentum model -- a slow copy version of the original model during training. By applying the momentum model combined with re-rendering pseudo labels during student training, we achieved an average of 84.1% Dice Score on five datasets (i.e., Kvarsir, CVC-ClinicDB, ETIS-LaribPolypDB, CVC-ColonDB, and CVC-300) with only 20% of the dataset used as labeled data. Our results surpass common practice by 3% and even approach fully-supervised results on some datasets. Our source code and pre-trained models are available at https://github.com/sun-asterisk-research/online learning ssl