UC-Hybrid: Uncertainty-based contrastive learning on hybrid network for medical image segmentation

Medical image segmentation has made remarkable progress with advances in deep learning technology, depending on the quality and quantity of labeled data. Although various deep learning model structures and training methods have been proposed and high performance has been published, limitations such as inter-class accuracy bias exist in actual clinical applications, especially due to the significant lack of small object performance in multi-organ segmentation tasks. In this paper, we propose an uncertainty-based contrastive learning technique, namely UncerNCE, with an optimal hybrid architecture for high classification and segmentation performance of small organs. Our backbone architecture adopts a hybrid network that employs both convolutional and transformer layers, which have demonstrated remarkable performance in recent years. The key proposal of this study addresses the multi-class accuracy bias and resolves a common tradeoff in existing studies between segmenting regions of small objects and reducing overall noise (i.e., false positives). Uncertainty based contrastive learning based on the proposed hybrid network performs spotlight learning on selected regions based on uncertainty and achieved accurate segmentation for all classes while suppressing noise. Comparison with state-of-the-art techniques demonstrates the superiority of our results on BTCV and 1K data.