A Multi-Organ Segmentation Network Based on Densely Connected RL-Unet

Abstract

The convolutional neural network (CNN) has been widely applied in medical image segmentation due to its outstanding nonlinear expression ability. However, applications of CNN are often limited by the receptive field, preventing it from modeling global dependencies. The recently proposed transformer architecture, which uses a self-attention mechanism to model global context relationships, has achieved promising results. Swin-Unet is a Unet-like simple transformer semantic segmentation network that combines the dominant feature of both the transformer and Unet. Even so, Swin-Unet has some limitations, such as only learning single-scale contextual features, and it lacks inductive bias and effective multi-scale feature selection for processing local information. To solve these problems, the Residual Local induction bias-Unet (RL-Unet) algorithm is proposed in this paper. First, the algorithm introduces a local induction bias module into the RLSwin-Transformer module and changes the multi-layer perceptron (MLP) into a residual multi-layer perceptron (Res-MLP) module to model local and remote dependencies more effectively and reduce feature loss. Second, a new densely connected double up-sampling module is designed, which can further integrate multi-scale features and improve the segmentation accuracy of the target region. Third, a novel loss function is proposed that can significantly enhance the performance of multiple scales segmentation and the segmentation results for small targets. Finally, experiments were conducted using four datasets: Synapse, BraTS2021, ACDC, and BUSI. The results show that the performance of RL-Unet is better than that of Unet, Swin-Unet, R2U-Net, Attention-Unet, and other algorithms. Compared with them, RL-Unet produces significantly a lower Hausdorff Distance at 95% threshold (HD95) and comparable Dice Similarity Coefficient (DSC) results. Additionally, it exhibits higher accuracy in segmenting small targets.

1. Introduction

With the rapid development of deep learning and computer vision technology, medical image segmentation has been widely applied in assisted diagnosis and intelligent medical treatment, considerably enhancing the effectiveness and accuracy of diagnosis [1]. It can extract meaningful semantic information from original medical image data, identify pixels of diseased organs, and obtain information features of these diseased parts [2,3,4]. Early medical image segmentation methods, such as template matching technology, edge detection, and traditional machine learning technology, have achieved certain results; however, due to the difficulty of feature representation, medical image segmentation is still a challenging problem. Recently, good progress has been made in medical image segmentation based on deep learning, which has aroused a lot of attention. Among them, the method based on Convolutional Neural Network (CNN) is widely used in medical image segmentation and has achieved good results.

Based on the AlexNet [5] network structure, the Fully Convolutional Network (FCN) [6] converts the full connection layer into the convolutional layer, and increases the dimension of the feature map by means of up-sampling, so as to realize the end-to-end semantic segmentation at the pixel level. However, this method has limited capability in capturing fine-grained features. This creates a significant challenge for tasks that require precise segmentation (e.g., medical image segmentation). Ronneberger et al. [7] proposed a landmark U-net network architecture on the basis of FCN. The encoding stage and decoding stage that correspond to each other in the internal components form the overall architecture of the network. In the encoding process, image features are extracted from the down-sampled images. In the process of decoding, images are up-sampled to gradually recover the size of the image [8]. The pooling operation and continuous convolution kernel in the encoding stage result in the loss of some feature information of the image. However, in the decoding stage, the up-sampled feature map merges with the front-end information from the skip connections, thereby enriching the image’s detailed features. This U-shaped structure makes U-Net achieve a major breakthrough in various image segmentation work and produces many related research algorithms. For example, Unet3+ [9] and Attention U-net [10] enable the decoder to integrate more comprehensive feature information by adding the feature information of different scales or introducing attention mechanisms. The V-Net [11] network integrates 3D convolution and U-Net architecture, and proposes a volumetric full-convolution 3D neural network segmentation method, which has achieved good results on prostate MRI images.

Although the CNN-based method has good local feature representation ability, due to the inherent restrictions of convolution transportation, it is challenging to learn global information and contextual semantic information interaction, so U-Net has certain limitations in remote relationship modeling and complex prediction tasks [12,13]. Some work has adopted the deep convolutional layer to expand the receptive field, introduce an attention mechanism into network structure, or adopt a pyramid network [14,15,16,17], etc., which may lead to further increases in model complexity and processing cost, and still need to further improve the ability of global information interactive capture.

