The UNet network consists of an encoder and a decoder. The encoder part follows the classical structure of convolution networks. The convolution block consists of two repeated 3 × 3 convolutions. Each convolution is followed by a ReLU activation function and a 2 × 2 maximum pooling operation with a step of 2 for downsampling. In each downsampling step, the number of feature channels is doubled. Each step in the decoder upsamples the size of feature maps by 2, and the number of feature map channels is reduced to half using 2 × 2 convolution (deconvolution). Feature maps from the encoder are directly passed to the decoder with skip connections. After concatenation, there is a convolution block to reduce the number of channels. At the end of the decoder, there is a 1 × 1 convolution layer which is designed for the output.

UNet obtains its energy function by combining the pixel-level softmax function calculated for the last layer of the feature map with the cross-entropy loss function. The definition of the softmax function is as follows:

where represents the activation function in feature channel k at the pixel position x ∈ Ω, Ω ⊂ ℤ². K is the number of classes and p_k(x) is the approximated maximum-function. p_k(x) ≈ 1 for the k that has the maximum activation a_k(x), and p_k(x) ≈ 0 for all other k. The cross entropy then penalizes at each position the deviation of p_l(x)(x) from 1 using:

where l : Ω → {1, ..., K} is the true label of each pixel, and w : Ω → ℝ is a weight map which can make some pixels more important than the others while training (7).

Attention gate is a mechanism that can be merged into any existing CNN architecture. Let xl={xil} be the activation map of the chosen layer l ∈ {1, ..., L}, where each xil represents a pixel-by-pixel feature vector of length F_l (i.e., the number of feature-maps in layer l). For every xil, AG will calculate the coefficient αl={αil}=1n, αil∈[0,1], in order to identify the key region of the feature map and only reserve the parts that are related to specific tasks. The output of the attention gate is:

in which each vector is scaled by the corresponding attention coefficient (10).

In this study, inspired by the UNet framework, a new prostate segmentation network is proposed. The architecture of the whole network is shown in Figure 1. First, the first half of the network (i.e., encoder) is used to extract features from 2d slices of MRI images of 3d tissues. The second half (i.e., decoder) is then used to generate the predicted segmentation results, where each type of label is segmented in different channels.

The network proposed in this article consists of two parts, an encoder and a decoder. The encoder consists of five convolution blocks, four Max Pooling blocks, and a spatial dilated convolution pyramid (ASPP) module. In the first four convolution blocks, each convolution block consists of two 2d convolution layers followed by group normalization (GN) and ReLu activation function. The fifth convolution block adds a dropout layer on the basis of the first four to prevent overfitting. An ASPP module is added at the bottom of the encoder to further extract features. Each Max Pooling block performs a maximum pooling to achieve downsampling of the feature map by 2. The decoder consists of four upsampling modules, four attention gate (AG) modules, three convolution blocks, and a 2d convolution layer for output. Each upsampling module uses the nearest neighbor interpolation to upsample the feature map. The final 2d convolution layer is responsible for outputting segmentation results.

As for batch normalization (BN), it will get a decline in performance if the batchsize is too small; on the other hand, big batchsize will consume a lot of memory, especially when large-size images are input into the network. In order to get better segmentation results while the batchsize is small, we decided to use group normalization (GN) instead of BN.

The structure of the first four convolution blocks and the convolution blocks used in the decoder part is shown in Figure 2.

Using group normalization solves the internal, covariate, and shift problems, and the result is much better than batch normalization when the batchsize is small (in our experiments the batchsize is set to 4). The definition of group normalization is:

The input feature map x is divided into several groups according to the channel, and the mean and standard deviation of each group are calculated, respectively. γ and β are learnable parameters.

In the encoder part, the input image will go through four such convolution blocks to extract its features. After each of those four convolution blocks, the max pooling module will use the maximum pooling to downsample the feature map (by 2), the number of channels of the feature map remains unchanged, and the length and width become half of the original, so as to further extract features and reduce the number of parameters.

