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

INTRODUCTION

Accurate heart segmentation in medical imaging is important for precise diagnosis and effective treatment planning, especially in cases of congenital heart disease (CHD). CHD includes a variety of structural anomalies present at birth, affecting the heart’s architecture and function. The complexity of these malformations requires detailed analysis of cardiac anatomy to guide clinical decisions,^[1,2] and makes accurate segmentation essential for understanding the spatial relationships between structures and planning surgical or interventional procedures.^[3,4]

Accurate segmentation maps allow the reconstruction of 3D heart models that support surgical planning and simulation by clarifying anatomical relationships and aiding preoperative decision-making,^[5,6] which enhances surgical outcomes.^[7,8] Moreover, it is essential for quantitatively assessing cardiac function, such as measuring ventricular volumes and ejection fractions,^[9] which support disease monitoring and evaluation of therapeutic efficacy.^[10,11]

Despite its importance, whole-heart segmentation remains challenging due to the variability of cardiac anatomy, particularly in CHD cases,^[12] as well as imaging issues such as noise, low contrast, and artifacts in computed tomography (CT) and magnetic resonance imaging (MRI).^[13] Furthermore, manual segmentation, the current gold standard, is labor-intensive, time-consuming, and susceptible to variations between observers, making it impractical for routine clinical workflows.^[14]

Deep learning has advanced cardiac segmentation considerably. Convolutional neural networks (CNNs), particularly U-Net and its variants, have achieved state-ofthe-art performance in segmenting cardiac structures from 3D medical images.^[15,16] However, their complexity means these methods are computationally expensive and often unsuitable for real-time applications.^[17]

Recently, the you only look once (YOLO) family of models has garnered attention in computer vision due to their balance of speed and accuracy in object detection tasks.^[18] Though initially developed for 2D detection problems, YOLO has been successfully adapted for several medical applications, including tumor detection, organ localization, and lesion classification.^[19,20] Its efficiency suggests potential for time-sensitive clinical use in cardiac imaging. Recent studies have explored its utility in various 3D imaging contexts, such as lung nodule detection and brain tumor segmentation, suggesting that similar approaches could be applied to cardiac imaging.^[21,22]

This study aims to explore the feasibility of a YOLO-based framework for whole-heart segmentation in 3D CT scans, particularly in CHD cases. Using YOLO for localization and segmentation, we aim to evaluate whether YOLO can provide a lightweight alternative to traditional segmentation networks in scenarios where rapid inference and reduced computational demand are essential.

The rest of the paper is organized as follows. Section 2 reviews existing heart segmentation methods, including both traditional techniques and recent deep learning approaches, with a focus on CHD cases. Section 3 introduces our proposed method, highlighting the application of YOLO in medical imaging, identifying current research gaps, and comparing our approach with 3D segmentation methods. Section 4 outlines the materials and experimental setup. Section 5 presents the results, showcasing the model’s performance through quantitative metrics and visual examples. Finally, Section 6 concludes the study with key findings and future directions.

MATERIAL AND METHODS

Experimental setup

In this work, the methodology was structured into two distinct phases:

Phase 1: Heart localization

The first phase of this work focused on heart localization, which is a critical step for reducing computational complexity and improving segmentation accuracy. A YOLOv8n object detection model was used to localize the heart and its great arteries within CT scans. YOLOv8n is a lightweight variant of YOLOv8, designed for efficient object detection with minimal computational overhead.

A custom data preparation process was implemented, wherein for each CT scan, an algorithm calculated the ground truth bounding box of the heart based on its label. Subsequently, we scanned through each slice of the CT scan along all axes, saving those slices as images when heart pixels constituted more than 30% of the total pixels in the slice. This approach focused on relevant data, eliminating non-essential information.

Following this preparation, we obtained a total of 8940 training samples, with an additional 1873 samples reserved for validation and 794 for testing. After the data preparation phase, the YOLOv8n model was trained for 500 epochs on a batch size of 64 images per iteration using the free hardware configuration provided by Google Colaboratory, which includes an Intel(R) Xeon(R) CPU @ 2.00GHz, a Tesla T4 Graphics processing unit (GPU) with 16 GB of VRAM, and 12 GB of RAM.

Phase 2: Heart segmentation

