Domain and Task-Focused Example Selection for Data-Efficient Contrastive Medical Image Segmentation

TL;DR

This paper introduces PolyCL, a novel self-supervised contrastive learning framework for medical image segmentation that effectively learns from limited annotations without requiring pixel-level labels, improving accuracy and transferability.

Contribution

PolyCL leverages inherent image relationships for contrastive learning in a task-related manner, integrating SAM for refinement and volumetric segmentation, advancing data-efficient medical image segmentation.

Findings

PolyCL outperforms baseline methods in low-data scenarios.

PolyCL achieves superior cross-domain segmentation accuracy.

SAM integration enhances mask refinement and volumetric segmentation.

Abstract

Segmentation is one of the most important tasks in the medical imaging pipeline as it influences a number of image-based decisions. To be effective, fully supervised segmentation approaches require large amounts of manually annotated training data. However, the pixel-level annotation process is expensive, time-consuming, and error-prone, hindering progress and making it challenging to perform effective segmentations. Therefore, models must learn efficiently from limited labeled data. Self-supervised learning (SSL), particularly contrastive learning via pre-training on unlabeled data and fine-tuning on limited annotations, can facilitate such limited labeled image segmentation. To this end, we propose a novel self-supervised contrastive learning framework for medical image segmentation, leveraging inherent relationships of different images, dubbed PolyCL. Without requiring any pixel-level annotations or unreasonable data augmentations, our PolyCL learns and transfers context-aware discriminant features useful for segmentation from an innovative surrogate, in a task-related manner. Additionally, we integrate the Segment Anything Model (SAM) into our framework in two novel ways: as a post-processing refinement module that improves the accuracy of predicted masks using bounding box prompts derived from coarse outputs, and as a propagation mechanism via SAM 2 that generates volumetric segmentations from a single annotated 2D slice. Experimental evaluations on three public computed tomography (CT) datasets demonstrate that PolyCL outperforms fully-supervised and self-supervised baselines in both low-data and cross-domain scenarios. Our code is available at https://github.com/tbwa233/PolyCL.