Align then Refine: Text-Guided 3D Prostate Lesion Segmentation

TL;DR

This paper introduces a novel multi-encoder U-Net architecture for 3D prostate lesion segmentation from MRI, leveraging text-guided alignment, heatmap calibration, and localized boundary refinement to improve accuracy.

Contribution

It proposes a new multi-encoder U-Net with innovative loss functions and a confidence-gated refiner, achieving state-of-the-art results in prostate lesion segmentation.

Findings

Outperforms previous methods on the PI-CAI dataset

Enhances multi-modal fusion with text-guided alignment

Improves localized boundary accuracy with confidence gating

Abstract

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.