The second phase of this work focused on heart segmentation. For this stage, we wanted to test the improvements made in the YOLOv11 model; therefore, the YOLOv11n-seg, YOLOv11s-seg, and YOLOv11m-seg variants were chosen to compare their performance. These models are segmentation variants of YOLOv11, designed to predict pixel-wise masks in addition to bounding boxes.

To prepare the data for segmentation, we transformed the heart masks from the ImageCHD dataset after cropping it from the scans in the first phase (heart localization) into YOLO annotation format. This involved converting the masks for each class (heart parts) into polygon coordinates, which were then used to train the models. The models were trained with the goal of detecting and segmenting seven key anatomical heart regions: left and right atria, left and RVs, Myo, Ao, and PA, from the background, and each model’s training process was monitored through key metrics, including segmentation loss, box loss, classification loss, and detection-specific metrics such as mean average precision (mAP)@50, mAP@50-95, precision, and recall for both bounding box (B) and mask (M) predictions.

After the data preparation phase, a total of 76362 training images and 9420 images for the validation set were obtained. We trained the YOLOv11n-seg and YOLOv11s-seg models for 500 epochs on a batch size of 144 and 80 images per iteration, respectively. Due to time constraints, we train YOLOv11m-seg for 224 epochs on a batch size of 48 images. The hardware configuration used in this phase is the one provided by Kaggle with GPU T4 × 2 accelerator enabled in order to exploit the multi-GPU feature included in YOLO models training. This phase aimed to achieve precise segmentation of the heart structures within the localized regions identified in the first phase. This two-phase approach leverages the strengths of the YOLO architecture in both localization and segmentation tasks, designed to address the complexities inherent in CHD cases. Figure 1 shows an overview of our approach, while the flowchart depicted in Figure 2 provides a graphical representation of our methodology.

Close

Figure 1:

Overview of the used method where the red box contains the heart area.

Export to PPT

Close

Figure 2:

Methodology flowchart.

Export to PPT

RESULTS

For the heart and great arteries localization phase, the results showed that the model achieved satisfactory accuracy in localizing the heart in CT scan images, demonstrating the effectiveness of both YOLOv8n architecture and the data preparation strategy in this study. The model performance has been evaluated using mAP metrics at an intersection over union (IoU), where the model gives 99.5% for mAP50 and 81.168 % for mAP50–95.

The results are illustrated in Figure 3, and in Figure 4, we can see the evaluation results on a random batch of the validation set.

Close

Figure 3:

The final results after you only look oncev8n training for 500 epochs (the x axis) (box_loss, cls_loss, and dfl_loss) for both training and validation sets and precision, recall, mAP50, mAP50-95 metrics (the y axis). Blue dots represent loss/metric value for each epoch, red dots represent the smothed curve

Export to PPT

Close

Figure 4:

You only look oncev8n evaluation results where (a) a random batch from the validation set and its ground truth. (b) The model prediction on the validation set batch.

Export to PPT

For the segmentation phase, due to the complexity of the task, the results were low in terms of segmentation accuracy.

The results of YOLOv11n-seg, YOLOv11s-seg and YOLOv11m-seg training are respectively illustrated in Figures 5-7 that show the box-loss, seg-loss, cls-loss, and dfl-loss curves for both training and validation sets, and also show the precision, recall, mAP50, mAP50-95 metrics curves for both objects detection (metrics/precision[B], metrics/recall[B], metrics/mAP50[B], metrics/mAP50-95[B] respectively) and segmentation (metrics/precision[M], metrics/recall[M], metrics/mAP50[M], metrics/mAP50-95[M] respectively). For YOLOv11n-seg and from Figure 5, all training losses show a good decreasing pattern throughout the 500 epochs, indicating that the model is effectively learning from the training data. The curves relatively flatten out, suggesting convergence on the training set. For the metrics curves, training metrics consistently improve, reflecting the loss decrease. Precision, Recall, and mAP for both bounding boxes (B) and masks (M) increase and start to plateau, confirming learning, even if it is still low. On the other hand, the validation behavior shows that the validation box loss remains relatively flat after an initial adjustment. The classification loss is also relatively stable. However, the distribution focal loss for validation-loss (distribution focal loss, related to bounding box refinement for validation samples) shows a slight but steady increase after epoch ~150.

