Affinity-Graph-Guided Contractive Learning for Pretext-Free Medical Image Segmentation with Minimal Annotation

TL;DR

An affinitygraph-guided semi-supervised contrastive learning framework (Semi-AGCL) is proposed by establishing additional affinity-graph-based supervision signals between the student and teacher network, to achieve medical image segmentation with minimal annotations without pretext.

Abstract

The combination of semi-supervised learning (SemiSL) and contrastive learning (CL) has been successful in medical image segmentation with limited annotations. However, these works often rely on pretext tasks that lack the specificity required for pixel-level segmentation, and still face overfitting issues due to insufficient supervision signals resulting from too few annotations. Therefore, this paper proposes an affinity-graph-guided semi-supervised contrastive learning framework (Semi-AGCL) by establishing additional affinity-graph-based supervision signals between the student and teacher network, to achieve medical image segmentation with minimal annotations without pretext. The framework first designs an average-patch-entropy-driven inter-patch sampling method, which can provide a robust initial feature space without relying on pretext tasks. Furthermore, the framework designs an affinity-graph-guided loss function, which can improve the quality of the learned representation and the model generalization ability by exploiting the inherent structure of the data, thus mitigating overfitting. Our experiments indicate that with merely 10% of the complete annotation set, our model approaches the accuracy of the fully annotated baseline, manifesting a marginal deviation of only 2.52%. Under the stringent conditions where only 5% of the annotations are employed, our model exhibits a significant enhancement in performance surpassing the second best baseline by 23.09% on the dice metric and achieving an improvement of 26.57% on the notably arduous CRAG and ACDC datasets.

Figures (2)

• Figure 1: The proposed framework. For labeled data, we directly use the supervised loss $\mathcal{L}_{sup}$ to update the student network. For unlabeled data, we first slice the image into patches, then bridge an affinity graph loss $\mathcal{L}_{AGG}^{PL}$ between pseudo labels of student and teacher networks, and also design a new loss $\mathcal{L}_{AGG}^{RW}$ using the reweighting hard negative sample based on the edge of affinity graph. In the affinity-graph-based losses, we use low $A_{ii}$ to construct a negative hard sample and try to pull positive pairs closer (increase $A_{ii}$) and push negative pairs away. Besides, the blue arrows use labeled data, and the rest (black arrows) are unlabeled data; we use a mixture of labeled and unlabeled data, so it is a semisupervised task rather than a self-supervised task. SA: Strong Augmentation, WA: Weak Augmentation.
• Figure 2: The visualization of the proposed framework and baselines on the CRAG, LA and dataset. The first and second rows are the segmentation results with labeled ratios of $5\%$ and $10\%$, respectively. The red boxes indicate that our method outperforms other baselines. GT: Ground Truth.