FSOU-Net: Feature supplement and optimization U-Net for 2D medical image segmentation

4.1.1 Datasets

In order to evaluate the performance and robustness of the proposed method, three public datasets of medical images from different domains were selected for the experiments: the RITE dataset [45], the Kvasir-SEG dataset [44], and the GlaS dataset [43]. The experimental results show that the proposed method achieves better segmentation results on all three datasets, with different degrees of improvement compared to the current more advanced CNN networks. In the next sections, we first present the model performance evaluation metrics used for the experiments and the details of the experimental implementation; then we describe the significance of the datasets for clinical diagnosis in practice and show the experimental results on each dataset. Finally, to show that the two proposed methods can work together, we perform ablation experiments on the GlaS dataset for the proposed model. Details of the datasets, the number of training and testing samples used and their availability are shown in Table 1.

Image segmentation is actually a pixel-level classification task. Statistical analysis of the classification correctness of each pixel value of the segmentation result allows to evaluate the segmentation performance of the model. This is also a common way of evaluating medical image segmentation models. The common metrics of semantic segmentation, IoU (Intersection over Union), Dice (Dice coefficient), Se (Sensitivity), and Acc (Accuracy), are used as the evaluation metrics of the model in the experiments. Among them, IoU and Dice metrics are the main evaluation metrics for semantic segmentation. Their formulas are as follows:

𝐼𝑜𝑈 = T ⁢ P F ⁢ P + T ⁢ P + F ⁢ N 𝐷𝑖𝑐𝑒 = 2 ⁢ T ⁢ P F ⁢ P + 2 ⁢ T ⁢ P + F ⁢ N S ⁢ e = T ⁢ P T ⁢ P + F ⁢ N 𝐴𝑐𝑐 = T ⁢ P + T ⁢ N T ⁢ P + T ⁢ N + F ⁢ P + F ⁢ N

Here, for each pixel point of the segmentation result TP is true positive, which is predicted to be a positive sample and also a positive sample in ground truth; TN is true negative, which is predicted to be a negative sample and also a negative sample in ground truth; FP is false positive, which is predicted to be a positive sample but a negative sample in ground truth; FN is false negative, which is predicted to be a negative sample but a positive sample in ground truth. Theoretically, the higher the values of IoU, Dice, Se, and Acc, the higher the efficiency of the algorithm segmentation.

4.1.3 Loss function

The loss function in the experiment uses Binary Cross Entropy Dice loss (BCEDiceLoss), which is a loss function that integrates Binary Cross Entropy loss (BCELoss) and Dice loss (DiceLoss). BCELoss is a commonly used loss function for binary classification. Assuming that P i denotes the prediction result of the model, G i denotes the ground truth of the image, G i , j denotes ground truth the pixel value at position i th , j th . P i , j denotes the pixel value at position i th , j th of the model output image, h denotes the height of the image, and w denotes the width of the image, BCELoss can be described as follows.

L 𝐵𝐶𝐸 = 𝐵𝐶𝐸𝐿𝑜𝑠𝑠 ⁢ ( P i , G i ) = - ∑ i = 1 w ∑ j = 1 h [ ( G i , j ⁢ log ⁡ P i , j ) + ( 1 - G i , j ) ⁢ log ⁡ ( 1 - P i , j ) ]

V-Net [46] proposed DiceLoss, which aims to cope with the scenario of strong imbalance between positive and negative samples in semantic segmentation. The calculation of the Dice metric has been introduced in the previous section; in fact, it is intuitive to understand that the dice coefficient calculates the similarity between the target region of the segmentation output by the model and the target region in ground truth. Assuming that I denotes the intersection of the target region in the predicted image and the target region in ground truth, and U denotes the union of the target region in the predicted image and the target region in ground truth, the DiceLoss can be expressed as follows:

L 𝑑𝑖𝑐𝑒 = 𝐷𝑖𝑐𝑒𝐿𝑜𝑠𝑠 ⁢ ( I , U ) = 1 - I + ε U - I + ε

where ε is the smoothing factor. BCEDiceLoss can be expressed as follows:

L 𝐵𝐶𝐸𝐷𝑖𝑐𝑒 = α ⁢ L 𝐵𝐶𝐸 + β ⁢ L 𝑑𝑖𝑐𝑒

where α and β are the weights of the two losses, respectively, and are not zero. In the experiment, we set α to 0.5 and β to 1.

4.1.4 Implementation details

The experiments were trained on NVIDIA GeForce RTX 3060, a GPU with 12 GB of memory. The implementation of this experiment is based on PyTorch. For the three public datasets, we use Adam as the model optimizer with weight decay set to 1e-4. The learning rate is set to 1e-3 during training, and we use simultaneous BN to train the model. The batch size is set to 8 for the other two datasets, except for the Kvasir-SEG dataset which the batch size is set to 16. The optimal parameters for training are not saved until the loss in the validation set is no longer decreasing and the test data are predicted.

4.2.1 Comparison on RITE

The RITE (Retinal Images vessel Tree Extraction) is a database that enables comparative studies on segmentation or classification of arteries and veins on retinal fundus images, which is established based on the public available DRIVE database (Digital Retinal Images for Vessel Extraction). The Angle, width and length of retinal blood vessels can be extracted from the segmented image. This information can be used to diagnose various ophthalmic diseases, such as atherosclerotic choroidal neovascularization [18].

