Uncertainty-aware Diffusion and Reinforcement Learning for Joint Plane Localization and Anomaly Diagnosis in 3D Ultrasound

2025-06-30

AI Summary

This study addresses intelligent diagnosis of congenital uterine anomalies (CUAs) by jointly optimizing automatic coronal-plane localization and anomaly classification in 3D ultrasound. We propose a novel framework integrating denoising diffusion models with reinforcement learning: (1) a local-global guided, text-conditioned diffusion model ensures robust coronal-plane localization; (2) a key-frame reinforcement learning strategy with unsupervised reward design adaptively selects diagnostically critical slices; and (3) uncertainty-aware probabilistic calibration coupled with adaptive attention weighting enables effective multi-planar feature fusion and enhances classification reliability. Evaluated on a large-scale clinical 3D uterine ultrasound dataset, our method reduces coronal-plane localization error by 32.7% and achieves 96.4% classification accuracy for CUAs—substantially outperforming state-of-the-art approaches. The source code is publicly available.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Knowledge Representation and Reasoning: Diagnosis and Abductive Reasoning
• Computer Vision: Diffusion Models for Vision

Application Category

• Responsible Web: Machine-in-the-loop, human agency and autonomy
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation

Abstract

Congenital uterine anomalies (CUAs) can lead to infertility, miscarriage, preterm birth, and an increased risk of pregnancy complications. Compared to traditional 2D ultrasound (US), 3D US can reconstruct the coronal plane, providing a clear visualization of the uterine morphology for assessing CUAs accurately. In this paper, we propose an intelligent system for simultaneous automated plane localization and CUA diagnosis. Our highlights are: 1) we develop a denoising diffusion model with local (plane) and global (volume/text) guidance, using an adaptive weighting strategy to optimize attention allocation to different conditions; 2) we introduce a reinforcement learning-based framework with unsupervised rewards to extract the key slice summary from redundant sequences, fully integrating information across multiple planes to reduce learning difficulty; 3) we provide text-driven uncertainty modeling for coarse prediction, and leverage it to adjust the classification probability for overall performance improvement. Extensive experiments on a large 3D uterine US dataset show the efficacy of our method, in terms of plane localization and CUA diagnosis. Code is available at https://github.com/yuhoo0302/CUA-US.

Problem

• Automated localization of uterine planes in 3D ultrasound
• Accurate diagnosis of congenital uterine anomalies (CUAs)
• Uncertainty-aware modeling for improved classification performance

Innovation

• Denoising diffusion model with adaptive guidance
• Reinforcement learning for key slice extraction
• Text-driven uncertainty modeling for prediction