Close

Figure 5:

Training results for you only look oncev11n-seg where the x axis represent the epochs and the y axis represent loss/metrics for each epoch. Blue dots represent loss/metric value for each epoch, red dots represent the smothed curve

Export to PPT

Close

Figure 6:

Training results for you only look oncev11s-seg where the x axis represent the epochs and the y axis represent loss/metrics for each epoch. Blue dots represent loss/metric value for each epoch, red dots represent the smothed curve

Export to PPT

Close

Figure 7:

Training results for you only look oncev11m-seg where the x axis represent the epochs and the y axis represent loss/metrics for each epoch. Blue dots represent loss/metric value for each epoch, red dots represent the smothed curve.

Export to PPT

The most concerning plot is the val/seg-loss curve, where the validation segmentation loss decreases initially but then starts to increase significantly from around epoch 100 onwards. This is a classic indicator of overfitting, specifically on the segmentation task. The model is learning intricate details like small and adjacent parts or noise specific to the training set’s segmentation masks that do not generalize to the unseen validation data.

For validation metrics, we can observe that they increase initially but plateau relatively early (around epoch 100–150) and remain stagnant or even slightly decrease afterward. The final validation mAP50-95 for masks is around 0.25, which is quite low for segmentation, suggesting poor pixel-level accuracy and difficulty with varying IoU thresholds. The mAP50 is higher (~0.45), indicating the model is better at getting the general location and shape correct but struggles with precise boundary delineation. The divergence between steadily improving training metrics and plateauing/worsening validation metrics/losses confirms overfitting.

DISCUSSION

Similar to the YOLOv11n-seg model, all training losses decrease steadily by looking to Figure 6, and training metrics improve throughout, indicating effective learning on the training set. The final training metrics (e.g., precision, mAP) seem slightly higher than the “n” model, which is expected given the larger capacity of the “s” variant. For the validation part the val/box-loss is perhaps slightly more stable than in the nano model. However, the critical val/seg-loss shows the exact same pattern of initially decreasing and then sharply increasing after ~150 epochs, indicating significant overfitting on segmentation, just like the “n” model. The val/dfl-loss also increases, suggesting overfitting in localization refinement. The validation metrics plateau similarly to the “n” model, perhaps reaching slightly higher peak values (e.g., mask mAP50-95 around 0.27 versus 0.25, mask mAP50 around 0.44 versus 0.45 - very similar). Despite the slightly better peak, the overall pattern and the final performance are comparable to the “n” model, and the overfitting issue persists strongly. The gap between training and validation performance remains large.

As shown in Figure 7, we can observe that, as with the previous models, all training losses decrease consistently throughout the 224 epochs, indicating effective learning and convergence on the training dataset. The training metrics for both bounding boxes (B) and masks (M) show steady improvement, plateauing towards the end of training. The peak training metrics are comparable or slightly higher than the “s” model, as expected from the increased capacity. However, the model is still suffering from the overfitting problem on the segmentation task that is not resolved by increasing model size to “m”.

To understand more about models’ performance, we use their confusion matrices, while metrics like mAP provide an overall score, the confusion matrix offers a wider view of performance for each individual class. Furthermore, it helps in identifying specific misclassification patterns by revealing which classes are being confused with each other. Moreover, it provides the number of these errors that help us in comparing the models and diagnosing their weaknesses for future improvements. Figures 8-10 represent confusion matrices for YOLOv11n-seg, YOLOv11s-seg, and YOLOv11m-seg, respectively.

Close

Figure 8:

Confusion matrices for you only look oncev11n-seg model.

Export to PPT

Close

Figure 9:

Confusion matrices for you only look oncev11s-seg model.

Export to PPT

Close

Figure 10:

Confusion matrices for you only look oncev11m-seg model.

Export to PPT

