Domain and Task-Focused Example Selection for Data-Efficient Contrastive Medical Image Segmentation
About
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.
Related benchmarks
| Task | Dataset | Result | Rank | |
|---|---|---|---|---|
| Liver Segmentation | LiTS (Liver Tumor Segmentation) | Dice89.6 | 29 | |
| Liver Segmentation | TotalSegmentator | Dice Coefficient75.3 | 22 | |
| Liver Segmentation | MSD cross-domain (test) | Dice Coefficient76.7 | 22 | |
| Kidney Segmentation | BTCV cross-domain | Dice Coefficient49.2 | 10 | |
| Kidney Segmentation | KiTS 2023 (test) | Dice Score0.957 | 10 |