Motivated by the prosperous implementation of Transformer in work related to NLP, the ViT [18] algorithm makes only minor changes to the whole image classification process in an effort to directly apply the typical Transformer structure to pictures. The entire image will be broken up into smaller image blocks as part of the ViT algorithm, and the linear embedding sequence of these smaller image blocks will then be supplied to the net as the input of Transformer. When compared to other convolution-based algorithms, it reaches state-of-the-art (SOTA) performance. With the deepening of the research, the combination of CNN and Transformer has achieved good results. Chen et al. [19] offered the Trans-Unet method in 2021, which generates a feature graph using CNN as a tool for extracting features and feeds it into Transformer to extract global data, and uses a cascading upper sampler to ensure accurate prediction. This algorithm has achieved excellent results in multiple medical image datasets. In 2021, Cao et al. proposed the Swin-Unet [20] algorithm based on Transformer to be applied to the medical image segmentation task, which significantly decreases the computational burden of the algorithm and achieves high segmentation accuracy. However, because Swin-Unet can only learn single-scale context features in the training process, the amount of computation is too large and the induction bias module for processing local information is lacking [21]. Some subsequent methods, such as nnFormer, MISSFormer, DSTransUnet, TransDeepLab, and DAE-Former [22,23,24,25,26], adopt a pure Transformer architecture. Specifically, transformer blocks are employed on both the encoder and decoder sides to capture more global characteristics and fuse multi-scale information through skip connections.

The Transformer architecture model, distinguished from the local dependency bias of convolutional neural networks due to its excellent global modeling capabilities and ability to capture long-range dependencies, has garnered significant attention from researchers. However, its high computational cost has also deterred some. As a result, in recent years, research in medical image segmentation has shifted towards hybrid Transformer architectures with skip connections.

Zhao et al. argued that segmentation methods based on patch division might disrupt local coherent features and introduce noise. Building on this premise, they proposed a progressive sampling module [27]. Unlike traditional segmentation models, this approach achieved promising results in small object segmentation tasks, although it did not perform as well in large object segmentation tasks. Rahman et al. innovatively combined visual transformers with graph convolution, proposing a graph-based cascaded attention architecture [28] to address the limitations of visual transformers in handling local spatial information correlations. Zhou et al. introduced a 3D Transformer architecture called nnFormer [29] based on multiple attention mechanisms, which balanced long-range and spatial dependencies in the self-attention mechanism. Besides, Huang et al. redefined the skip connections between the encoder and decoder, enhancing global dependencies further through a bridge-like connection method [30].

Inspired by these works, we suggest an RL-Unet algorithm based on Swin-Unet to solve the aforementioned issues. It redesigns the up-sampling module and adds a local inductive bias module to the Swin-Transformer block, resulting in a densely connected double up-sampling module that can fully learn multi-scale information and increase the boundary segmentation accuracy of the algorithm. The suggested method has perfect segmentation accuracy and a strong generalization capability, according to extensive testing on the segmentation datasets for the synapse multiple abdominal organs and BraTS2021 datasets. Our contributions, specifically, can be summed up as follows:

• We add a local induction bias module to the Swin-Transformer block to assist the RLSwin-Transformer module learn local features and remote dependencies, and replace the MLP module with the Res-MLP module to prevent feature loss during transmission, thus increasing segmentation accuracy.
• We design a new double up-sampling module, which includes bilinear up-sampling and expansion convolution with different expansion rates for feature extraction and image restoration.
• A dense connection structure is designed at the decoding end, and an attention mechanism is introduced to restore resolution and generate segmentation map.
• To significantly increase the segmentation accuracy of small targets, a novel loss function is proposed.

The remainder of the paper is organized as follows: In Section 2, we introduce the overall details of the proposed network structure and design the loss function. Section 3 consists of a description of the datasets used, as well as evaluation scenarios and metrics, and discusses the experiments and results. The paper is finally summarized in Section 4.

2. Method

In this section, we will discuss the details of our proposed network architectures for three models: Residual Local Induction Bias-Unet (RL-Unet), RLSwin-Transformer Module, and Densely Connected Double Up-Sampling Modules (DUSM). Furthermore, we introduce a unique loss function that enhances the performance of these models.

2.1. Integral Structure

The RL-Unet is a U-shaped network that is structured with an encoder, decoder, and skip connections. In the process of residual connection, local inductive bias module is added to build RLSwin-Transformer module. Meanwhile, inspired by TFCNs [31] network, the MLP module is changed into Res-MLP module. It can learn local feature information better and realize cross-group information blending. In the decoder stage, a new densely linked double up-sampling module is proposed, which contains two existing up-sampling methods: one is bilinear up-sampling and the other is introduced by Deeplabv3+ [32] ASPP module. Through the combination of the two, multi-scale information can be effectively learned and model generalization ability can be improved. Figure 1 depicts the RL-Unet’s general structure.

Firstly, in the coding part, the medical image data are separated into non-overlapping image blocks of equal size and input into Patch Partition, and the image blocks are embedded and encoded to complete the feature dimension conversion and other operations. Then, the converted image blocks are input into the RLSwin-Transformer module for feature learning, and the down-sampling function is implemented by Patch Merging [33,34,35]. The feature extraction part is divided into four stages. As the network depth increases, the size of the image blocks halves, and the number of channels increases by twice. The output of the four stages is processed by a standard 1 × 1 convolution to obtain the four hierarchical features of the RLSwin-Transformer (ST1, ST2, ST3, and ST4). In the decoding part, a dense connection double up-sampling module is proposed, which can restore the image’s original resolution and complete the final segmentation prediction through cross-scale connection and attention module.