The RITE dataset [45] is a subset of the DRIVE dataset. It contains 40 images, and the original dataset divides them into a training dataset and a test dataset, each containing 20 images. On the basis of the original training data set, we divide the training data set and validation data set, and the partition results are shown in Table 1. We extended the training data by using horizontal and vertical flips during training. The results of the experiments are shown in Table 2.

OPEN IN VIEWER

The RITE dataset is characterized by extremely fine branches in the target region, which also leads to the generally low segmentation accuracy of most models. the multiple down-sampling design used by U-Net may lose fine-grained semantic information, which is one of the drawbacks of the U-Net model, and for small targets, it tends to fail in recognition. The proposed model can supplement the semantic information of fine-grained targets in shallow semantic features, and thus obtain better segmentation performance. As we expect, it can be seen from the table that the proposed model obtains better segmentation performance compared to the pre-improved U-Net model, with the main metrics IoU and Dice being 0.9% and 0.7% higher than U-Net, respectively. the Se and Acc metrics are 2% and 0.04% higher, respectively.

Figure 5 visualizes the segmentation performance of the advanced model in the field of semantic segmentation for three datasets. For image A of the RITE dataset, the proposed model can segment finer vascular branches compared to the other networks, which we believe is due to the shallow feature supplement structure in FSOU-Net. In the segmentation results of image B, a trunk branch of the vascular network is completely segmented by the proposed FSOU-Net, while this trunk can be identified but not completely segmented in other advanced models. The proposed model can be better identified and has better segmentation performance compared with other advanced networks.

Colorectal cancer is a common type of cancer in humans, and polyps are the precursor of colorectal cancer. The risk of polyps increases with age. Half of people who get colonoscopies at age 50 are found to have polyps. Early detection of the disease can greatly improve the survival rate of colon cancer patients, and polyp detection is particularly important [44].

Kvasir-SEG is an open-access dataset of gastrointestinal polyp images and corresponding segmentation masks, manually annotated by a medical doctor and then validated by an experienced gastroenterologist. The dataset has a total of 1000 original images [44]. We divided them into training set, validation set and test set by 8:1:1 after random sorting, and did comparison experiments between the proposed model and some advanced models on this dataset, and the experimental results are shown in the Table 3.

The boundaries of the target regions in the Kvasir-SEG dataset are difficult to determine, which places higher demands on the model. For image C of Fig. 5, the models in the experiment all detect the location of the target region, but FSOU-Net outperforms the other networks in segmenting the contours of the target region. Image D has four target regions that need to be segmented, and only FSOU-Net segments them all. The disadvantage is that two target regions are connected together, which is an area for improvement.

As can be seen from Table 3, our method has the highest performance in four metrics, IOU reaches 76.3%, Dice reaches 84.1%, Se reaches 86.5%, and Acc reaches 94.63%. Compared to the original U-Net before improvement, the proposed method improves the IOU metric by 2.8%, Dice by 2.3%, Se by 2.5%, and Acc by 0.14%.

Gland is an important histological structure in human body. It exists in most organs and is the main structure that secretes carbohydrates and proteins. Malignant tumors arising from glandular epithelium are the most prevalent form of cancer. The morphology of glands is important to evaluate the degree of malignancy of glands [43]. Accurate gland segmentation is essential to obtain reliable morphological statistics.

The GlaS datasets used in the experiments are provided by a team of pathologists at the University Hospitals Coventry and Warwickshire, UK. The GlaS dataset was segmented into 85 training images (37 benign and 48 malignant) and 80 test images (37 benign and 43 malignant). On the basis of the original training data set, we divide the training data set and validation data set, and the partition results are shown in Table 1. We extended the training data by using horizontal and vertical flips during training. The experimental results are shown in Table 4.

Compared with the other two datasets, the GlaS dataset has relatively more segmentation targets, and the number of segmentation targets in the images varies greatly, and this dataset requires the model to have better segmentation of the boundaries of the target regions. As can be seen from the visualization in Fig. 5, the proposed model has a better segmentation effect on the boundary information compared with the original U-Net due to the supplementation of fine-grained semantic information in the shallow features and a certain degree of optimization of the deep semantic features. The experimental results show that the main metrics IOU and Dice are improved by 0.3% and 0.2, and Se and Acc are improved by 1.1% and 0.05%, respectively.

To verify that the two proposed methods can work together, we performed ablation experiments on the GlaS data set for the shallow feature supplement method and the deep feature optimization method. The experimental process adopts 5-fold cross validation, and the average value of the 5-fold result is taken as the final experimental result. The ablation study results are shown in Table 5.

OPEN IN VIEWER

From the table, it can be seen that with the addition of shallow feature supplementation alone or deep feature optimization both can improve the segmentation effect. Adding the shallow feature supplementation method alone improves the main metrics IOU and Dice by 0.2% and 0.16%, respectively, while adding the deep feature optimization method alone improves the main metrics IOU and Dice by 0.24% and 0.19%, respectively. The simultaneous addition of the two proposed methods has a better segmentation effect than the single addition of one, and the IOU and Dice metrics are improved by 0.34% and 0.24%, respectively. This indicates that the two proposed methods can work in concert to jointly improve the segmentation effect of the network.

As shown in Fig. 6, we compare the training process of the three best-performing models on each of the three datasets, and we can see that the proposed model always maintains the highest IOU metrics, while the convergence rate of IOU metrics also remains high among several best models.