CNN-based medical image segmentation

In the domain of medical image segmentation, algorithms based on Convolutional Neural Networks (CNNs) have made considerabl progress. U-Net [10] is a classic CNN-based image segmentation architecture, which consists of an encoder and a decoder, forming a symmetrical U-shaped structure. Skip connections with channel concatenation are employed between the encoder and decoder to fuse shallow and deep features. However, U-Net can only extract feature information from each individual slice, while ignoring the inter-slice feature correlations and resulting in low smoothness of the segmentation results. To overcome this limitation, Özgün Çiçek et al. [34] proposed 3D-Unet in 2016, which can directly integrate multiple slices into a 3D image for training, which improves improving the continuity of variations between adjacent image masks. Residual U-Net [12] incorporates residual connections to facilitate information flow and gradient propagation. Han et al. introduced FCN [35], a fully convolutional network that enables dense pixel-wise prediction of input images. It incorporates skip connections to fuse features at different scales, which can strengthen the model’s ability to perceive objects of different scales and capture fine details, thereby improving segmentation accuracy. While these CNN-based networks have achieved remarkable success in image segmentation, they primarily focus on local perception and feature extraction and share parameters between each convolutional layer, which may limit their capability to capture long-range contextual information. For certain medical image segmentation tasks, global context information is vital for accurate segmentation decisions, and CNNs may not fully capture these global dependencies.

Attention mechanism for medical imaging

The attention mechanism has also been extensively applied and researched in the field of image segmentation, leading to significant advancements. O. Oktay et al. [11] proposed AttnUnet, which introduced attention mechanisms into the skip connections of U-Net. By dynamically adjusting the weights of feature maps in the skip connections, the network focuses more on crucial information, thereby enhancing segmentation accuracy. In 2017, the Transformer model [13] emerged, utilizing multi-head self-attention mechanisms to capture global contextual information more precisely. Subsequently, Dosovitskiy et al. [16] introduced Transformer to the image processing domain and proposed Vision Transformer (ViT). ViT divides the input image into multiple pixel patches and flattens them into one-dimensional sequences, effectively leveraging global image features. This approach transforms image segmentation tasks into one-dimensional sequence prediction tasks. However, compared to Convolutional Neural Networks (CNNs), Transformers perform poorly in extracting detailed features and are prone to losing fine-grained content in liver segmentation tasks. As a consequence, researchers have begun exploring the combination of CNNs and Transformers for medical image segmentation tasks.

Chen et al. [17] proposed TransUnet, which concatenates multiple Transformer encoder layers and connects them with the CNN layers in the U-Net encoder to retain the global contextual perception of Transformers and the local contextual perception of CNNs. Gao et al. [18] introduced UTNet, which extracts partial features using convolutional layers and incorporates a novel self-attention mechanism to reduce network complexity. Wang et al. [19] proposed TransBTS, incorporating four layers of ViT’s Transformer modules in the network to fuse contextual features and perform classification output through CNN modules. However, this approach increases the computational burden while enhancing segmentation accuracy. Unetr [20] removes the CNN layers in the encoder and adopts 12 Transformer encoder layers to learn the sequential representation of input volumes, effectively capturing global-scale information. Liu et al. [21] introduced Swin Transformer, which utilizes hierarchical attention mechanisms and cross-window attention mechanisms to enable Transformers to be more applicable in image processing. Subsequently, Tang et al. [22] improved the Swin Transformer module and combined it with CNNs, proposing SwinUnetr, which achieved excellent performance in 3D medical image segmentation tasks. The Transformer model models dependencies in the input sequence through self-attention mechanisms, making it highly effective in handling global information and long-range dependencies. However, when dealing with small objects that may have fewer pixels and local context, the self-attention mechanisms may not fully capture the detailed information of these small objects, leading to suboptimal performance of Transformer models in small object segmentation tasks.

GANs in image segmentation

Traditional generative models are mainly based on probability density estimation, requiring manual specification of prior probability distributions and loss functions, and are susceptible to problems such as overfitting. Hence, Goodfellow et al. [30] proposed Generative Adversarial Networks (GANs), which enable generative models to better generate data similar to real data. To improve the stability and precision of generator output, Mirza et al. [36] introduced conditional Generative Adversarial Networks (cGANs) based on GANs. Compared to GANs that only accept noise vectors as input, cGANs introduce conditional vectors, allowing the generator to approximate the real data distribution more closely. By introducing cGANs into image segmentation tasks, the precision of segmentation results can be further improved.

