GU-Net: Causal relationship-based generative medical image segmentation model

Abstract

Due to significant anatomical variations in medical images across different cases, medical image segmentation is a highly challenging task. Convolutional neural networks have shown faster and more accurate performance in medical image segmentation. However, existing networks for medical image segmentation mostly rely on independent training of the model using data samples and loss functions, lacking interactive training and feedback mechanisms. This leads to a relatively singular training approach for the models, and furthermore, some networks can only perform segmentation for specific diseases. In this paper, we propose a causal relationship-based generative medical image segmentation model named GU-Net. We integrate a counterfactual attention mechanism combined with CBAM into the decoder of U-Net as a generative network, and then combine it with a GAN network where the discriminator is used for backpropagation. This enables alternate optimization and training between the generative network and discriminator, enhancing the expressive and learning capabilities of the network model to output prediction segmentation results closer to the ground truth. Additionally, the interaction and transmission of information help the network model capture richer feature representations, extract more accurate features, reduce overfitting, and improve model stability and robustness through feedback mechanisms. Experimental results demonstrate that our proposed GU-Net network achieves better segmentation performance not only in cases with abundant data samples and relatively simple segmentation targets or high contrast between the target and background regions but also in scenarios with limited data samples and challenging segmentation tasks. Comparing with existing U-Net networks with attention mechanisms, GU-Net consistently improves Dice scores by 1.19%, 2.93%, 5.01%, and 5.50% on ISIC 2016, ISIC 2017, ISIC 2018, and Gland Segmentation datasets, respectively.

Keywords

• Medical image segmentation
• Convolutional neural network
• Attention mechanism
• Causal reasoning
• Interactive training