Deep EM with Hierarchical Latent Label Modelling for Multi-Site Prostate Lesion Segmentation
Authors: Wen Yan, Yipei Wang, Shiqi Huang, Natasha Thorley, Mark Emberton, Vasilis Stavrinides, Yipeng Hu, Dean Barratt
Organizations: UCL Hawkes Institute; Department of Medical Physics and Biomedical Engineering, University College London, Gower St, London WC1E 6BT, U.K. · Centre for Medical Imaging, Division of Medicine, University College London, Gower St, London WC1E 6BT, U.K. · Division of Surgery and Interventional Science, University College London, 10 Pond St, London NW3 2PS, U.K. · Cancer Institute, Urology Department, UCLH, University College London, 16-18 Westmoreland St, London W1G 8PH, U.K. · Radiology Department, Imperial College Healthcare, The Bays, S Wharf Rd, London W2 1NY, U.K.
Label variability is a major challenge for prostate lesion segmentation. In multi-site datasets, annotations often reflect centre-specific contouring protocols, causing segmentation networks to overfit to local styles and generalise poorly to unseen sites in inference. We treat each observed annotation as a noisy observation of an underlying latent 'clean' lesion mask, and propose a hierarchical expectation-maximisation (HierEM) framework that alternates between: (1) inferring a voxel-wise posterior distribution over the latent mask, and (2) training a CNN using this posterior as a soft target and estimate site-specific sensitivity and specificity under a hierarchical prior. This hierarchical prior decomposes label-quality into a global mean with site- and case-level deviations, reducing site-specific bias by penalising the likelihood term contributed only by site deviations. Experiments on three cohorts demonstrate that the proposed hierarchical EM framework enhances cross-site generalisation compared to state-of-the-art methods. For pooled-dataset evaluation, the per-site mean DSC ranges from 29.50% to 39.69%; for leave-one-site-out generalisation, it ranges from 27.91% to 32.67%, yielding statistically significant improvements over comparison methods (p<0.039). The method also produces interpretable per-site latent label-quality estimates (sensitivity alpha ranges from 31.5% to 47.3% at specificity beta approximates 0.99), supporting post-hoc analyses of cross-site annotation variability. These results indicate that explicitly modelling site-dependent annotation can improve cross-site generalisation.
Figures & tables
Figure 1: Overview of proposed deep mixed EM framework with hierarchical latent label-quality parameters. The E-step infers a latent clean mask, and the M-step updates the segmentation network and site-specific sensitivity and specificity.
Methods
Site 1
Site 2
Site 3
Dice(%)
HD95(mm)
Dice(%)
HD95(mm)
Dice(%)
HD95(mm)
Split A: Pooled patient-level held-out split
UNet
38.63±10.41
17.04±2.96
27.47±9.62^*
21.25±4.50
33.49±6.98^*
22.06±3.92
Bootstrap
38.49±9.94
16.97±3.14
26.69±10.59^*
22.29±5.12^*
33.14±7.05^*
21.41 ± 5.84
Site-EM
36.46±11.72^*
18.41±7.07^*
23.74±9.73^*
22.99±5.52^*
34.62±6.15
21.86± 4.91
HierEM
39.69 ± 11.88
16.92 ± 3.29
29.50 ± 9.13
20.46 ± 5.58
35.60 ± 7.85
21.66±4.11
Table 1: Sites 1–3 contained 252, 87, and 265 cases in the pooled patient-level held-out set, and 1201, 391, and 1265 cases in the LOSO experiments, respectively. ∗ denote HierEM is significantly better than other baselines based on statistical tests with a significance level of α=0.05 .
Figure 2: Left: Risk-coverage curves on three-site datasets under LOSO experiments, where smaller AURC represents lower risk; Right: Estimated sensitivity of baselines and HierEM with spec≃0.99 . Such high specificity is expected because lesions occupy only a small fraction of the image.
Medical image segmentation delineates organs, lesions, and other structures in scans to support measurement, diagnosis, treatment planning, and research.
Automated 3D segmentation of prostate lesions from biparametric MRI (bp-MRI) is essential for reliable algorithmic analysis, but achieving high precision remains challenging. Volumetric methods must combine multiple modalities while ensuring anatomical consistency, but current models struggle to integrate cross-modal information reliably. While vision-language models (VLMs) are replacing the currently used architectural designs, they still lack the fine-grained, lesion-level semantics required for effective localized guidance. To address these limitations, we propose a new multi-encoder U-Net architecture incorporating three key innovations: (1) an alignment loss that enhances foreground text-image similarity to inject lesion semantics; (2) a heatmap loss that calibrates the similarity map and suppresses spurious background activations; and (3) a final-stage, confidence-gated multi-head cross-attention refiner that performs localized boundary edits in high-confidence regions. A phase-scheduled training regime stabilizes the optimization of these components. Our method consistently outperforms prior approaches, establishing a new state-of-the-art on the PI-CAI dataset through enhanced multi-modal fusion and localized text guidance. Our code is available at https://github.com/NUBagciLab/Prostate-Lesion-Segmentation.
Cuiling Sun, Linkai Peng, Adam Murphy +9
Department of Radiology, Northwestern University, Chicago, USA · Department of Urology, Northwestern University, Chicago, USA · Center for Cancer Research, National Cancer Institute, Bethesda, USA
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.
Junqing Yang, Natasha Thorley, Ahmed Nadeem Abbasi +4
UCL Hawkes Institute; Department of Medical Physics and Biomedical Engineering, University College London, UK. · Centre of Medical Imaging, University College London, UK. · Department of Oncology, Aga Khan University Hospital, Pakistan. +1
Accurate lesion segmentation in mammography is important for computer-aided breast cancer diagnosis, but dense pixel-level annotations remain costly and limited. Existing public segmentation datasets usually cover only a small set of lesion types with manually delineated masks, whereas larger mammography datasets often provide only image-level diagnostic labels. This annotation mismatch makes it difficult to train segmentation models that can generalize to rare or weakly annotated lesion categories. To address this problem, we study open-vocabulary mammographic lesion segmentation under a mixed-supervision setting, where pixel-level masks are available for base lesion categories and image-level labels are available for a broader set of categories. We propose a two-branch framework that jointly learns from segmentation data and classification data. The segmentation branch learns lesion boundaries from annotated masks, while the classification branch introduces additional image-level semantic supervision. To better transfer image-level supervision to lesion-relevant regions, we further introduce a Lesion-wise Attentive Classifier that maintains category-specific feature prototypes and generates lesion-aware attention maps for foreground feature pooling. This design reduces the influence of background tissue and encourages more compact lesion representations. Experiments on multiple mammography datasets show that the proposed method consistently improves open-vocabulary segmentation baselines on both base and novel lesion categories, while reducing the gap to fully supervised models that require dense target-domain masks. Qualitative results further indicate that our method produces more complete lesion localization for small, low-contrast, and visually subtle mammographic findings.
Di Dai, Bo Liu, Liwei Wang
State Key Laboratory of General Artificial Intelligence, School of Intelligence Science and Technology, Peking University, Beijing, China