P. Isola et al. [37] proposed Pix2PixGAN, which utilizes the U-Net network as the generator and incorporates the PatchGAN architecture in the discriminator. The discriminator conducts local region-based discrimination instead of whole-image discrimination, allowing it to focus more on image details and local structures, thus enhancing sensitivity to image details and achieving precise mapping from input images to target domain images. However, Pix2Pix-GAN demands a large amount of paired image data for training, i.e., matching input images with target domain images. In the domain of medical imaging, obtaining a large-scale paired dataset is arduous and costly. Beji A et al. [24] proposed a two-stage architecture, Seg2GAN. Seg2GAN generates synthetic data distributions to improve the quality of input images, and then uses the new data for segmentation. Experiments conducted on datasets such as retinal and coronary vessels show that using the new synthetic data results in better segmentation performance compared to models trained only on original real images. Guo K et al. [25] addressed issues such as model collapse, gradient vanishing, and convergence failure in GANs by proposing MedGAN, which uses Wasserstein loss as a convergence measure to assess the convergence of the generator and discriminator. The model’s advantages in terms of convergence, training speed, and the visual quality of generated samples have been validated on multiple datasets, including scabies and blisters. Additionally, GAN-based medical image segmentation algorithms may encounter problems such as mode collapse and mode oscillation, and the training process is frequently unstable, especially in the initial phases and parameter selection of the model. This may lead to difficulty in achieving convergence during the training process, requiring more time and computational resources to obtain satisfactory segmentation results.

Transformer in generative adversarial networks

Inspired by the notions of Transformers and GANs, researchers have put forward a series of networks that combine Transformer and GANs [38–43]. Esser et al. [14] introduced a novel self-attention mechanism that takes into account both local and global contextual information to handle long-range dependencies in images, thereby better grasping the structure and details of the images. They combined it with a convolutional GANs model to learn context-rich visual feature codebooks, facilitating high-resolution image synthesis. Wang et al. [38] proposed TransGAN, which applies the Transformer architecture to image generation tasks, converting it into an unconditional GANs framework. By capturing long-range dependencies and modeling global context, TransGAN can generate high-quality medical images. However, when applied to medical image segmentation tasks, TransGAN might face the challenge of insufficient explicit guidance at the pixel-level annotation. P. Kumar et al.[26] proposed a digital artwork restoration method based on Generative Adversarial Networks (GANs) that incorporates dual (spatial and channel) attention and channel transformers. The method uses spatial and channel attention layers in the encoder part of the generator and employs channel transformers in the skip connections. This approach achieves multi-scale feature fusion and reduces the semantic gap between the encoder and decoder layers. The method outperforms existing artwork restoration techniques on two datasets. Transformers can help networks better capture global information in images, particularly for structures that cover large spatial extents in segmentation tasks. The application of GAN in medical image segmentation tasks enables the network to generate more realistic segmentation results. However, combining Transformers and GANs can result in increased computational complexity, requiring more computational resources and time for training and inference, fundamentally limiting its applicability.

Generator

The network architecture of the generator is a sophisticated three-dimensional segmentation network specifically designed for end-to-end segmentation of 3D medical images. As depicted in Fig 1A, the architecture features a carefully constructed encoder-decoder structure. The encoder is composed of four sequential downsampling layers that progressively contract the spatial dimensions of the input image while expanding the depth. Starting from an initial image size of C×D×256×256, where C represents the number of channels, D the depth, and 256 × 256 the spatial dimensions, the encoder reduces this to 256×D/16×16×16 through a series of convolutional and pooling operations. This process is crucial for extracting and compressing high-level semantic features from the input medical images, allowing the network to capture complex patterns and structures inherent in the data.

The decoder, in contrast, is designed with four upsampling layers that gradually expand the feature maps back to the original image dimensions. The decoder aims to reconstruct the image to match the pixel-level segmentation results with the ground truth labels. To ensure consistency, the shape of the interpolated image aligns with that of the image before pooling in the corresponding encoder layer. This process is facilitated by skip connections, which bridge the encoder and decoder layers, allowing for the transfer of fine-grained low-level features that complement the high-level features processed in the decoder.

To enhance the feature extraction capabilities, the network incorporates Multiple-Angle (MA) attention mechanisms. These mechanisms leverage information from different axes to ensure comprehensive feature capture from various perspectives. Convolutional kernels throughout the network are uniformly set to 3×3, which simplifies the model while preserving high segmentation accuracy [44]. Batch Normalization (BN) layers [45] are employed to normalize the outputs of the convolutions, which helps to accelerate the training process and improve the generalization performance of the network. Given the extensive number of parameters in the generator, the ReLU activation function is applied after both upsampling and downsampling operations. This choice of activation function is instrumental in mitigating overfitting by retaining only the positive values of the input features, thus enhancing the model’s robustness.

Dicscriminator