2.2. RLSwin-Transformer Module

In the process of medical image segmentation, the main role of the RLSwin-Transformer module is to replace the traditional convolutional module for feature learning and information extraction. The RLSwin-Transformer module structure is shown in Figure 2. Features are input into the normalization layer (LN) to ensure the stability of data feature distribution; thenk feature information is learned by the window-based multi-head self-attention (W-MSA) module. Finally, W-MSA output is input into LN and residual multi-layer perceptrons (Res-MLP) successively. At the same time, the local induction bias (LIB) module is introduced in the process of residual connection to obtain the final output. The other layer is a similar process [27,36].

After adding local inductive bias module into RLSwin-Transformer module, on the one hand, rich contextual semantic information can be obtained by stacking convolutional layers with different receptive fields without changing scale. On the other hand, RLSwin-Transformer module can calculate adjacent pixels through the convolution operation to obtain local details of edges and corners. The calculation process of RLSwin-Transformer module is summarized as follows:

$$\begin{array}{cc}\hfill & {\widehat{z}}^l=W\_MSA(LN(z^{l-1}))+z^{l-1}\hfill \end{array}$$

(1)

$$\begin{array}{cc}\hfill & z^l=ResMLP(LN(z^l))+LIB({\widehat{z}}^l)\hfill \end{array}$$

(2)

$$\begin{array}{cc}\hfill & {\widehat{z}}^{l+1}=SW\_MSA(LN(z^l))+z^l\hfill \end{array}$$

(3)

$$\begin{array}{cc}\hfill & z^{l+1}=ResMLP(LN(z^{l+1}))+LIB({\widehat{z}}^{l+1})\hfill \end{array}$$

(4)

where l represents the number of index layers, ẑ^l represents the output of W-MSA module, ẑ^l+1 indicates the output of SW-MSA module, and z^l and z^l+1 represent the output of the Res-MLP module. Equations (5) and (6) illustrate the specific structure of Res-MLP clearly.

$$\begin{array}{cc}\hfill & z^{″}=GeLU(W^{\prime }{z}^{\prime }+b^{\prime })\hfill \end{array}$$

(5)

$$\begin{array}{cc}\hfill & ResMLP(z^{\prime })=z^{\prime }+GeLU(W^{″}{z}^{″}+b^{″})\hfill \end{array}$$

(6)

where $z^{\prime }$ represents the output of layernorm; W and b represent weight matrix and bias vector, respectively; and GeLU is the GeLU nonlinear activation function.

The local induction bias module is embedded in the residual connection process of the RLSwin-Transformer module. The right half of Figure 2 provides a detailed description of the precise structure. The Seq2Img layer performs a simple reconstruction function, converting the input one-dimensional sequence into a two-dimensional feature ${\widehat{x}}^l\in {\mathbb{R}}^{C\times H\times W}$.

$$M_{Res},M_{Cov}=CSp{({\widehat{x}}_{i,:,:}^l)}_{(i=C/2)}$$

(7)

where $M_{Res}\in {\mathbb{R}}^{C/2\times H\times W}$ and $M_{Cov}\in {\mathbb{R}}^{C/2\times H\times W}$ represent the feature maps of residual connections and feature extraction branches, respectively. Let $CSp(·)$ be the channel split module. ${\widehat{x}}_{i,:,:}^l$ indicates that input feature map is truncated at the i-th channel in the channel dimension.

Then channel split module evenly divides the channel number of input features into two parts: one is used for feature extraction and the other is used for residual connection. In the feature extraction part, the feature map is input into three Conv layers with the receptive field of 1 × 1, 3 × 3, 1 × 1 to extract local multi-scale features $M_{Lm}\in {\mathbb{R}}^{C/2\times H\times W}$.

The 3 × 3 convolution has been deeply optimized on GPUs, offering higher computational density compared to other convolution kernels [37]. Therefore, we opted for the most efficient combination of 1 × 1, 3 × 3, and 1 × 1 convolutions. This combination was chosen to balance between capturing fine details and maintaining computational efficiency.

$$x^l=CSh{({(M_{Res},M_{Lm})}_{i,:,:})}_{(i=C/2)}$$

(8)

where $x^l\in {\mathbb{R}}^C\times H\times W$ represents the output feature map of the channel shuffle module. Let $CSh(·)$ be the channel shuffle module. ${(M_{Res},M_{Lm})}_{i,:,:}$ represents that the output feature map of the feature extraction part and the residual connection part are spliced in the channel dimension.

Afterwards, the features are spliced on the channels, and the channels of the stacked feature graphs are reordered through the channel shuffle module to realize the feature fusion among the groups. Finally, the Img2Seq layer converts the two-dimensional images with multi-scale features into one-dimensional sequences and incorporates them into Res-MLP to enhance the RLSwin-Transformer module to learn local features and remote dependencies.

