Enhance Polyp Segmentation via Supervised Contrastive Learning
摘要
Accurate segmentation of colorectal polyps is crucial for the early diagnosis and prevention of colorectal cancer. However, automatic polyp segmentation remains a challenging task due to the high diversity of polyps in shape, size, color, and texture, as well as their low contrast with the surrounding mucosal background. Existing deep learning-based segmentation methods, which primarily rely on pixel-wise supervised losses, may struggle to learn sufficiently discriminative features to distinguish polyps from the background. This paper proposes a novel polyp segmentation framework, namely Sup-Polyp, that enhances the model’s feature representation capability by introducing supervised contrastive learning. Specifically, building upon a standard segmentation network, we first design a projection head to map the features extracted by the backbone into a dedicated embedding Space. Secondly, we construct a polyp memory bank to store and dynamically update the average feature representations of polyp regions during training. Finally, in each training iteration, we retrieve features from the memory bank and, in conjunction with the image features from the current batch, optimize them using an InfoNCE loss function. This contrastive loss pulls the features of polyp regions to move closer to the features in the memory bank within the embedding space, while simultaneously pushing them away from background features. Experimental results on public polyp segmentation datasets demonstrate that our proposed method can effectively improve segmentation performance without additional inference costs.