We designed a discriminator network comprising three convolutional layers, as illustrated in Fig 1B, specifically to differentiate between real and synthetic input images. The network takes as input a CT image that is concatenated with a label indicating whether it is real or generated. This concatenated input is processed through three downsampling layers, which progressively reduce the spatial dimensions while increasing the depth of the feature maps. Each convolutional layer applies a series of filters to extract and learn high-level feature representations from the input image. The downsampling operations are designed to capture complex patterns and textures that are crucial for distinguishing real images from those generated by the generator.

After the convolutional operations, the extracted features are fed into a fully connected layer, which ultimately uses a sigmoid activation function to produce a probability score. This score reflects the likelihood that the input sample is real data. A score close to 1 indicates high confidence in the input being real, while a score close to 0 suggests that the input is more likely a fake sample generated by the generator. This probabilistic output is critical for guiding the training process, ensuring that the discriminator effectively challenges the generator to improve the quality of its outputs.

To maintain the balance between the discriminator and the generator and to prevent the discriminator from becoming too dominant, which could lead to training instability, a dropout layer with a rate of 0.5 is incorporated after each downsampling layer. This dropout layer randomly omits a portion of neurons during training, which helps to regularize the model and prevent overfitting. By introducing this stochastic element, the network maintains robustness and ensures that the discriminator’s learning does not overshadow the generator’s progress.

Multi-axis attention mechanism (MA)

The Slice Feature Projection Module (FP) is introduced in MA, which projects the 3D medical images onto a 2D plane along a specific axis, enhancing the proportion of target information. If the proportion of target segmentation content before projection is O(1/n³), after projection to the 2D plane, the proportion is improved to O(1/n²), reducing the loss of critical information. The FP structure, as illustrated in Fig 2A, obtains relevant information on the projection axes by globally average-pooling and max-pooling the input features (C×D×H×W) and then summing them. This process generates the key (K) and query (Q) values. The shapes differ when projected onto different axes planes. When projected onto the axial plane, the shape is C×H×W; when projected onto the sagittal plane, the shape is C×D×H; and when projected onto the coronal plane, the shape is C×D×W. The specific formulas are defined as shown in Eq 1.

(1)

Download:

• PNG
  larger image
• TIFF
  original image

Fig 2. Overview of multi-axis feature extraction structure.

In MA-cGAN, we first employ slice feature projection (FP) to project the 3D image feature information along axes onto a 2D plane. Then, the obtained 2D feature information from multiple directions is fused with the input features. Finally, a learnable factor is applied to weight the fused output, enabling the comprehensive integration of slice feature information from different axes. (a) Slice Feature Projection(FP) (b) Multi-Axis Attention Mechanism(MA).

https://doi.org/10.1371/journal.pone.0312105.g002

The value tensor (V)is obtained by applying a 1×1 convolution to the input features, preserving its shape (C×D×H×W). The key (K)is processed using two 3×3 convolutions and concatenated with the query (Q). The tensor K is subjected to two consecutive 3x3 convolutions followed by concatenation with Q. Then, two 1x1 convolutions and one Unsqueeze operation are applied to expand the dimensions and replicate along the depth dimension, ensuring the shape remains consistent with the input features, resulting in an attention matrix R (C×D×H×W). The attention matrix R is then broadcast-multiplied with the value (V), and the multiplication result is concatenated with the input features after expanding its dimensions. This process generates a mixed attention matrix H (C×2×D×H×W). Finally, the mixed attention matrix H is summed along the depth dimension and averaged, resulting in the two-dimensional projected feature information (C×1×1×1).

Three FP modules were employed in MA, projecting the information of three-dimensional medical images onto axial, sagittal, and coronal planes respectively. This approach facilitated the acquisition of two-dimensional projection features in three different directions. The two-dimensional projected feature information is fused with the input feature through element-wise multiplication, followed by weighted summation using three learnable factors to achieve the fusion of multidimensional features, as illustrated in Fig 2B. MA effectively utilizes the multi-angle information of the input image, preserving more detailed features. The formulation for the fusion process is given by Eq 2.

(2)

Whereas α represents the three projection directions, β_α represents the learnable weighted factors in each direction, FP(·) represents the two-dimensional feature information obtained by projecting along a specific direction, and x represents the input features.

The training process of Generative adversarial networks (GANs)

The training process of the generative adversarial network is illustrated in Fig 3, involving iterative training between the generator and the discriminator. The generator’s task is to generate fake labels that resemble the real labels as much as possible, with the expectation that the discriminator will classify the fake labels as real. Initially, the CT image is passed through the generator segmentation network, which performs a series of operations such as downsampling, upsampling, and skip connections, to output fake labels. The CT image is concatenated with the fake labels and fed into the discriminator. The loss between the discriminator’s output and a matrix of all ones is computed, and then backpropagation is used to update the weights of the generator, enhancing its ability to generate more realistic fake labels. This completes a single training step for the generator. After completing a single training step for the generator, its parameters are frozen, and training is performed for the discriminator.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 3. Overview of alternating adversarial training process.

