Automatic Segmentation of the Prostate on MR Images based on Anatomy and Deep Learning

Abstract

Accurate segmentation of the prostate has many applications in the detection, diagnosis and treatment of prostate cancer. Automatic segmentation can be a challenging task because of the inhomogeneous intensity distributions on MR images. In this paper, we propose an automatic segmentation method for the prostate on MR images based on anatomy. We use the 3D U-Net guided by anatomy knowledge, including the location and shape prior knowledge of the prostate on MR images, to constrain the segmentation of the gland. The proposed method has been evaluated on the public dataset PROMISE2012. Experimental results show that the proposed method achieves a mean Dice similarity coefficient of 91.6% as compared to the manual segmentation. The experimental results indicate that the proposed method based on anatomy knowledge can achieve satisfactory segmentation performance for prostate MRI.

1. INTRODUCTION

Prostate segmentation has many applications in clinical diagnosis and treatment of prostate diseases, especially prostate cancer¹. A manual delineation on prostate is subject to inter-observer variability and is time consuming. Semiautomatic segmentation methods expect the user to provide the interactive information^2–4. Automated segmentation may be robust and fast^5–6. In recent years, deep learning methods are gradually applied and designed for prostate segmentation. Guo et al proposed a new deformable MR prostate segmentation method by unifying the stacked sparse auto-encoders with the sparse patch matching⁷. Tian et al used a pretrained fully convolutional network (FCN) with fine-tuning for prostate MRI segmentation⁸. Milletari et al presented a volumetric convolutional neural network called V-Net for the prostate segmentation using 3D convolutional layers with an objective function directly based on the Dice coefficient⁹. Ma et al propose an automatic prostate segmentation method by combining the convolutional neural network and multi-atlas refinement¹⁰. Yu et al proposed a 3D fully convolutional network with long and short residual connections for automated prostate segmentation from MR images¹¹. Nie et al designed an attention model based semi-supervised deep networks to segment prostate from MRI¹².Wang et al proposed a three-dimensional FCN with deep supervision and group dilated convolution to segment the prostate on MRI¹³. Zhang et al proposed a novel architecture, namely Z-net, for segmenting prostate from MRI, assembling five pairs of Z-block and decoder Z-block with different sizes and numbers of feature maps¹⁴. Si et al proposed a multi-step segmentation method for prostate MR image based on deep reinforcement learning¹⁵. Jia et al proposed a hybrid discriminative network consisting of a 3D segmentation decoder using channel attention block and an auxiliary 2D boundary decoder guiding the segmentation network for prostate segmentation in MR images¹⁶. Jia et al propose a 3D adversarial pyramid anisotropic convolutional deep neural network for prostate segmentation in MR images¹⁷. Liu et al proposed a multisite network for improving prostate segmentation by learning robust representations, leveraging multiple sources of data¹⁸. Rundo et al proposed a novel convolutional neural network, called USE-Net, for the prostate zonal segmentation task, incorporating Squeeze-and-Excitation (SE) blocks into U-Net¹⁹.

In this paper, we propose an automatic segmentation method for the prostate on MR image based on the anatomical structure of prostate. According to the anatomy, the prostate is between the rectum and the pelvic bone and is a walnut shaped organ with a roughly elliptical cross section. In addition to the intensity information, we involve the anatomical information as the constraint loss into the convolutional neural networks, therefore, for improving the performance of the prostate segmentation.

2. METHOD

We propose an automatic segmentation model based on the anatomy of the prostate. We utilize the Dice loss, location loss, and shape loss to constrain the segmentation process. The overall architecture of our method is shown in Figure 1. Our method takes the U-Net as the base network structure to extract the deep features, and uses the gray and anatomical information as the loss function for improving the performance.

Figure 1.

Overview of the proposed method.

2.1. Deep feature generation