2.3. Densely Connected Double Up-Sampling Modules (DUSM)

To solve the problem that the size or shape of the target organ or lesion varies from patient to patient in medical images, a new densely connected double up-sampling module is proposed in this paper. In particular, based on previous studies on linear attention mechanisms [38], joint spatial attention (JSA) and joint channel attention (JCA) are enhanced to improve the spatial and channel relationships of semantic information. Additionally, down-sampling connection and double up-sampling connection are used to further integrate multi-scale features and improve feature representation capability. As shown in Figure 3. DUSM has four multi-scale features (DM1, DM2, DM3, DM4) through cross-scale connections (i.e., down-sampled connection and up-sampled connection) and attention modules.

Densely connected double up-sampling module (DUSM): Input is divided into two parts. One part is further extracted by expansion convolution with different expansion rates; then, up-sampling is performed by transposed convolution with step size of 2. In the other part, the number of channels is changed by 1 × 1 convolution first; then, bilinear up-sampling is carried out through ReLU activation function. Finally, the output from the two portions is combined on one channel, cutting the total number of channels in half. It can be defined as follows:

$$\begin{array}{cc}\hfill & U_i^j(X)=U_i^j{(X)}_r+U_i^j{(X)}_l\hfill \end{array}$$

(9)

$$\begin{array}{cc}\hfill & U_i^j{(X)}_l=f_{\alpha }(f_{\theta }(f_u^1(X))+f_u^6(X)+f_u^{12}(X)+f_u^{18}(X))\hfill \end{array}$$

(10)

$$\begin{array}{cc}\hfill & U_i^j{(X)}_r=f_{\theta }(f_{\delta }(f_{\sigma }(f_u^1(X))))\hfill \end{array}$$

(11)

where X is the input vector, $f_u$ is an expansion convolution with different expansion rates, $f_{\theta }$ represents a 1 × 1 convolution with step 1, $f_{\alpha }$ is a transposed convolution with step 2, $f_{\sigma }$ is a ReLU activation function, $f_{\delta }$ means bilinear up-sampling, and + indicates channel fusion.

Down-sampling connection: As shown in Figure 3b, output obtained at each stage of encoder end is down-sampled by means of convolution so that it can be fused to each corresponding DM module of decoder end through skip connection. It can be defined as follows:

$$D_m^n(X)=f_{\sigma }(f_a(X)+f_a(f_b(X)))$$

(12)

where $f_a$ is a 3 × 3 convolution with step 2 and $f_b$ is a 3 × 3 convolution with step 1.

Joint channel attention (JCA) module: The long-distance relationship between channel dimensions is extracted using an upgraded channel attention module, displayed in Figure 4. This process can be described as Equation (13).

$$CA(x)={\displaystyle \frac{{\displaystyle \sum _{c}M(x)+\left(M{(x)}_{c,n}{\left(\frac{M(x)}{∥M(x)∥}\right)}^T\right)\frac{M(x)}{∥M(x)∥}}}{{\displaystyle N+{\left(\frac{M(x)}{∥M(x)∥}\right)}^T\sum _c{\left(\frac{M(x)}{∥M(x)∥}\right)}^T}}}$$

(13)

where M(x) represents the reshaping operation that flattens the space size and N is the number of pixels in the input feature map.

Joint spatial attention (JSA) module: As illustrated in Figure 5, based on the linear attention mechanism, spatial attention is used to describe the long-distance dependence connection in the dimension of space, which can be defined as Equation (14).

$$SA(x)={\displaystyle \frac{{\displaystyle \sum _{c}V(x)+\left(softmax\left(\left(Q{\left(x\right)}^{T}K(x)\right)\right)\right)V(x)}}{{\displaystyle N+Q(x)\sum _c{\left(K(x)\right)}^T}}}$$

(14)

where Q(x), K(x), and V(x) are vector matrices obtained by convolution, $Q\in R^{N\times D}$, $K\in R^{N\times D}$, and $V\in R^{N\times D}$.

Feature polymerization: the four polymerization characteristics (DM1, DM2, DM3, DM4) can be calculated as follows:

$$\begin{array}{cc}\hfill & DM1=U_2^1(DM2)+ST1+U_2^1(JCA(ST3))+U_3^1(DM3)\hfill \end{array}$$

(15)

$$\begin{array}{cc}\hfill & DM2=JCA(ST2)+U_4^2(JSA(ST4))+U_4^2(DM4)\hfill \end{array}$$

(16)

$$\begin{array}{cc}\hfill & DM1=JSA(ST3)+D_2^3(JCA(D_1^2(ST1)))\hfill \end{array}$$

(17)

$$\begin{array}{cc}\hfill & DM4=ST4+D_3^4(JSA(D_2^3(ST2)))\hfill \end{array}$$

(18)

2.4. New Loss Function

Following are some examples of loss functions that are widely applied to medical image segmentation:

$$\begin{array}{cc}\hfill & {\mathcal{L}}_{FL}=-{\displaystyle \frac{1}{N}}\sum _{i=1}^N\sum _{a=1}^A\lambda {\left(1-y_{pred}\right)}^{\gamma }{y}_{true}log(y_{pred})\hfill \end{array}$$

(19)

$$\begin{array}{cc}\hfill & {\mathcal{L}}_{DL}={\displaystyle \frac{1}{A}}\sum _{a=1}^A{\displaystyle \frac{1}{N}}\left(1-{\displaystyle \frac{{\sum }_{i=1}^N{y}_{true}{y}_{pred}}{N}}{y}_{true}^2+\sum _{i=1}^N{y}_{pred}^2\right)\hfill \end{array}$$

(20)

$$\begin{array}{cc}\hfill & {\mathcal{L}}_{CE}=-\sum _{i=1}^N{y}_{true}log{y}_{pred}\hfill \end{array}$$

(21)

where y_true represents the Ground Truth, y_pred represents the predicted segmentation result, N is the size of the training batch, and A is the number of categories separated. The FL loss balance factor is represented by $\lambda$ while the FL loss exponential factor is represented by $\gamma$.

• Focal Loss (FL): Increasing the weight of difficult-to-classify samples in the loss function is the same as using the FL [39] loss function. As a result, the loss of the comparatively simple to anticipate sample is reduced while the loss of the comparatively hard to guess sample is bigger. Consequently, the emphasis is put more on challenging-to-classify samples.
• Dice Loss (DL): A metric function called Dice Loss [40] is used to assess how comparable two samples are. The value is between 0 and 1, and the higher it is, the more related the two are.
• Cross Entropy Loss (CE): CE [41] expresses the distribution relationship between the predicted value and the true value, and the smaller the cross entropy, the more similar the probability distribution between the two.

Several experiments show that using only DL or CE on the improved RL-Unet model is not optimal for the segmentation results of the Synapse abdominal multi-organ segmentation dataset. Additionally, this paper also studies the segmentation of the RL-Unet model on the BraTS2021 dataset. When there is only backdrop and foreground, if several pixels of the little target are incorrect, Dice will vary visibly and have an impact on the training results. As a result, the prediction of the DL loss function for small targets is not reliable. Therefore, a novel loss function is designed in this paper, as shown in Equation (22), and Equation (23) is adjusted on the basis of the original loss function.

$$\begin{array}{cc}\hfill & {\mathcal{L}}_1=\alpha {\mathcal{L}}_{CE}+\beta {\mathcal{L}}_{DL}\hfill \end{array}$$

(22)

$$\begin{array}{cc}\hfill & {\mathcal{L}}_2=4log{\mathcal{L}}_{DL}^2+log{\mathcal{L}}_{FL}^2\hfill \end{array}$$

(23)

where $\alpha$ and $\beta$ are hyperparameters that can be fine-tuned and, here, we choose 0.3 and 0.7, respectively.

3. Experiments

3.1. Evaluation Metrics and Datasets

In this paper, the mean Dice-Similarity Coefficient (DSC) and mean Hausdorff Distance (HD) are utilized as evaluation measures. Both of them describe the similarity measures between two sample sets, but DSC focuses more on the segmentation accuracy of the internal filling part while HD is more sensitive to the segmentation boundary. They are defined as follows:

$$\begin{array}{cc}\hfill & DSC={\displaystyle \frac{2\left|A\cap B\right|}{\left|A\right|+\left|B\right|}}\hfill \end{array}$$

(24)

$$\begin{array}{cc}\hfill & HD(A,B)=max(h(A,B),h(B,A))\hfill \end{array}$$

(25)

$$\begin{array}{cc}\hfill & h(A,B)=\underset{a\in A}{max}\left\{\underset{b\in B}{min}∥a-b∥\right\}\hfill \end{array}$$

(26)

$$\begin{array}{cc}\hfill & h(B,A)=\underset{b\in B}{max}\left\{\underset{a\in A}{min}∥b-a∥\right\}\hfill \end{array}$$

(27)

where A represents the label diagram of the medical image and B represents the prediction graph of algorithm segmentation. Equation (24) is bidirectional HD, and Equations (25) and (26) are unidirectional HD from A to B and from B to A, respectively.

This paper verifies the segmentation effect of RL-Unet on the Synapse abdominal multi-organ segmentation, BraTS2021, ACDC, and BUSI datasets.

Synapse dataset: The abdominal CT scan dataset consists of 30 cases, focusing on the segmentation of various abdominal organs. Every CT scan consists of 85–198 slices of 512 × 512 pixels. We randomly partition the dataset into 18 scans for training (2212). Validation scans were 12 axial slices and 12 scans. We only segmented eight abdominal organs, including the aorta, gallbladder (GB), left kidney (KL), right kidney (KR), liver, pancreas (PC), spleen (SP), and stomach (SM).