In the alternating adversarial training process, the generator model generates samples to deceive the discriminator model, while the discriminator model provides feedback by accurately distinguishing between the generated samples and real samples. The generator and discriminator compete and cooperate, iteratively optimizing the adversarial loss function between them until the generator can generate realistic samples.

https://doi.org/10.1371/journal.pone.0312105.g003

The task of the discriminator is to distinguish between real labels and fake labels generated by the generator. The CT image is used as a conditioning factor and is concatenated with either real or fake labels before being input into the discriminator. Based on the authenticity of the input labels, the loss is computed between the discriminator’s output and a matrix of all ones or all zeros. Through backpropagation, the discriminator’s weights are updated to improve its ability to discriminate between real and fake labels, completing a single training step for the discriminator.

By iteratively training the generator and the discriminator, the quality of the generated labels and the discriminative accuracy of the discriminator are continuously improved. This process continues until the generated fake labels are sufficiently close to the real labels, and the discriminator is no longer able to accurately distinguish between real and fake labels, achieving a Nash equilibrium. This marks the completion of the training process for the generative adversarial network. At this point, the generated fake labels possess high quality and can be used for high-precision liver and tumor segmentation in unlabeled CT images.

Loss function

In traditional generative adversarial network, the loss function is defined as Eq (3): (3)

Where E represents the expectation value of the distribution function, x represents the preprocessed CT image, z represents the noise input to the generator, G(z) represents the generated fake image by the generator, and D(x) represents the discriminator’s probability of judging an image as real. For the discriminator D, the goal is to classify real data x as true and fake data G(z) as false, which means D(x) should approach 1 and D(G(z)) should approach 0. The objective is to maximize V(D,G). On the other hand, for the generator G, the objective is to have the discriminator classify fake data G(z) as real, which means D(G(z)) should approach 1. The objective is to minimize V(D,G). This optimization approach uses the JS divergence, which can result in fixed gradients and prevent updates to the generator when the two distributions have no overlap or negligible overlap. Wasserstein GAN [46] addresses this issue by using the Wasserstein distance, which is defined in Eq (4) as follows: (4)

Where inf represents the maximum lower bound, ∏(P,Q) represents the joint distribution set combining distributions P and Q, and E(x,y) represents the expectation. Unlike JS divergence, the Wasserstein distance not only considers the differences between distributions but also takes into account the minimum cost of transforming one distribution into another. It can measure the distance between two distributions. We removed the log function from the loss function, which reduces the issues of gradient vanishing and exploding, resulting in a more stable training process.

Tversky Loss was proposed by Tversky et al. [47] in 2017 to address the issue of high precision but low recall in image segmentation caused by class imbalance. The formula for Tversky Loss is given by Eq (5): (5)

Where y_i in T(y,G(z,y)) represents the i-th voxel as a lesion voxel, 1−y_i represents the i-th voxel as a normal voxel. G(x)_i represents the probability of predicting the i-th voxel as a lesion voxel, and 1−G(x)_i represents the probability of predicting the i-th voxel as a normal voxel. (1−y_i)⋅G(x)_i represents the false positive case, y_i⋅(1−G(x)_i) represents the false negative case, and α and β respectively denote the parameters controlling the weight of false positive and false negative.

We incorporated additional conditional information and introduced Tversky Loss into the loss function, effectively addressing the issue of class imbalance and further enhancing the segmentation performance of the generator. The loss function is defined as Eq (6) below: (6)

Dataset

We conducted image segmentation experiments on the LiTS dataset [48]. The dataset consists of 131 abdominal CT image files with annotations for liver and liver tumors. Each file varies in terms of liver size, shape, and pathological features. The abdominal CT images and corresponding annotations have a shape of depth × width × height (D × W × H), where both W and H are 512, and D ranges from 86 to 987. In our experiments, we excluded 21 CT images with inaccurate annotations, resulting in a total of 110 CT images. Using random sampling, we selected 80 images as the training set, 10 images as the validation set, and the remaining 20 images as the test set.

