Multi-scale graph harmonies: Unleashing U-Net’s potential for medical image segmentation through contrastive learning

Abstract

Medical image segmentation is essential for accurately representing tissues and organs in scans, improving diagnosis, guiding treatment, enabling quantitative analysis, and advancing AI-assisted healthcare. Organs and lesion areas in medical images have complex geometries and spatial relationships. Due to variations in the size and location of lesion areas, automatic segmentation faces significant challenges. While Convolutional Neural Networks (CNNs) and Transformers have proven effective in segmentation task, they still possess inherent limitations. Because these models treat images as regular grids or sequences of patches, they struggle to learn the geometric features of an image, which are essential for capturing irregularities and subtle details. In this paper we propose a novel segmentation model, MSGH, which utilizes Graph Neural Network (GNN) to fully exploit geometric representation for guiding image segmentation. In MSGH, we combine multi-scale features from Pyramid Feature and Graph Feature branches to facilitate information exchange across different networks. We also leverage graph contrastive representation learning to extract features through self-supervised learning to mitigate the impact of category imbalance in medical images. Moreover, we optimize the decoder by integrating Transformer to enhance the model’s capability in restoring the intricate image details feature. We conducted a comprehensive experimental study on ACDC, Synapse and BraTS datasets to validate the effectiveness and efficiency of MSGH. Our method achieved an improvement of 2.56–13.41%, 1.04–5.11% and 1.77–3.35% of dice on the three segmentation tasks respectively. The results demonstrate that our model consistently performs well compared with state-of-the-art models. The source code is accessible at https://github.com/Dorothywujie/MSGH.