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

2026-04-20

AI Summary

This study addresses the limitations of existing biparametric MRI methods for 3D segmentation of prostate lesions, which suffer from insufficient cross-modal fusion and a lack of lesion-level semantic guidance. To overcome these challenges, the authors propose a multi-encoder U-Net architecture incorporating a text-guided alignment mechanism. The framework introduces an alignment loss to enhance foreground semantics, a heatmap loss to suppress background interference, and a confidence-gated multi-head cross-attention module to refine local boundaries. A staged training strategy is employed to effectively integrate multimodal imaging data with vision-language priors. Evaluated on the PI-CAI dataset, the method achieves a new state-of-the-art performance, significantly improving both segmentation accuracy and anatomical consistency.

Technology Category

• Computer Vision: Multi-modal Vision
• Machine Learning: Multimodal Learning
• Natural Language Processing: Language Grounding & Multi-modal NLP

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

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.

Problem

• 3D prostate lesion segmentation
• biparametric MRI
• multi-modal fusion
• vision-language models
• lesion-level semantics

Innovation

• text-guided segmentation
• multi-encoder U-Net
• alignment loss
• heatmap loss
• confidence-gated cross-attention