We employ 3D U-Net as the base network structure for extracting the deep features of the prostate. This network contains an extraction and a restoration path each with four resolution steps. In the extraction path, each layer contains two 3 × 3 × 3 convolutions with strides of one each followed by a batch normalization layer (BN) and rectified linear unit layer (ReLu), which does not reduce the resolution, and then a 2 × 2 × 2 max pooling with strides of two in each dimension. In the restoration path, each layer consists of an upconvolution of 2 × 2 × 2 by strides of two in each dimension, followed by two 3 × 3 × 3 convolutions each followed by a BN and a ReLu. Skip connections from layers of equal resolution in the extraction path provide the necessary high-resolution features to the restoration path. In the last layer a 1 × 1 × 1 convolution reduces the number of output channels to the number of ground truth which is 1 in our case.

2.2. Loss functions

We combine three different loss functions together for improving the segmentation performance.

Dice Loss:

There is a serious category imbalance due to the wide difference in the amount of foreground and background for the prostate on MR images. Dice coefficient can solve this problem, so we convert Dice coefficient as the basic loss function of our neural network. The Dice loss L_Dice between two binary volumes can be written as

$$L_{\mathit{\text{Dice}}}=1-\frac{2{\sum }_i^N{p}_i{g}_i}{{\sum }_i^N{p}_i^2+{\sum }_i^N{g}_i^2},$$ (1)

where the p_i and g_i are the i-th voxel in the predicted segmentation volume and ground truth volume, respectively, and N is the number of the voxels.

Location Loss:

The prostate MR images are unevenly distributed in density and contain noise that they are often misclassified. Since the prostate is between the rectum and the pelvic bone, it is roughly in the middle of the MR images. Hence, we use the location loss function to remove the irrelevant regions. The minimum bounding rectangle (MBR) of the ground truth is used to determine the position of the prostate in the whole MR image. Then, the Dice loss between ground truth and segmentation are calculated at the corresponding position. In this way, even without post-processing, the accuracy of the model can be improved and the model can fit the data more quickly. L_Location can be written as

$$L_{\mathit{\text{Location}}}=L_{\mathit{\text{Dice}}}(\mathit{\text{MBR}}(\mathit{\text{segmentation}}),\mathit{\text{MBR}}(\mathit{\text{groundTruth}}))$$ (2)

Shape Prior Loss:

Since the edge of the prostate in MR image is very fuzzy, we may obtain the segmented prostate with the bad boundary. The low-dimensional features contain some concrete information such as color, corner and contour, while the high-dimensional features contain some abstract features such as shape. The auto-encoder machine (AE) can extract the high-dimensional features effectively, so it is used to obtain the shape-related information. We adopt the simplest form of down-sampling and up-sampling. Input the ground truth into the AE and extract the high-dimensional features such as shape, then restore to the original size after up-sampling. By comparing the predicted results with ground truth, the network can well extract the shape features of ground truth. Figure 2 shows the process of obtaining the shape information.

Figure 2.

An auto-encoder machine for obtaining shape information

We use the auto-encoder machine to obtain the shape prior knowledge of the segmentation target. We extract shape-related information from the ground truth and the predicted segmentation using the encoder part, and then calculate their Kullback-Liebler (KL) divergence. The loss of shape prior can be written as

$$L_{\mathit{\text{Shape}}}=e^{D_{KL}\left(V_{\mathit{\text{seg}}},V_{gt}\right)},$$ (3)

where V_seg and V_gt are the feature vector generated from the encoder. D_KL(.) is the KL divergence. KL divergence is used to measure the similarity between two feature vectors.

Total Loss:

We combine the three loss functions in a weighted sum. The total loss can be written as

$$L_{\mathit{\text{total}}}={\lambda }_1{L}_{\mathit{\text{Dice}}}+{\lambda }_2{L}_{\mathit{\text{Location}}}+{\lambda }_3{L}_{\mathit{\text{Shape}}},$$ (4)