To reduce the impact of irrelevant tissues in CT images on the segmentation results, we performed two preprocessing operations on the dataset: ① Irrelevant tissue removal: Along the depth direction of the slices, we identified the start and end indices of the labeled voxel points and removed the slices outside this index range, thereby reducing the interference from other abdominal organs during liver segmentation. After preprocessing, the image depth (D) values ranged from 28 to 312. ② Windowing: We set the window level of the CT images to 70 Hounsfield Units (HU) and the window width to 100 HU. This adjustment resulted in the intensities of the CT images ranging from -30 HU to 170 HU. Tissues with intensities below -30 HU were displayed as black, while those above 170 HU were displayed as white. This setting reduced the influence of excessively high or low voxel intensities on the experiment and improved the contrast of the liver and tumors in the CT images. The effects of the preprocessing are illustrated in Fig 4A. After CT image preprocessing, irrelevant slices to liver segmentation were removed to reduce interference from other abdominal organs during the liver segmentation process. Simultaneously, the grayscale display range was redefined, biasing the image towards higher density values, thus highlighting the density differences in liver tissue and further enhancing image clarity. This processing enables clearer visualization of structures with different densities and facilitates differentiation and identification of low-density abnormal tissues such as tumors, thereby improving the network’s capability for feature extraction of the liver and tumors.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 4. Data processing results. before segmentation.

The images were subjected to windowing to enhance the contrast between the tumor and surrounding tissues, making the target region more prominent. After segmentation, a connected component analysis was performed to remove noise and undesired regions that did not conform to the expected shape, resulting in a more continuous, complete, and accurate segmentation. (a) Data pre-processing (b) Data post-processing.

https://doi.org/10.1371/journal.pone.0312105.g004

There is only one liver in the human body. If there are two or more regions in the liver segmentation result, it indicates that false positive regions may be present in the segmentation results. To address this, we performed connected component analysis on the liver segmentation result, dividing it into multiple connected components. Assuming that the largest connected component has x voxels, if there are connected components with voxel counts less than rate×x in the segmentation result, they are removed. In this study, we set the rate to 0.3. The post-processing effect is illustrated in Fig 4B. After connected component analysis and voxel removal operations, significant changes were observed in the data content of the liver segmentation result. By removing small volume connected components, the influence of noise voxels was reduced, leading to a significant decrease in noise levels in the data and an enhancement of the saliency and recognizability of the liver structures. Additionally, the overall distribution of voxel values became more concentrated around liver tissue representation, reinforcing the representation of liver structures in the data and improving the accuracy of the segmentation result.

Evaluation metrics

We evaluated the experimental results and the model using multiple metrics, including Dice score [49], 95 Hausdorff distance (HD95) [50], average surface distance (ASD) [51], and Jaccard Index(Jaccard) [52].

In medical image segmentation, the Dice coefficient is used to measure the similarity between the segmentation result and the ground truth. Let S represent the segmentation result and G represent the ground truth. The formula for the Dice coefficient is given by Eq (7): (7)

The Dice coefficient calculated using this formula ranges from 0 to 1, where 1 indicates perfect overlap, meaning the segmentation result is identical to the ground truth, and 0 indicates no overlap. The Dice coefficient is one of the standard metrics for evaluating the performance of segmentation algorithms.

Hausdorff distance is used to measure the maximum error between the segmentation result and the ground truth. It reflects the extreme differences between the segmentation result and the ground truth, making it particularly suitable for measuring boundary accuracy. Let S and G denote the segmentation result and the ground truth, respectively. The definition of Hausdorff distance is given by Eq (8): (8)

To avoid unreasonable distances caused by certain points being too far from the main cluster, this study calculates the distance of the furthest 95% of pixels as a metric (HD95) to ensure numerical stability. The distance values are bounded by the 95th percentile, making the distance measurement more robust and less sensitive to outliers. A smaller HD95 indicates that the two sets are closer to each other, indicating a better segmentation result.

ASD represents the average distance between the segmented results and the ground truth edges. Compared to HD95, ASD provides a more comprehensive assessment of segmentation accuracy. A smaller ASD indicates a closer match between the segmentation results and the ground truth labels. Let S represent the segmentation result and G represent the ground truth. The formula for ASD is given by Eq (9): (9)

However, Hausdorff focuses on the distance between the most dissimilar point pairs, whereas ASD is based on the average distance between all points. Therefore, ASD can provide a more comprehensive assessment of the accuracy of segmentation results.

The Jaccard measures the overlap between the predicted and ground truth values. A larger Jaccard indicates a higher degree of overlap between the two sets, indicating a smaller difference between the segmentation results and the ground truth annotations. Let S represent the segmentation result and G represent the ground truth. The formula for the Jaccard index is given by Eq (10): (10)

Implementation details

We implemented MA-cGAN on PyTorch 1.8.0. All models were trained on an RTX 3090 with 24GB of memory. We used AdamW [53] as the optimizer with an initial learning rate of 2e-4. During training, we monitored the average Dice coefficient of the liver and tumor on the validation set as the evaluation metric. If the Dice coefficient did not improve after 30 epochs, we halved the learning rate and continued training. Training was stopped either after 60 epochs without improvement or when the number of epochs reached 300. The input images were resized to a width and height resolution of 256×256, and the batch size was set to 2. After training, we visualized the experimental segmentation results using ITK-SNAP for comparison.