BraTS2021 dataset: BraTS2021 includes multi-parameter MRI scans from 2000 patients. MRI scans have four modal data sizes: Flair, t1ce, t1, and t2, each of which is 240 × 240 × 155 shared segmentation labels. The dataset was arbitrarily split into a training set, verification set and test set at a ratio of 8:1:1. Methods of data enhancement were pruning, rotation, flipping, Gaussian noise, contrast transformation, and brightness enhancement.

ACDC dataset: The ACDC dataset comprises 100 patients focusing on the segmentation of the right ventricle cavity, left ventricle myocardium, and left ventricle cavity. The segmentation labels for each case include left ventricle (LV), right ventricle (RV), and myocardium (MYO). We use 70 cases (1930 axial slices) for training: 10 for validation and 20 for testing.

BUSI dataset: The BUSI (Breast Ultrasound Image) dataset, categorized into normal, benign, and malignant classes, serves as a comprehensive resource for both classification and segmentation tasks in the analysis of breast ultrasound images. For training purposes, only the normal and malignant categories from the BUSI dataset are selected.

3.2. Experiment Settings

Synapse dataset and BraTS2021 Dataset. The pre-training model is initialized on the ImageNet dataset with an initial learning rate of 0.05. A clustering learning rate strategy is employed. The maximum training epoch is set to 150, batch size is 32, and the SGD optimizer with a momentum of 0.9 and weight decay of 1 × ${10}^{-4}$ is used for RL-Unet. Many different data enhancement methods are used to expand the data diversity, including rotating the image 30 degrees, 45 degrees, 90 degrees, 120 degrees, and 150 degrees; moving the image (15,15), (−15, −15), (−15, 15), and (15, −15); and using enhancement factors 0.9 and 1.3 to enhance image tone.

ACDC dataset and BUSI dataset. We train each model for a maximum of 150 epochs, with a batch size of 32. For RL-Unet, we set the input resolution to 224 × 224. For data augmentation, we use random flipping and rotation. The combined weighted Cross-entropy (0.3) and DICE (0.7) loss functions are optimized.

3.3. Experiment Results

The parameters of DSC and HD obtained by the segmentation of eight abdominal organs by RL-Unet in the Synapse dataset are compared with the classical segmentation networks ViT, U-Net, Swin-Unet, Attention Unet (AttU-Net), GCASCADE, nnFormer, MISSFormer, and Recurrent Residual Convolutional Neural Network based on U-Net (R2U-Net). The experimental results are shown in Table 1.

Table 1. Segmentation results of different methods on the SYNAPSE dataset.

Methods	Aorta	Gallbladder	Kidney (L)	Kidney (R)	Liver	Pancreas	Spleen	Stomach	DSC	HD
ViT [18]	44.38	39.59	67.46	62.94	89.21	43.14	75.45	69.87	61.50	39.61
U-Net [7]	89.18	65.42	76.86	70.64	93.35	55.37	89.80	76.01	77.03	39.70
Swin-Unet [20]	86.34	63.45	81.13	75.69	93.63	56.83	87.94	73.22	77.28	26.93
AttU-Net [10]	88.59	64.42	81.73	76.77	93.99	63.68	89.56	71.18	78.74	29.54
R2U-Net [42]	88.26	68.97	76.94	71.36	91.86	57.36	87.36	74.88	77.12	32.12
GCASCADE [28]	81.55	79.75	83.06	64.46	62.89	64.43	82.87	89.39	76.05	12.34
nnFormer [29]	92.04	70.17	86.57	86.25	96.84	83.35	90.51	86.83	86.57	10.63
MISSFormer [30]	86.99	68.65	85.21	82.00	94.41	65.67	91.92	80.81	81.96	18.20
RL-Unet (ours)	86.95	67.48	82.68	79.47	93.84	55.97	90.21	77.41	80.13	22.07

Table 1. Segmentation results of different methods on the SYNAPSE dataset.

Methods	Aorta	Gallbladder	Kidney (L)	Kidney (R)	Liver	Pancreas	Spleen	Stomach	DSC	HD
ViT [18]	44.38	39.59	67.46	62.94	89.21	43.14	75.45	69.87	61.50	39.61
U-Net [7]	89.18	65.42	76.86	70.64	93.35	55.37	89.80	76.01	77.03	39.70
Swin-Unet [20]	86.34	63.45	81.13	75.69	93.63	56.83	87.94	73.22	77.28	26.93
AttU-Net [10]	88.59	64.42	81.73	76.77	93.99	63.68	89.56	71.18	78.74	29.54
R2U-Net [42]	88.26	68.97	76.94	71.36	91.86	57.36	87.36	74.88	77.12	32.12
GCASCADE [28]	81.55	79.75	83.06	64.46	62.89	64.43	82.87	89.39	76.05	12.34
nnFormer [29]	92.04	70.17	86.57	86.25	96.84	83.35	90.51	86.83	86.57	10.63
MISSFormer [30]	86.99	68.65	85.21	82.00	94.41	65.67	91.92	80.81	81.96	18.20
RL-Unet (ours)	86.95	67.48	82.68	79.47	93.84	55.97	90.21	77.41	80.13	22.07

