Deep Learning-based Heart Localization and Segmentation for Congenital Heart Disease Diagnosis Using You Only Look Once

Congenital heart disease (CHD) presents significant diagnostic challenges due to complex anatomical variations. Accurate whole-heart segmentation from 3D computed tomography (CT) is important for treatment planning but remains difficult. This paper introduces and evaluates a two-phase deep learning pipeline leveraging you only look once (YOLO) architectures for efficient heart localization and segmentation in CHD cases using the ImageCHD dataset. The first phase uses YOLOv8n for heart localization, achieving high accuracy with 99.5% mean average precision (mAP)@50 and 81.168% mAP@50-95 by utilizing a custom slice-filtering data preparation strategy. The second phase uses YOLOv11-seg variants (n, s, m) for pixel-wise segmentation of seven cardiac structures within the localized regions. While training metrics indicated effective learning, validation results revealed significant limitations for the segmentation task across all the models. Key challenges included overfitting, evidenced by increasing validation loss and low mask mAP@50-95 (plateauing around 0.26–0.27), and difficulty in distinguishing foreground small structures from background, confirmed by confusion matrix analysis. Notably, increasing model size did not resolve these core issues. Despite the segmentation challenges, this study demonstrates the strong potential of YOLOv8 for rapid medical object localization and explores the feasibility and potential of YOLOv11-seg for whole-heart segmentation in CHD. Future work should focus on advanced augmentation, regularization, and potentially alternative architectures to improve segmentation robustness for clinical applicability.