Quantitative evaluations

To evaluate the segmentation performance of the MA-cGAN method, we conducted comparisons on the LiTS dataset’s test set. The comparison of segmentation results from multiple models is shown in Table 1.

Download:

• PNG
  larger image
• TIFF
  original image

Table 1. Quantitative segmentation results on the LiTS dataset.

https://doi.org/10.1371/journal.pone.0312105.t001

It can be observed that the segmentation performance of the MA-cGAN network surpasses that of previously proposed networks. After data preprocessing and segmentation post-processing, the liver model achieves a Dice score of 96.95, and the tumor model achieves a Dice score of 78.53. These scores are 0.41 (UTNet [18]) and 1.64 (SwinUnetr [22]) higher than the second-best segmentation methods, respectively. MA-cGAN achieves excellent segmentation performance while utilizing a smaller parameter count. This indicates that our proposed multi-axis attention mechanism can better focus on small regions of interest and exhibit good segmentation performance on the irregular boundaries of the liver and tumors.

We also visualized the changes in average Tversky Loss on the validation set during the training process of 3D-Unet [34], AttnUnet [11], TransBTS [19], SwinUnetr [22], and MA-cGAN. The descending curves are illustrated in Fig 5.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 5. Comparison of Tversky Loss on the validation set during the training process of different models.

After convergence, MA-cGAN exhibits a lower Tversky Loss, demonstrating its excellent segmentation performance for liver segmentation tasks.

https://doi.org/10.1371/journal.pone.0312105.g005

It can be observed that MA-cGAN exhibits a relatively slower rate of Tversky Loss reduction. This is due to the unique nature of the generative adversarial network, where the optimization objectives of the generator and discriminator are inherently conflicting. The generator aims to generate samples that resemble real data distributions, while the discriminator aims to differentiate between the generated samples and real samples. This adversarial training process makes the training of MA-cGAN more challenging compared to traditional supervised learning models. However, with increasing training iterations, the segmentation accuracy of MA-cGAN gradually improves, and its Tversky Loss becomes lower than the other compared models. This indicates that although MA-cGAN requires a longer convergence time during training, it achieves higher segmentation accuracy in medical segmentation tasks, minimizing the occurrence of false negatives and reducing the potential for diagnostic errors by medical professionals. This improvement enhances the accuracy and reliability of medical diagnosis.

In order to better evaluate the performance of each model in liver and tumor segmentation, we created box plots for comparison, as shown in Fig 6A and 6B.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 6. The box plots of liver and tumor segmentation results.

The box plots demonstrate that MA-cGAN exhibits better segmentation accuracy and stability compared to other networks. (a) Boxplot of Dice Liver Values (b) Boxplot of Dice Tumor Values.

https://doi.org/10.1371/journal.pone.0312105.g006

These two figures respectively illustrate the Dice coefficient values for different network models in liver and tumor segmentation tasks. Fig 6A shows the box plot of Dice coefficients for liver segmentation, while Fig 6B presents the box plot for tumor segmentation. The figures include performances of 3D-Unet, AttUnet, TransUnet, UTNet, TransBTS, Unetr, SwinUnetr, and our proposed MA-cGAN model.

In the liver segmentation task, the MA-cGAN model demonstrates the highest distribution of Dice coefficients, with a median close to 0.98, indicating its superior performance in liver segmentation. Additionally, the box for the MA-cGAN model is taller, suggesting a high stability and a narrow range of fluctuation in its predictions, further confirming its consistency and reliability across different samples. In contrast, other models exhibit some shortcomings. The Dice coefficient distributions for the 3D-Unet and AttUnet models are lower, with some outliers, indicating potential underperformance in certain cases. The performances of the TransUnet, UTNet, and TransBTS models are relatively similar, with means and medians around 0.95. Although these models show stable performance, they slightly lag behind MA-cGAN in terms of highest values. The Unetr and SwinUnetr models perform slightly worse than TransUnet, UTNet, and TransBTS, and the presence of outliers in the 3D-Unet and TransBTS models suggests that their prediction results may deviate from the normal range in specific situations.

In the tumor segmentation task, the MA-cGAN model again exhibits excellent performance, with a median Dice coefficient close to 0.85 and a tall box, indicating stability and consistency in tumor segmentation. In comparison, the other models show significant variability in performance. The Dice coefficient distributions for the 3D-Unet and AttUnet models are lower, with multiple outliers, highlighting instability in tumor segmentation. The performances of the TransUnet and UTNet models are relatively similar, with means and medians around 0.75. Although these models are stable, they fall short of MA-cGAN in terms of the highest values. The TransBTS, Unetr, and SwinUnetr models perform slightly worse than TransUnet and UTNet, with multiple outliers, indicating that their predictions may deviate from the normal range in certain situations.

