AG-SE 3D U-Net: A Dual-Attention Deep Learning Framework for Automated Whole Prostate Gland Segmentation in Volumetric MRI
Downloads
Segmentation of the prostate gland from three dimensional (3D) magnetic resonance imaging (MRI) is an essential step in prostate cancer diagnosis and treatment planning. Nevertheless, whole-gland segmentation accurately continues to be a challenge due to the high degree of inter-patient anatomical variability, low soft tissue contrast, and variation in signal intensities across different MRI protocols. This is particularly true at the apex, base, and peripheral zones as boundary error could impact clinical decision making. In this study, we propose the use of a volumetric deep learning architecture called AG-SE 3D U-Net that incorporates attention gates (AG) in all encoder-decoder skip connections and squeeze-and-excitation (SE) blocks in each convolutional block to improve the ability of the model to represent features both spatially and by channels. To ensure stable training of the network, group normalization and a composite loss function that combines binary cross-entropy, dice and focal Tversky losses have been used.
We performed a four variant ablation study (baseline, +SE, +AG, +AG+SE) on 66 paired 3D T2-weighted MRI volumes from the ProstateX database using fivefold stratified cross validation with test time augmentation (TTA). We observed a mean DSC of 0.8697±0.0143, IoU of 0.7730±0.0210, and HD95 of 3.46±0.69mm, with the highest DSC value of 0.8891 in Fold 4.When compared to the baseline 3D U-Net (DSC: 0.8553±0.0192; IoU: 0.7517±0.0273; HD95: 3.37±0.28mm), our proposed model had better segmentation performance based on overlap metrics. Additionally, qualitative evaluation of the images revealed sharper borders and less leakage into adjacent structures. Therefore, our results demonstrate that AG-SE 3D U-Net represents a viable strategy for automatic prostate segmentation.
[1] L. E. O. Jacobson et al., “Prostate MR image segmentation using a multi ‑ stage network approach,” Int. Urol. Nephrol., no. 0123456789, 2025, doi: 10.1007/s11255-025-04763-0.
[2] Z. I. A. Khan, N. Yahya, K. Alsaih, M. Isam, and F. Meriaudeau, “Algorithms of MRI Prostate Regions : A Review,” vol. 4, pp. 1–28, 2021, doi: 10.1109/ACCESS.2021.3090825.
[3] D. K. Lee et al., “Three-dimensional convolutional neural network for prostate MRI segmentation and comparison of prostate volume measurements by use of artificial neural network and ellipsoid formula,” Am. J. Roentgenol., vol. 214, no. 6, pp. 1229–1238, 2020, doi: 10.2214/AJR.19.22254.
[4] Y. J. Chou, B. E. Jong, and Y. C. Tsai, “Integrating PSA Change with PSA Density Enhances Diagnostic Accuracy and Helps Avoid Unnecessary Prostate Biopsies,” Diagnostics, vol. 15, no. 16, pp. 1–11, 2025, doi: 10.3390/diagnostics15162027.
[5] S. M. Rezaeijo, S. J. Nesheli, M. F. Serj, M. Javad, and T. Birgani, “Segmentation of the prostate , its zones , anterior fibromuscular stroma , and urethra on the MRIs and multimodality image fusion using U-Net model,” vol. 12, no. 10, pp. 4786–4804, 2022, doi: 10.21037/qims-22-115.
[6] A. Vaitulevičius, J. Bernatavičienė, J. Markevičiutė, I. Naruševičiūtė, M. Trakymas, and P. Treigys, “Advancements in prostate zone segmentation: integrating attention mechanisms into the nnU-Net framework,” Mach. Learn. Sci. Technol., vol. 5, no. 4, 2024, doi: 10.1088/2632-2153/ad7f24.
[7] Y. Zhao et al., “SCU-Net : Shape Constraint U-Net for Prostate Segmentation in MR Images,” vol. 12032, pp. 1–8, 2022, doi: 10.1117/12.2611327.
[8] Y. Wang et al., “Deep Attentive Features for Prostate Segmentation in 3D Transrectal Ultrasound,” IEEE Trans. Med. Imaging, vol. 38, no. 12, pp. 2768–2778, 2019, doi: 10.1109/TMI.2019.2913184.
[9] C. Zijiang Sang1, a, Yeqin Shao2, b*, Changyan Xu1, “Multi-branch Context Awareness Network for Prostate MRI Segmentation Multi-branch Context Awareness Network for Prostate MRI Segmentation”, doi: 10.1088/1742-6596/2562/1/012008.
[10] Q. Du, L. Wang, and H. Chen, “OPEN A mixed Mamba U ‑ net for prostate segmentation in MR images,” Sci. Rep., pp. 1–12, 2024, doi: 10.1038/s41598-024-71045-7.
[11] A. Gu and T. Dao, “Mamba: Linear-Time Sequence Modeling with Selective State Spaces,” no. Kalman 1960, pp. 1–36, 2024, [Online]. Available: http://arxiv.org/abs/2312.00752
[12] M. Arshad, C. Wang, M. Wajeeh Us Sima, J. A. Shaikh, S. Alkhalaf, and F. Alturise, “RaNet: a residual attention network for accurate prostate segmentation in T2-weighted MRI,” Front. Med., vol. 12, no. 4, 2025, doi: 10.3389/fmed.2025.1589707.
[13] Ö. Çiçek, A. Abdulkadir, S. S. Lienkamp, T. Brox, and O. Ronneberger, “3D U-net: Learning dense volumetric segmentation from sparse annotation,” Lect. Notes Comput. Sci. (including Subser. Lect. Notes Artif. Intell. Lect. Notes Bioinformatics), vol. 9901 LNCS, pp. 424–432, 2016, doi: 10.1007/978-3-319-46723-8_49.
[14] F. Milletari, N. Navab, and S. A. Ahmadi, “V-Net: Fully convolutional neural networks for volumetric medical image segmentation,” Proc. - 2016 4th Int. Conf. 3D Vision, 3DV 2016, pp. 565–571, 2016, doi: 10.1109/3DV.2016.79.
[15] F. Isensee, P. F. Jaeger, S. A. A. Kohl, J. Petersen, and K. H. Maier-Hein, “nnU-Net: a self-configuring method for deep learning-based biomedical image segmentation,” Nat. Methods, vol. 18, no. 2, pp. 203–211, 2021, doi: 10.1038/s41592-020-01008-z.
[16] O. Oktay et al., “Attention U-Net: Learning Where to Look for the Pancreas,” no. Midl, May 2018, doi: https://doi.org/10.48550/arXiv.1804.03999.
[17] L. Rundo et al., “USE-Net: Incorporating Squeeze-and-Excitation blocks into U-Net for prostate zonal segmentation of multi-institutional MRI datasets,” Neurocomputing, vol. 365, pp. 31–43, 2019, doi: 10.1016/j.neucom.2019.07.006.
[18] S. Bhandary et al., “IB-U-Nets: Improving medical image segmentation tasks with 3D Inductive Biased kernels,” 2022, doi: https://doi.org/10.48550/arXiv.2210.15949.
[19] P. C. Quihui-Rubio, D. Flores-Araiza, M. Gonzalez-Mendoza, C. Mata, and G. Ochoa-Ruiz, “FAU-Net: An Attention U-Net Extension with Feature Pyramid Attention for Prostate Cancer Segmentation,” in Lecture Notes in Computer Science (including subseries Lecture Notes in Artificial Intelligence and Lecture Notes in Bioinformatics), 2024, pp. 165–176. doi: 10.1007/978-3-031-47640-2_14.
[20] Y. Bouslimi, T. Ben Aïcha Gader, and A. K. Echi, “Deep Learning-based Detection, Segmentation of Prostate Cancer from mp-MRI Images,” Electron. Lett. Comput. Vis. Image Anal., vol. 22, no. 1, pp. 52–76, 2023, doi: 10.5565/REV/ELCVIA.1620.
[21] L. Coroamă et al., “A Light , 3D UNet-based Architecture for fully Automatic Segmentation of Prostate Lesions from T2-MRI Images Abstract :,” pp. 1–12, 2024, doi: 10.2174/1573405620666230522151445.
[22] S. G. Armato et al., “PROSTATEx Challenges for computerized classification of prostate lesions from multiparametric magnetic resonance images,” J. Med. Imaging, vol. 5, no. 04, p. 1, 2018, doi: 10.1117/1.jmi.5.4.044501.
[23] M. Larobina and L. Murino, “Medical Image File Formats,” J. Digit. Imaging, vol. 27, no. 2, pp. 200–206, 2014, doi: 10.1007/s10278-013-9657-9.
[24] A. L. Simpson et al., “A large annotated medical image dataset for the development and evaluation of segmentation algorithms,” 2019, doi: https://doi.org/10.48550/arXiv.1902.09063 Focus to learn more.
[25] S. Yadav and S. Shukla, “Analysis of k-Fold Cross-Validation over Hold-Out Validation on Colossal Datasets for Quality Classification,” in 2016 IEEE 6th International Conference on Advanced Computing (IACC), 2016, pp. 78–83. doi: 10.1109/IACC.2016.25.
[26] L. R. Dice, “Measures of the Amount of Ecologic Association Between Species,” Ecology, vol. 26, no. 3, pp. 297–302, Feb. 1945, doi: 10.2307/1932409.
[27] X. Li, P. S. Morgan, J. Ashburner, J. Smith, and C. Rorden, “The first step for neuroimaging data analysis: DICOM to NIfTI conversion.,” J. Neurosci. Methods, vol. 264, pp. 47–56, May 2016, doi: 10.1016/j.jneumeth.2016.03.001.
[28] R. T. Shinohara et al., “Statistical normalization techniques for magnetic resonance imaging,” NeuroImage Clin., vol. 6, pp. 9–19, 2014, doi: https://doi.org/10.1016/j.nicl.2014.08.008.
[29] L. G. Nyul, J. K. Udupa, and X. Zhang, “New variants of a method of MRI scale standardization,” IEEE Trans. Med. Imaging, vol. 19, no. 2, pp. 143–150, 2000, doi: 10.1109/42.836373.
[30] G. Wang, W. Li, S. Ourselin, and T. Vercauteren, “Automatic brain tumor segmentation using cascaded anisotropic convolutional neural networks,” Lect. Notes Comput. Sci. (including Subser. Lect. Notes Artif. Intell. Lect. Notes Bioinformatics), vol. 10670 LNCS, pp. 178–190, 2018, doi: 10.1007/978-3-319-75238-9_16.
[31] S. Zhu, C. Yu, and J. Hu, “Regularizing deep neural networks for medical image analysis with augmented batch normalization,” Appl. Soft Comput., vol. 154, p. 111337, 2024, doi: https://doi.org/10.1016/j.asoc.2024.111337.
[32] F. Pedregosa et al., “Scikit-learn: Machine Learning in Python,” J. Mach. Learn. Res., vol. 12, no. null, pp. 2825–2830, 2011, doi: 10.5555/1953048.2078195.
[33] G. Centineo et al., “Prostate MRI Segmentation: A Comparative Analysis of U-Net and E-Net Deep Learning Models,” in Image Analysis and Processing - ICIAP 2025 Workshops, E. Rodolà, F. Galasso, and I. Masi, Eds., Cham: Springer Nature Switzerland, 2026, pp. 299–309.
[34] M. Montazerolghaem, Y. Sun, G. Sasso, and A. Haworth, “U-Net Architecture for Prostate Segmentation: The Impact of Loss Function on System Performance,” Bioengineering, vol. 10, no. 4, 2023, doi: 10.3390/bioengineering10040412.
Copyright (c) 2026 American Journal of Technology Advancement
This work is licensed under a Creative Commons Attribution 4.0 International License.
