Prototype Contrastive Consistency Learning for Semi-Supervised Medical Image Segmentation

TL;DR

This paper introduces PCCS, a semi-supervised medical image segmentation method that uses prototype contrastive learning with uncertainty estimation to improve segmentation accuracy with limited labeled data.

Contribution

The paper proposes a novel prototype contrastive learning framework with uncertainty estimation and a prototype updating mechanism for semi-supervised medical image segmentation.

Findings

PCCS outperforms state-of-the-art methods on medical segmentation tasks.

The use of signed distance and uncertainty maps improves prototype quality.

Uncertainty-consistency loss enhances the reliability of unlabeled data utilization.

Abstract

Medical image segmentation is a crucial task in medical image analysis, but it can be very challenging especially when there are less labeled data but with large unlabeled data. Contrastive learning has proven to be effective for medical image segmentation in semi-supervised learning by constructing contrastive samples from partial pixels. However, although previous contrastive learning methods can mine semantic information from partial pixels within images, they ignore the whole context information of unlabeled images, which is very important to precise segmentation. In order to solve this problem, we propose a novel prototype contrastive learning method called Prototype Contrastive Consistency Segmentation (PCCS) for semi-supervised medical image segmentation. The core idea is to enforce the prototypes of the same semantic class to be closer and push the prototypes in different semantic classes far away from each other. Specifically, we construct a signed distance map and an uncertainty map from unlabeled images. The signed distance map is used to construct prototypes for contrastive learning, and then we estimate the prototype uncertainty from the uncertainty map as trade-off among prototypes. In order to obtain better prototypes, based on the student-teacher architecture, a new mechanism named prototype updating prototype is designed to assist in updating the prototypes for contrastive learning. In addition, we propose an uncertainty-consistency loss to mine more reliable information from unlabeled data. Extensive experiments on medical image segmentation demonstrate that PCCS achieves better segmentation performance than the state-of-the-art methods. The code is available at https://github.com/comphsh/PCCS.