In summary, the MA-cGAN model demonstrates outstanding performance in both liver and tumor segmentation tasks, achieving the highest Dice coefficients and the smallest range of fluctuations, proving the superiority of our proposed network in these tasks. Although other models also perform well, they are slightly inferior to the MA-cGAN model in terms of stability and maximum values, and exhibit certain deficiencies.

Qualitative results

To further validate the effectiveness of the MA-cGAN network in liver and tumor segmentation tasks, we visualized the segmentation results on the LiTS dataset. The tumor models in the dataset contain both large target regions and small, blurry regions with irregular surfaces, which pose high demands on the model’s ability to segment multi-scale targets. By visually comparing the segmentation results, we can better evaluate the performance of each model in liver and tumor segmentation tasks. The comparisons are shown in Fig 7.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 7. Visual comparisons of the segmentation results were performed on the LiTS dataset with other methods.

MA-cGAN demonstrated remarkable performance in both tumor and liver segmentation tasks, particularly excelling in handling small targets and irregular regions, thereby exhibiting more accurate segmentation performance.

https://doi.org/10.1371/journal.pone.0312105.g007

The comparative results reveal that the liver models segmented by 3D-Unet [34] and AttnUnet [11] exhibit a relatively rough surface, and there are more false positive regions in the tumor models. In contrast, the liver segmentation results obtained by MA-cGAN exhibit a smoother surface, fewer false positive regions in the tumor models, and finer edges, which are closer to the real samples. On the other hand, networks such as UTNet [18], TransBTS [19] and Unetr [20], achieve good segmentation performance, but the introduction of Transformer modules significantly increases the number of parameters and computational requirements, demanding more computational resources. MA-cGAN, on the other hand, demonstrates excellent performance while utilizing fewer parameters and computational resources, thereby proving its suitability for medical-assisted treatments.

We further compared the segmentation results using two-dimensional cross-sectional views to better observe any false positive or false negative regions present in the segmentation results. The comparative results are shown in Fig 8.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 8. Visual comparison of 2D section.

Compared to other networks. The segmentation results of MA-cGAN contain fewer false positive and false negative regions, resulting in higher segmentation accuracy.

https://doi.org/10.1371/journal.pone.0312105.g008

Fig 8 illustrates the comparative results of our proposed MA-cGAN model against several common models for liver and tumor segmentation tasks. The first two rows depict liver segmentation results, while the last two rows show tumor segmentation results. The white areas in the figure represent the Ground Truth (GT) labels, and the red areas denote the segmentation results of each model.

In liver segmentation, the MA-cGAN model exhibits results closest to the GT, demonstrating its superior performance in accurately segmenting the liver. In contrast, other models show varying degrees of error in their segmentation results. For instance, 3D-Unet and AttUnet models exhibit noticeable false negative regions along the liver edges (white areas not covered by red), indicating inaccuracies in certain edge recognition scenarios. Although TransUnet, UTNet, and TransBTS models approximate the overall shape of GT, they still exhibit some false positive areas (red areas extending beyond white), suggesting potential over-segmentation in specific cases. Unetr and SwinUnetr models show significant errors in certain regions, particularly along the liver edges. The MA-cGAN model excels in liver segmentation by accurately capturing both the edges and detailed parts of the liver, significantly reducing false positive and false negative regions.

In tumor segmentation, the MA-cGAN model similarly performs exceptionally well, with segmentation results closely matching GT, accurately identifying tumor shapes and positions. Conversely, other models show considerable discrepancies in tumor segmentation tasks. 3D-Unet and AttUnet models display numerous false negative areas in tumor segmentation, especially along tumor edges, indicating challenges in accurate tumor recognition. While TransUnet and UTNet models achieve closer results to GT, they still exhibit false positive areas in certain regions, indicating potential over-segmentation in specific scenarios. TransBTS, Unetr, and SwinUnetr models demonstrate evident false positive and false negative occurrences in specific areas. The MA-cGAN model excels in tumor segmentation by accurately identifying both the edges and detailed parts of tumors, significantly reducing false positive and false negative regions.

In summary, the MA-cGAN model demonstrates outstanding performance in both liver and tumor segmentation tasks, accurately capturing the edges and detailed parts of the target regions while significantly reducing false positive and false negative regions. This validates the superiority of our proposed network for these two tasks. While other models also perform well, they exhibit slightly inferior stability and accuracy compared to the MA-cGAN model, with some shortcomings.

Ablation

Effect of the hyperparameters in Tversky loss.

