Machine Learning: Science and Technology (Jan 2024)
Advancements in prostate zone segmentation: integrating attention mechanisms into the nnU-Net framework
Abstract
Prostate cancer is one of the most lethal cancers in the world. Early diagnosis is essential for successful treatment of prostate cancer. Segmentation of prostate zones in magnetic resonance images is an important task in the diagnosis of prostate cancer. Currently, the state-of-the-art method for this task is no-new U-Net. In this paper, a method to incorporate the attention U-Net architecture into no-new U-Net is proposed and compared with a classical U-net architecture as research. The experimental results indicate that there is no significant statistical difference between the proposed modification of no-new U-Net with the generalizability of the attention mechanism or the ability to achieve more accurate results. Moreover, two novel workflows are proposed for prostate segmentation, transitional zone segmentation and peripheral zone calculation workflow, and separate models for peripheral zone and transitional zone segmentation workflow. These workflows are compared with a baseline single peripheral zone and transitional zone segmentation model workflow. The experimental results indicate that separate models for peripheral zone and transitional zone segmentation workflow generalizes better than the baseline between data sets of different sources. In peripheral zone segmentation separate models for peripheral zone and transitional zone segmentation workflow achieves 1.9% higher median Dice score coefficient than the baseline workflow when using the attention U-Net architecture and 5.6% higher median Dice score coefficient when using U-Net architecture. Moreover, in transitional zone segmentation separate models for peripheral zone and transitional zone segmentation workflow achieves 0.4% higher median Dice score coefficient than the baseline workflow when using the attention U-Net architecture and 0.7% higher median Dice score coefficient when using U-Net architecture. Meanwhile, prostate segmentation, transitional zone segmentation and peripheral zone calculation workflow generalizes worse than the baseline. In peripheral zone segmentation prostate segmentation, transitional zone segmentation and peripheral zone calculation workflow achieves 4.6% lower median Dice score coefficient than the baseline workflow when using the attention U-Net architecture and 3.6% lower median Dice score coefficient when using U-Net architecture. In transitional zone segmentation prostate segmentation, transitional zone segmentation and peripheral zone calculation workflow achieves a similar median Dice score coefficient to the baseline workflow.
Keywords