The confusion matrix of the YOLOv11n-seg model in Figure 8 provides insight into how well the network is able to classify pixels across different anatomical classes. The diagonal elements, which represent the true positives, show a reasonably accurate classification for several heart structures. Among these, the LA and PA stand out with the highest number of correctly identified pixels. Other structures, such as the LV, RV, RA, Myo, and AO, also demonstrate adequate correct classification, despite the lower pixel counts. However, a major issue is revealed in the off-diagonal elements, particularly concerning confusion with the background. A large number of pixels belonging to actual cardiac structures are misclassified as background, which is clearly evident from the last row of the matrix. For example, 24307 pixels from the PA are wrongly identified as background. Furthermore, the last column of the matrix highlights a critical false-positive issue, where background pixels are misclassified as various structures, such as 11093 background pixels predicted as PA and 10656 as LA. This suggests that the model has significant difficulty in effectively distinguishing small heart regions from the surrounding background, resulting in both false negatives (missed structures) and false positives (incorrectly predicted structures).

In addition, there is notable inter-class confusion between anatomically adjacent structures. For instance, the LA is misclassified as PA in 3186 pixels, and PA is confused with LA in 5189 pixels. These errors imply that the model struggles with accurately delineating boundaries between nearby regions, likely due to subtle intensity variations or overlapping textures within the CT images. Overall, the YOLOv11n-seg model demonstrates a lack of precision in separating both the foreground structures and the background, indicating a need for better boundary modeling and spatial awareness.

As shown in Figure 9, the confusion matrix for the YOLOv11s-seg model exhibits similar behavior to its smaller counterpart, the “n” variant. The diagonal entries representing correctly classified pixels remain comparable, with some classes showing slight improvements. For example, PA classification improves marginally (18635 vs. 19852), while other classes, such as RV show a minor drop in correct predictions (5371 vs. 5185). These variations suggest small quantitative differences but no significant architectural advantage.

The issue of background confusion persists across this model as well. Although certain improvements are observed, such as fewer misclassified LV (4268 vs. 4253) and Myo (7194 vs. 7168) pixels, the number of LA and PA pixels predicted as background decreases notably (20831 vs. 20276 and 24307 vs. 23438, respectively). Background pixels misclassified as structures remain high and almost identical to the YOLOv11n-seg model. This reinforces the observation that the model, despite its increased capacity, continues to struggle with segmenting small, low-contrast, or spatially ambiguous boundaries between structures and background. Patterns of inter-structure confusion also remain consistent. The PA is still misclassified as LA in 4880 instances (slightly improved from 5189), while LA is misclassified as PA even more frequently (3186 vs. 3441). These consistent misclassifications suggest that increasing model complexity alone is insufficient to resolve ambiguities caused by overlapping anatomical features or limited contextual awareness.

From Figure 10 observation, we find that YOLOv11mseg, despite its increased capacity, suffers from the exact same weaknesses as the smaller variants, where it has the same difficulty in separating foreground small parts from background and confusion between adjacent structures without much improvement over previous models.

CONCLUSION

This study introduces a YOLO-based pipeline for wholeheart segmentation in CHD cases, combining the strengths of YOLOv8 for localization and YOLOv11-seg for mask prediction. For heart and its great vessels localization, YOLOv8n gives 99.5% for mAP50 and 81.168% for mAP50-95, which is considered a satisfactory accuracy that helps a lot in CHD automatic diagnosis and shows the potential of YOLO models in medical assisting tasks, and of course, there is room for future improvement. However, the segmentation task was not as successful as the previous phase, as the models demonstrated effective training behavior, as indicated by steadily decreasing training losses and improving performance metrics, but evaluation on the validation set revealed notable limitations. A key challenge observed was overfitting in the segmentation; at the same time, the validation mask mAP plateaued at low levels. The confusion matrix analysis revealed significant challenges in accurately distinguishing between foreground cardiac structures and background regions, suggesting difficulties in boundary detection, low-contrast region handling, or insufficient model confidence. Furthermore, there was notable misclassification between anatomically adjacent or visually similar regions, such as the left and right atria, highlighting insufficient fine structural discrimination. Scaling the model led to only minor gains and failed to resolve key issues such as overfitting and segmentation inaccuracies, suggesting limitations in architecture or training strategy. These findings suggest that, in its current form, YOLO is not competitive with benchmark segmentation frameworks such as nnUNet and 3D U-Net. However, the potential value of YOLO lies in its lightweight architecture and speed, which may be relevant for applications requiring rapid inference or when computational resources are constrained. This exploratory study highlights the trade-offs between efficiency and accuracy and underscores the need for further development, such as incorporating regularization strategies, improved data augmentation, or hybrid architectures that combine the strengths of detection and segmentation networks.