The sizes of α and β values in Tversky Loss have a certain impact on the segmentation performance of the MA-cGAN network. The values of these two parameters are typically in the range of 0, 1, and the sum of a and b is equal to 1. In this study, we conducted ablation experiments to investigate the effects of different α and β values on the performance of the MA-cGAN network. The experimental results are shown in Table 2.

Download:

• PNG
  larger image
• TIFF
  original image

Table 2. The impact of α and β on the performance of MA-cGAN.

https://doi.org/10.1371/journal.pone.0312105.t002

If α is too large, the model will prioritize accuracy over class balance, and it may tend to predict all regions as healthy tissues to maximize accuracy. However, this can result in a very low recall, potentially missing some true tumor lesions and leading to false negatives. These lesions may require treatment in the early stages, and missing the optimal treatment window could worsen the condition and even jeopardize the patient’s life. Additionally, this can cause a performance drop when dealing with minority classes, leading to a decrease in overall image segmentation performance.

On the contrary, if β is too large, it can lead to false positives. Although this may result in some unnecessary treatments, the impact is relatively smaller since these erroneous judgments can be corrected through further examination and diagnosis.

In medical prediction tasks, false negatives may pose greater risks and harm to patients compared to false positives. After conducting multiple ablation experiments, this study sets the value of α as 0.7 and β as 0.3 to achieve a balanced trade-off between precision and recall, adapting to the liver and tumor segmentation tasks.

Effect of the cGAN and MA module

To further validate the effectiveness of the proposed Conditional Generative Adversarial Network (cGAN) and Multi-Axis Attention (MA) mechanism, we used the U-Net segmentation network mentioned in the text as a baseline. We incorporated cGAN and MA separately into the baseline, and also conducted an ablation experiment with MA-cGAN to evaluate the performance improvement of cGAN and MA on liver and tumor segmentation. The quantitative segmentation results of the experiments are compared and presented in Table 3.

Download:

• PNG
  larger image
• TIFF
  original image

Table 3. Quantitative segmentation results of cGAN and MA modules.

https://doi.org/10.1371/journal.pone.0312105.t003

After each training epoch, we also computed the Tversky Loss between the segmentation results on the validation set and the corresponding ground truth labels. summed these values, and took their average. The trend of average Tversky Loss on the validation set during the training process is shown in Fig 9.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 9. Comparison of validation set Tversky loss during training with cGAN and MA Modules in the Baseline.

The downward trend of the Tversky Loss can be used to compare the segmentation accuracy and convergence speed of the network.

https://doi.org/10.1371/journal.pone.0312105.g009

It can be observed that although the baseline exhibits a rapid decrease in Tversky Loss during training, its performance capacity is limited, as it fails to adequately capture the data distribution and patterns. This limitation can result in the model’s inability to accurately learn the features of real data during the training process. Therefore, even after training, the baseline still exhibits relatively high Tversky Loss.

After introducing the cGAN module, the adversarial training between the generator and discriminator allows the generator to learn more intricate liver and tumor image features, resulting in more accurate segmentation boundaries and improved segmentation accuracy. However, this adversarial training process requires more iterations and training time to converge. Therefore, during the training process, the Tversky Loss of the cGAN model decreases at a slower rate compared to the Baseline model.

After incorporating the MA module into the baseline, MA can extract two-dimensional projection feature information along three axes (X, Y, Z) and compress and integrate the multi-axis information, thereby reducing the dimensionality of the features and enhancing the focus on key feature information. Furthermore, by adaptively adjusting the weights between different axes, MA can better capture the correlations between different axes and modalities, further improving the accuracy of liver and tumor segmentation.

GAN and MA have improved the segmentation performance of the network for the liver and tumor. The visual comparison of the segmentation results is shown in Fig 10.

Download:

• PNG
  larger image
• TIFF
  original image

Fig 10. Visualization comparison of segmentation results from GAN and MA ablation experiments.

Compared to the Baseline, the inclusion of GAN and MA both improved the segmentation performance of the network, leading to higher accuracy in the segmentation results for both liver and tumor.

https://doi.org/10.1371/journal.pone.0312105.g010

Compared to the ground truth labels, the segmentation results of the Baseline model only roughly preserve the segmentation shapes, and the accuracy in capturing edge details is relatively low. After incorporating GAN, the generator can generate images that are similar to the real liver segmentation results, closer to the actual liver shape and structure. The discriminator network of GAN provides feedback on the differences between the generated results and the real results, aiding the generator in learning more accurate liver boundaries and tumor regions. With the addition of MA, the network can capture multi-axis feature information for fusion, resulting in a greater focus on important regions within the liver or regions with pathologies, leading to higher segmentation accuracy in these areas. MA-cGAN combines GAN and MA, retaining the high similarity of GAN and the high attention of MA, resulting in liver and tumor segmentation results that are closer to the ground truth, with clearer edge details and improved segmentation performance.