After four rounds of downsampling, the feature map comes to the fifth convolution block. The structure of the last convolution block in the encoder is shown in Figure 3. Compared with the first four convolution blocks, this convolution block has an additional Dropout layer (where the p-value is set to 0.5). Experiment results show that adding Dropout can effectively alleviate overfitting and improve the final segmentation results.

The final part of the encoder is an ASPP module, which is added to extract further features, and its structure is shown in Figure 4. For the input feature map, ASPP uses dilated convolution with different dilation rates to process it (in this article, the dilation rates are set to 1, 6, 12, and 18), then concatenates the obtained results together, expands the number of channels, and finally reduces the number of channels to the desired value through a 1^*1 convolution layer.

The algorithm of the ASPP module is as follows:

The feature map obtained by the encoder will be sent to the decoder. The structure of the upsampling module is shown in Figure 5. We use the nearest neighbor interpolation for upsampling, which is defined as:

The size of the input feature map is W ^* H, and the size of the upsampled feature map is w ^* h. The pixel value of pixel (x, y) on the upsampled feature map equals the pixel value of pixel (W/w ^* x, H/h ^* y) on the original feature map.

In order to improve the ability to capture key regions and channels, we added a channel attention mechanism to our attention gate. We calculate the attention coefficient using:

Spatial attention:

where σ₁(x_i) = max(0, x_i) represents ReLu function, σ2(xi,c)=11+exp(-xi,c) denotes sigmoid function, and AG is represented by a set of parameters Θ_atts, including the linear transformations Wx∈ℝFl×Fint, Wg∈ℝFg×Fint, and W2∈ℝFint. The linear transformations are realized by 1 × 1 convolution.

Channel attention:

Different from spatial attention, in order to obtain the weight of each channel of the input feature map, the channel attention mechanism uses adaptive average pooling (the AvgPool part). In fact, adaptive average pooling works better than adaptive max pooling or use both of them in channel attention. The remaining linear transformations include two 1 × 1 convolutions: W0∈ℝFg×Fg/16 and W1∈ℝFg/16×Fg.

The output of the attention gate is:

The overall structure of our attention gate is shown in Figure 6:

The algorithm of the attention gate is as follows:

For multi-class segmentation tasks, we used dice coefficients for each class as the main metric to measure the segmentation effect. Dice is the most frequently used metric in medical image competition. It is a set similarity metric, which is usually used to calculate the similarity of two samples, and the threshold is [0, 1]. It is often used for image segmentation in medical images. The best result of segmentation is 1, and the worst result is 0. The dice coefficient is calculated as follows:

where pred is the set of predicted values, true is the set of true values, the molecule is the intersection between pred and true, and the denominator is the union of pred and true.

In addition to the dice coefficient, we also applied PPV and sensitivity metrics to our experiment to further measure the segmentation result.

The definitions of PPV and sensitivity are as follows:

The experiment used the Task05 prostate data set in the medical decathlon competition. Due to the small number of experimental samples, we used offline data augmentation to enrich the training set. Specifically, we first performed data augmentation on prostate MRI images and corresponding labels, including horizontal/vertical flipping, rotation, adding Gaussian noise, and adjusting brightness and contrast, so that the amount of data in the training set is expanded to three times of the original set, which effectively alleviates the problem of small dataset and insufficient training data. Then, some operations are used for data preprocessing, including uniform image size, normalization, and slicing.

Among them, 80% of the original dataset and the pseudo image obtained by data augmentation are used as the training set, and 20 % of the original dataset is used as the validation set. In order to compare the performance of different networks, we trained the network proposed in this study and five other UNet based networks on the same dataset and under same hyperparameters as a comparison. All experiments are based on python language and pytorch framework, carried out on a server equipped with RTX308 and Windows11 system. The parameters of our experiment are shown in Table 1.

The experiment results of all networks are shown in Tables 2–4.

From the tables mentioned above, it can be seen that the network we proposed in this study has achieved 0.807 and 0.907 dice scores in peripheral zone and transition zone of the prostate, respectively. It also achieved the best PPV and sensitivity scores. Compared with the other five UNet-based networks, the proposed method is better in the prostate segmentation task.

The segmentation map is shown in Figure 7, in which the green area represents the peripheral zone (PZ) and red area represents the transition zone (TZ).