where λ_1, λ_2, λ₃ are the weight coefficients which are set to 1, 1, and 0.5 in this paper, respectively.

3. EXPERIMENTS

3.1. Database

We evaluate our method on the 50 MRI volumes from the “PROMISE2012” challenge dataset²⁰. This dataset contains medical data acquired in different hospitals, using different equipment and different acquisition protocols.

3.2. Experiment details

We implemented the proposed method based on the Pytorch framework with one Nvidia V100 32G GPU. In the process of pre-processing, we unify the spatial resolution of each volume to 1.0 × 1.0 × 1.5 millimeters, and use Adam optimizer, the initial learning rate is 0.002 and decrease by a weight decay of 0.001 after 10, 30, 60 epochs. The batch size is 2 and the number of epochs is 100. All the volumes processed by the network have fixed size of 64 × 128 × 128 voxels. We conducted leave-one-out experiments for the prostate segmentation. We take each MR image as the testing sample in turn, and the 49 remaining samples as the training set.

3.3. Evaluation metrics

The proposed method was evaluated based on the manually labeled ground truth. Dice similarity coefficient (DSC) is used for the segmentation evaluation. The DSC formula is:

$$\mathit{\text{DSC}}=\frac{2\times TP}{(FP+TP)+(TP+FN)},$$ (5)

where TP, TN, FP, FN are the true positives, true negatives, false positives and false negatives, respectively.

The Hausdorff distance, HD(A, B) is used for measuring the surface distance between two surfaces A and B, and it is defined as

$$HD(A,B)=\text{max}(\text{max\hspace{0.17em}}d(a,B),\text{max\hspace{0.17em}}d(b,A)),$$ (6)

where max(u, v) is a function returning the bigger value of the u and v, and d(x, Y) is a distance of a pixel x to a surface Y, and it is defined as

$$d(x,Y)={\text{min}}_{y\in Y}\Vert x-y\Vert$$ (7)

3.4. Quantitative results

Figure 3 shows the quantitative segmentation results in the leave-one-out experiments. We can obtain a mean DSC of 91.6% with the standard deviation of 3.26%, and a mean HD of 5.18 mm with the standard deviation of 1.61 mm. Those results mean that our method can achieve good segmentation performance.

Figure 3.

The quantitative results, (a) Dice similarity coefficient (DSC) and (b) Hausdorff distance in mm (HD), of the proposed method in the leave-one-out experiments.

3.5. Qualitative Results

The qualitative evaluation results are shown in in Figure 4. The proposed automatic segmentation is close to the manual segmentation of the experienced radiologist.

Figure 4.

The qualitative evaluation of the proposed segmentation method. The blue curves are the manual segmented prostate by the experienced radiologist, while the red curves are the contours obtained by our automatic segmentation method.

3.6. Effectiveness of the anatomy

We test the effectiveness of the anatomy on the segmentation performance. We show the four pairs of segmented results in Figure 5, where the left segmented result does not use the anatomy information and the right one involves the anatomy information in each pair. In Figure 5, we can see that the segmented prostates involved the anatomy have a more complete region. That proves the anatomy knowledge is effect on the segmentation performance.

Figure 5.

The compared segmented results involved the anatomy and without the anatomy.

3.7. Comparison results

We compared our method with seven published segmentation methods on the same database. The comparison results are shown in Figure 6. The proposed method achieved a higher DSC as compared to the existing methods.

Figure 6.

The comparison between the proposed method and existing approaches.

4. CONCLUSION

In this paper, we propose an automatic segmentation method for the prostate on MR images based on the anatomy. 3D U-Net guided by anatomy knowledge is employed. The location and shape prior information of the prostate are involved as the knowledge to constrain the segmentations to coincide with object boundaries. We combine the Dice loss, location loss, and shape loss, together to obtain accurate segmentation of the prostate. Experimental results show that the proposed method could achieve satisfactory segmentation performance.