Compared with basic ViT and U-Net, DSC and HD parameters of RL-Unet have been greatly improved, which demonstrates the efficiency of the combination of CNN and Transformer. Secondly, compared with the classical medical segmentation algorithm U-Net, the parameters of DSC and HD of RL-Unet increased by 3.1% and 17.63 mm, respectively, indicating the feasibility of the segmentation effect of RL-Unet and its sensitivity to small organs. Compared with Swin-Unet, AttU-Net, and R2U-Net, the parameters of DSC and HD of RL-Unet are still greatly improved, which indicates that RL-Unet pays attention to local information and extracts and integrates multi-scale information in context, effectively improving the segmentation accuracy. Although our proposed algorithm exhibits a slightly lower DSC compared to the latest MISSFormer and nnFormer algorithms, we are unable to precisely reproduce the results of the original authors due to machine limitations and variations in data processing methods. To validate the effectiveness of our proposed algorithm, we conducted experiments on the BraTS2021, ACDC, and BUSI datasets.

As shown in Figure 6, we can see that the boundary of RL-Unet segmentation is clearer and smoother than that of other algorithms. It can be seen from the second line that several algorithms are missing the segmentation of pancreas, and the segmentation boundary of stomach in this dimension is poor, which overlaps with the prediction of pancreas. It can be seen from the three lines that several algorithms have a good segmentation effect and high segmentation accuracy for large targets such as the liver and stomach. Other algorithms resulted missing segmentation and over-segmentation on the pancreas. However, our proposed algorithm RL-Unet has a significantly better segmentation effect on the pancreas and small target organs than other algorithms, basically completely segmenting the pancreas and stomach. The fourth line reveals that the Attention Unet and Swin-Unet algorithms have a poor segmentation effect on the stomach, with a missing segmentation phenomenon, while RL-Unet and Attention Unet have a better segmentation effect on the stomach.

Despite the RL-Unet architecture’s capability to fuse multi-level semantic features, thereby preserving crucial edge information, its performance in pancreas segmentation remains suboptimal on certain images. This is primarily due to the small proportion of pancreas representation in the dataset, which leads to interference from surrounding organs and tissues during feature extraction. Consequently, a considerable amount of edge information is lost, compromising the model’s overall segmentation accuracy for the pancreas.

Figure 7a describes the average DSC scores obtained by different models with different training iterations on the Synapse dataset. As the number of training iterations increases, the average DSC score of our model is higher than that of other models. The DSC index, which calculates the similarity between the two samples, shows that our training results and labels are more similar, verifying the validity of our proposed model.

Figure 7b depicts the average HD resulting from different iterations of training using different models on the Synapse dataset. As can be seen from Figure 7b, the HD distance of our proposed model is superior to other models, significantly smaller than other models, and higher than the Att-Unet model only at 20 K, which may cause anomalies caused by isolated points, demonstrating that the result of dataset segmentation is good and close to the marked result.

As can be seen from

Table 2. Segmentation results of different methods on BRATS2021 dataset.

Methods	ET	TC	WT	DSC
U-Net [7]	83.92	87.71	90.70	87.44
U-Net++ [24]	84.21	87.76	91.30	87.76
U-Net3+ [9]	82.21	86.52	91.17	86.63
nnFormer [29]	83.93	87.82	91.52	87.74
AttU-Net [10]	85.08	87.76	91.53	88.12
Swin-Unet [20]	84.16	86.77	91.01	87.31
RL-Unet (ours)	85.27	88.12	91.52	88.36

, the RL-Unet algorithm proposed in this paper can still achieve a good segmentation effect on the BraTS2021 dataset, with DSC indexes reaching more than 80%, achieving similar or even surpassing results with U-Net, Attention Unet, Swin-Unet, and other classical medical image segmentation algorithms. It shows that the algorithm has good generalization ability and robustness.

To verify the generalization of the model, we performed another set of experiments on MRI images of the ACDC dataset.

Table 3. Segmentation results of different methods on ACDC dataset.

Methods	RV	Myo	LV	DSC
R50+AttnUnet [19]	87.58	79.20	93.47	86.75
R50+Unet [19]	87.10	80.63	94.92	87.55
Trans-Unet [19]	86.67	87.27	95.18	89.71
Swin-Unet [20]	85.77	84.42	94.03	88.07
MT-Unet [43]	86.64	89.04	95.62	90.43
PVTCASCADE [28]	89.97	88.90	95.50	91.46
RL-Unet (ours)	88.54	87.65	95.46	91.18

shows the results of average DSC comparison between our proposed method and previous advanced methods. It can be clearly seen that compared with previous models such as AttnUnet, SwinUnet, TransUnet, and MT-Unet, the proposed method has a higher average DSC score; compared with PVTCASCADE model, the average DSC score is slightly lower, but it is still quite capable.

