End-to-end prostate lesion segmentation and detection from undersampled diffusion-weighted MRI
Authors: Lukas T. Rotkopf, Marco Schlimbach, Kevin Kröninger, Stefan O. Schoenberg, Heinz-Peter Schlemmer, Jens Kleesiek, Moritz Rempe
Organizations: Institute for AI in Medicine (IKIM), University Hospital Essen, Girardetstraße 2, 45131, Essen, Germany. · Department of Physics, Technical University Dortmund, Otto-Hahn-Straße 4a, 44227, Dortmund, Germany. · Division of Radiology, German Cancer Research Center (DKFZ), 69120, Heidelberg, Germany. · Department of Radiology and Nuclear Medicine, University Medical Center Mannheim, Heidelberg University, Theodor-Kutzer-Ufer 1-3, 68167, Mannheim, Germany. · Cancer Research Center Cologne Essen (CCCE), University Medicine Essen, Hufelandstraße 55, 45147, Essen, Germany. · Site Essen, RACOON Study Group, Essen, Germany. · Partner Site Essen, German Cancer Consortium (DKTK), Hufelandstraße 55, 45147, Essen, Germany. · Medical Faculty and Faculty of Computer Science, University of Duisburg-Essen, 45141, Essen, Germany. · Department of Radiology, Brigham and Women’s Hospital, Harvard Medical School, Boston, USA.
Automated prostate lesion detection from diffusion-weighted MR imaging (DWI) conventionally relies on reconstructed images and does not fully use all acquisition information. We present an end-to-end network that uses complex multicoil DWI data and coil sensitivity information for prostate lesion segmentation. We included 269 fastMRI prostate examinations with additional voxel-level lesion annotations. Of these, 64 examinations contained PIRADS 4 or 5 lesions. The proposed network, VarNet-Seg, consists of a k-space variational network (VarNet) with six variational cascades and an image-space 3D segmentation U-Net. Comparator models used the same segmentation architecture with fully coil-combined images (CC-Seg), fixed root-sum-of-squares reconstruction (RSS-Seg), or fixed SENSE reconstruction (SENSE-Seg). We performed ablations with end-to-end training and data consistency (DC) enabled or disabled. We evaluated lesion segmentation using Dice and patient-level detection using AUROC in five-fold cross-validation with full sampling, simulated undersampling, added noise, or both. VarNet-Seg achieved the highest Dice at every undersampling level. At eight-fold acceleration, Dice was 0.22 for VarNet-Seg, 0.14 for SENSE-Seg and 0.16 for VarNet-Seg (no DC). The paired difference from SENSE-Seg was 0.08 (95% CI, 0.04-0.11). With full sampling, VarNet-Seg and RSS-Seg performed similarly (0.40 vs 0.39). Patient-level AUROC decreased from 0.75 with full sampling to 0.63 at eight-fold acceleration; corresponding RSS-Seg AUROCs were 0.76 and 0.64. The proposed multicoil network improved prostate lesion segmentation under simulated undersampling compared to conventional models. Matched ablations supported contributions from both end-to-end training and data consistency. Prospective evaluation is needed to establish whether this approach can be used clinically in accelerated prostate DWI.
Figures & tables
Figure 1: Model architectures and evaluation. (a) The illustrated core architectures use the same 3D segmentation U-Net in image space, with the k-space stage either fixed or trained jointly with segmentation. Red dashed contours indicate class-specific predictions and green contours the reference annotations. (b) Illustration of the evaluation setup with full sampling, undersampling, added noise and their combination. Sampling parameters are provided in Table S3 .
Figure 2: Segmentation and patient-level detection with undersampling. Dice for combined PIRADS 4+5 lesions is shown in the left column (a, c) and patient-level AUROC in the right column (b, d). The upper row compares CC-Seg, RSS-Seg, SENSE-Seg and VarNet-Seg. The lower row compares VarNet-Seg (frozen), VarNet-Seg (frozen, no DC), VarNet-Seg (no DC) and VarNet-Seg. Error bars show 95% bootstrap confidence intervals. Asterisks indicate VarNet-Seg–RSS-Seg comparisons in a and VarNet-Seg–VarNet-Seg (frozen, no DC) comparisons in c and d ( ∗p<0.05 , ∗∗p<0.01 , ∗∗∗p<0.001 ). The dashed vertical line separates the training range from higher undersampling factors.
Figure 3: Examples of PIRADS 3, 4 and 5 lesions with full sampling and four-fold undersampling. Contours show the indicated grade only: reference annotations in green and predictions as red dashed lines. Background images are SENSE b1000 reconstructions, with matching slices and crops across models. Dice values describe the displayed grade over the full prediction volume.
Figure 4: Overall segmentation Dice with added noise for CC-Seg, RSS-Seg, SENSE-Seg and VarNet-Seg. (a) Noise only. (b–d) Undersampling and added noise at SNRs of 20, 10 and 5 dB, respectively. Asterisks in b–d indicate VarNet-Seg–RSS-Seg comparisons ( ∗∗p<0.01 , ∗∗∗p<0.001 ).
Figure 5: Confusion matrices with full sampling for CC-Seg, RSS-Seg, SENSE-Seg and VarNet-Seg. Values are percentages, normalized by reference class within each patient and averaged across patients with that class. Patient counts are given in parentheses. BG, background.
Model
Parameters (M)
Training time (min)
Inference time (ms/volume)
CC-Seg
2.71
6.56
3.13
RSS-Seg
2.71
30.12
3.85
SENSE-Seg
2.71
27.23
4.15
VarNet-Seg
3.12
60.05
29.18
VarNet-Seg (no DC)
3.12
59.77
29.18
VarNet-Seg (frozen)
3.12
58.20
29.17
Table 1: Computational cost for training and inference measured across folds including parameter counts. Training time includes validation and, for the frozen variants, both reconstruction pretraining and segmentation training. Inference measures the GPU forward pass only.
Lesion category
Patients
No annotated lesion
130
PIRADS 3
75
PIRADS 4
41
PIRADS 5
23
Total
269
Table S1: Study cohort. Examinations with lesions are classified by the highest annotated DWI PIRADS score.
Component
Setting
Image dimensions
100×100 -pixel reconstruction matrix, central 64×64 -pixel crop for segmentation and 24-slice input windows.
VarNet module
Six cascades with separate weights, 24 channels and two residual blocks with in-plane dilations of 1 and 2, without pooling. Final layers initialized to zero, residual scales to 0.1 and data-consistency weights to 0.5.
Segmentation U-Net
Channels 24, 48, 96 and 192, with strides (1,1,1) , (1,2,2) and (2,2,2) . Two residual units per level, group normalization with two groups and LeakyReLU with slope 0.1. A linear layer maps 24 output features to four classes.
Batch and sampling
Batch size 16. Windows containing target lesions receive half the total sampling probability.
Loss weights
Dice loss 0.7 and focal loss 0.15. Reconstruction loss 0.1 for VarNet-Seg and VarNet-Seg (no DC).
Optimization
AdamW, learning rate 2×10−4 , weight decay 0.01 and gradient clipping at 1.0. Cosine decay to 10−6 over 100 epochs, with early-stopping patience 30.
Table S2: Model and training settings.
Undersampling factor
Central lines
Retained lines
1
—
100
2
4
52
4
8
31
6
5
20–21
8
4
16
10
3
12–13
Table S3: Sampling parameters for the 100 phase-encoding lines. Retained line counts include the central band.
Figure S1: Grade-specific Dice with undersampling for all seven models, including VarNet-Seg (frozen) and VarNet-Seg (frozen, no DC). (a) PIRADS 4 and (b) PIRADS 5. Error bars show 95% bootstrap confidence intervals. Asterisks indicate VarNet-Seg–RSS-Seg comparisons ( ∗p<0.05 , ∗∗p<0.01 , ∗∗∗p<0.001 ).
Computational Radiology & Artificial Intelligence Unit, Department of Physics and Computational Radiology, Oslo University Hospital, Oslo, Norway · Center for Lifespan Changes in Brain and Cognition, Department of Psychology, University of Oslo, Oslo, Norway · Division of Radiology and Nuclear Medicine, Oslo University Hospital, Oslo, Norway +5
the State Key Laboratory of Multimodal Artificial Intelligence Systems, Institute of Automation, Chinese Academy of Sciences, Beijing, 100190, Beijing, China · the School of Artificial Intelligence, University of Chinese Academy of Sciences, Beijing, 101408, Beijing, China · the Department of Radiology, Peking Union Medical College Hospital, Chinese Academy of Medical Sciences & Peking Union Medical College, Beijing, 100730, Beijing, China +5