Promptable segmentation with region exploration enables minimal-effort expert-level prostate cancer delineation

TL;DR

This work tackles the challenge of prostate cancer segmentation on MR images by introducing a promptable segmentation framework driven by reinforcement learning. The method combines a region-growing module with a surrogate voxel probability map and an RL agent that iteratively relocates seed points, guided by a reward that balances Dice improvement and voxel-wise entropy to encourage exploration of uncertain regions. Across PROMIS and PICAI datasets, the approach substantially outperforms fully automated baselines and strong promptable methods, approaching radiologist-level accuracy while reducing annotation effort by about a factor of 10. The results demonstrate the value of sample-specific optimisation via RL in medical image segmentation and provide open-source code for broader adoption and extension to other tasks.

Abstract

Purpose: Accurate segmentation of prostate cancer on magnetic resonance (MR) images is crucial for planning image-guided interventions such as targeted biopsies, cryoablation, and radiotherapy. However, subtle and variable tumour appearances, differences in imaging protocols, and limited expert availability make consistent interpretation difficult. While automated methods aim to address this, they rely on large expertly-annotated datasets that are often inconsistent, whereas manual delineation remains labour-intensive. This work aims to bridge the gap between automated and manual segmentation through a framework driven by user-provided point prompts, enabling accurate segmentation with minimal annotation effort. Methods: The framework combines reinforcement learning (RL) with a region-growing segmentation process guided by user prompts. Starting from an initial point prompt, region-growing generates a preliminary segmentation, which is iteratively refined through RL. At each step, the RL agent observes the image and current segmentation to predict a new point, from which region growing updates the mask. A reward, balancing segmentation accuracy and voxel-wise uncertainty, encourages exploration of ambiguous regions, allowing the agent to escape local optima and perform sample-specific optimisation. Despite requiring fully supervised training, the framework bridges manual and fully automated segmentation at inference by substantially reducing user effort while outperforming current fully automated methods. Results: The framework was evaluated on two public prostate MR datasets (PROMIS and PICAI, with 566 and 1090 cases). It outperformed the previous best automated methods by 9.9% and 8.9%, respectively, with performance comparable to manual radiologist segmentation, reducing annotation time tenfold.

Figures (2)

• Figure 1: An overview of the proposed promptable segmentation using RL.
• Figure 2: Samples from PROMIS segmented using our RL-PromptSeg approach.