As shown in

Table 4. Segmentation results of different methods on BUSI dataset.

Methods	Benign	Malignant	DSC
U-Net [7]	70.49	63.47	66.98
AttU-Net [10]	73.30	62.95	68.13
Swin-Unet [20]	73.31	62.13	67.72
U-Net++ [24]	75.56	65.52	70.54
U-Net3+ [9]	75.07	66.19	70.67
RL-Unet (ours)	74.92	66.09	70.51

, RL-Unet outperforms most SORT architectures on the BUSI dataset, achieving a higher DSC score. Notably, the performance of RL-Unet is slightly below that of U-Net3+. We attribute U-Net3+’s superior performance to its full-scale skip connections, which are particularly effective for accurately segmenting complex structures in medical images. To address the limitations of full-scale skip connections, we have introduced JCA and JSA into our architecture, enhancing its performance to a comparable level.

As illustrated in Figure 8, the detailed annotations and consistent segmentation shapes in the ACDC dataset enable most architectures, including RL-Unet, to achieve high segmentation accuracy and visually appealing results. In certain cases, U-net exhibits partial segmentation omissions. As observed in the third line, the segmentation results produced by RL-Unet and Attention Unet are more complete and closely aligned with the Ground Truth. In contrast, the segmentation contours of Swin-Unet and U-Net show partial omissions. In the classification of benign cases, most architectures achieve satisfactory segmentation, with smooth transitions along the segmentation boundaries. However, in the classification of malignant cases, most architectures are only able to delineate the approximate contours and often lack the ability to capture detailed features accurately. RL-Unet more effectively restores details, mitigating issues of over-segmentation and under-segmentation.

3.4. Ablation Experiment

The four parts of the proposed model are validated by ablation experiments. (1) The MLP in the Transformer module is replaced with Res-MLP and a local inductive bias module is introduced. (2) A new feature aggregation module DUSM is proposed. (3) A new attention module is introduced. (4) A new loss function is designed.

First, replace the MLP with Res-MLP in the Transformer module and add the local inductive bias module. The baseline model is Swin-Unet. As can be seen from

Table 5. An ablation experiment of each module.

Methods	DSC	HD
Baseline	77.28	26.93
Baseline+ResMLP	78.36	28.12
Baseline+ResMLP+DUSM	78.59	26.45
Baseline+ResMLP+DUSM+Attention	79.69	26.78
Baseline+ResMLP+DUSM+Attention+L2	80.13	22.07
Baseline+ResMLP+DUSM+Attention+L1	79.95	24.21

, the DSC index increased by 1.08%. Experiments show that the Res-MLP module and local inductive bias module make the semantic feature propagation of Transformer more comprehensive and increase the encoder’s understanding of global context information, thus improving the performance of the whole structure. Next, a new feature aggregation module was designed at the decoding end, and the index of DSC improved by 1.31%. Meanwhile, based on previous work on linear attention mechanisms, a new attention mechanism was introduced to enhance the channel and spatial relationship between semantic features. Experiments show that our proposed attention mechanism is effective. Finally, a novel loss function was proposed. When the loss function is applied to our proposed method, compared with the Swin-Unet and other algorithms, as shown in Figure 7c, it can be seen that the loss obtained by our proposed method is smoother and has stronger convergence, which verifies the effectiveness of our proposed algorithm. According to multiple experiments, the DSC score increased by 2.85% and the HD index also significantly improved, which proved the effectiveness of our proposed loss function.

3.5. Hyperparametric Learning Rate $\eta$

The experiment shows that with the change in learning rate, the experimental results will also change. In order to obtain the best learning rate, the parameter was selected as 0.01, 0.05, and 0.1 for experimental comparison, as shown in

Table 6. Comparison of segmentation results at different $\eta$.

$\mathit{\eta }$	DSC	HD
0.01	78.51	28.67
0.05	80.13	22.07
0.1	78.87	27.43

. Under other conditions being the same, it is found that the best segmentation results can be obtained when $\eta$ is 0.05. Consequently, $\eta$ = 0.05 was used as the starting point for every experiment in this research.

4. Conclusions

In this paper, we propose an RL-Unet algorithm based on the U-Net and Swin-Transformer models. It integrates multi-scale feature information and inductive bias module for local feature processing, and further improves prediction accuracy through multi-scale information fusion selection. We evaluate the segmentation results of our approach on four datasets, namely, the synapse multi-organ dataset and the BraTS2021 dataset. Additionally, we conduct ablation studies on each component of our design to analyze its individual contributions. The experimental results demonstrate that our algorithm outperforms traditional medical image segmentation algorithms in terms of segmentation results and indices. RL-Unet exhibits suboptimal performance on small sample and fine-grained segmentation tasks. Future work will investigate lightweight modules, which aim to mitigate computational burden while preserving the model’s expressive capability, to enhance the focus on challenging regions and pixels.