Domain generalization for prostate segmentation in transrectal ultrasound images: A multi-center study

Abstract

Prostate biopsy and image-guided treatment procedures are often performed under the guidance of ultrasound fused with magnetic resonance images (MRI). Accurate image fusion relies on accurate segmentation of the prostate on ultrasound images. Yet, the reduced signal-to-noise ratio and artifacts (e.g., speckle and shadowing) in ultrasound images limit the performance of automated prostate segmentation techniques and generalizing these methods to new image domains is inherently difficult. In this study, we address these challenges by introducing a novel 2.5D deep neural network for prostate segmentation on ultrasound images. Our approach addresses the limitations of transfer learning and finetuning methods (i.e., drop in performance on the original training data when the model weights are updated) by combining a supervised domain adaptation technique and a knowledge distillation loss. The knowledge distillation loss allows the preservation of previously learned knowledge and reduces the performance drop after model finetuning on new datasets. Furthermore, our approach relies on an attention module that considers model feature positioning information to improve the segmentation accuracy. We trained our model on 764 subjects from one institution and finetuned our model using only ten subjects from subsequent institutions. We analyzed the performance of our method on three large datasets encompassing 2067 subjects from three different institutions. Our method achieved an average Dice Similarity Coefficient (Dice) of 94.0±0.03 and Hausdorff Distance (HD95) of 2.28 mm in an independent set of subjects from the first institution. Moreover, our model generalized well in the studies from the other two institutions (Dice: 91.0±0.03; HD95: 3.7 mm and Dice: 82.0±0.03; HD95: 7.1 mm). We introduced an approach that successfully segmented the prostate on ultrasound images in a multi-center study, suggesting its clinical potential to facilitate the accurate fusion of ultrasound and MRI images to drive biopsy and image-guided treatments.

Keywords: Continual learning segmentation; Deep learning; Gland segmentation; Prostate MRI; Targeted biopsy; Transrectal ultrasound.

Fig. 1:

Example of prostate outlines (red) on ultrasound images acquired at three different institutions. Note the large variations in prostate shape, contrast, field of view, and the presence of artifacts (e.g., inhomogeneous intensity distributions). Moreover, the prostate boundary is not always clearly visible or easily distinguishable from the neighboring structures (orange arrows).

Fig. 2:

Flowchart of the dilated Residual UNet with coordination attention block and 2.5D input (128×160 of three neighboring slices used to provide spatial context). The output of the model is the probability map of the prostate segmented in the central slice. The attention blocks assist the model in reducing uncertainty in prostate borders, which are well known to be challenging to segment due to a lack of distinct borders. I_s is the input data. E and D refers to encoder and decoder blocks. $\widehat{\mathbf{Y}}$ is the predicted segmentation output.

Fig. 3:

The overall pipeline for continual prostate gland segmentation. The segmentation model M_s is trained on cohort C1 in the first step, then the model is finetuned on cohorts C2 and C3. During optimization, the consistency of the features latent space of M_s and M_t1 is maintained by including the L2 norm loss (eq. 2). The same steps are repeated for finetuning the model trained on M_s and M_t1 for M_t2 model. E₁, E₂, D₁, D₂ are the encoders and decoders and z and z′ are the latent space features for the source and target models. Y is the ground-truth and $\widehat{\mathbf{Y}}$ is the predicted segmentation output.

Fig. 4:

Visual comparison of prostate segmentation results produced by different methods for a sample patient in cohort C1-test. From left to right are the input TRUS slices (column (a)), the UNet 2.5D predictions (column (b)), Attention-UNet 2.5D (column c), Nested-UNet 2.5D (column (d)), DAFNet (Wang et al., 2019) (column (e)) and our proposed CoordDR-UNet 2.5D (column (f) for different slices from the apex (top row) to base (bottom row). The blue contours show the ground-truth segmentation outlined by an expert urologist.

Fig. 5:

Visual comparison of CoordDR-UNet segmentation output with 2D, 3D, and 2.5D as the input. Each row shows apex, mid-gland, and base slices from different subjects from the cohort C1-test dataset.

Fig. 6:

3D visualization of the segmentation results in a TRUS volume. CoordDR-UNet output with (a) 2D, (b) 2.5D and (c) 3D representations. The ground-truth mask is shown on a blue surface.

Fig. 7:

3D visualization of the surface distance (in voxel) between segmented surface and ground truth. Each row shows one subject. Different colors represent different surface distances. From left to right are (a) ground-truth, (b) UNet 2.5D, (c) Attention-UNet 2.5D, (d) Nested-UNet 2.5D, (e) DAFNet3D (Wang et al., 2019) and (f) our proposed CoordDR-UNet 2.5D. Our method consistently performs well on the whole prostate surface.

Fig. 8:

Box plots of Dice measure for the CoordDR-UNet, CoordDR-UNet + W/o pretrained, CoordDR-UNet with standard finetuning, and CoordDR-UNet + KDL segmentation approaches of cohort C2-test and C3-test. SS: statistically significant (P ≤ 0.05), NS: not significant (P > 0.05).

Fig. 9:

Visual comparison of segmentation results produced by different methods for three patients (best, average, worst) in cohort C2-test. From left to right are the 2D TRUS slices (column (a)), CoordDR-UNet 2.5D prediction without finetuning (column (b)), CoordDR-UNet 2.5D trained on ten TRUS images without any pretraining (column (c)), and our proposed CoordDR-UNet 2.5D finetuned on ten TRUS images with knowledge distillation loss (column (d)). Dice and HD95 scores are shown for each patient.

Fig. 10:

Visual comparison of segmentation results produced by different methods for three patients (best, average, worst) in cohort C3-test. From left to right are the 2D TRUS slices (column (a)), CoordDR-UNet 2.5D direct prediction without finetuning (column (b)), CoordDR-UNet 2.5D trained on ten TRUS images without any pretraining (column (c)), and our proposed CoordDR-UNet 2.5D finetuned on ten TRUS images with knowledge distillation loss (column (d)). Dice and HD95 scores are also shown for each patient.

Fig. 11:

3D visualization of the surface distance (in voxel) between segmented prostate and ground truth. Different colors represent different surface distances. The top row shows the ground-truth surface masks for three different patients from cohort C1-test (a), C2-test (c), and C3-Test (b). The bottom row shows the computed HD surface distance between CoordDR-UNet 2.5D predictions and ground-truth.

Fig. 12:

Continual segmentation results for domain generalization. M_s: model trained using the train data from cohort C1, M_t1 : model M_s, finetuned using 10 TRUS images from cohort C2, and M_t2 : model M_t1, finetuned using 10 TRUS images from cohort C3.

Fig. 13:

(a) CoordDR-UNet 2.5 performance with different number of slices on cohort C1-test set. (b) CoordDR-UNet + LKD finetuning performance with different number of cases on cohort C2-test and C3-test data. (c) CoordDR-UNet + LKD performance with different λ values on cohort C2-test data.