A ViTUNeT-based model using YOLOv8 for efficient LVNC diagnosis and automatic cleaning of dataset

2.3.1 YOLOV8 for LV detection

Training a YOLO model for object detection involves labeling the images in a defined format. These labels are saved in accompanying text files, with each line corresponding to an object in the image. The format used is as follows:

Object-class is an integer representing the object class; x is x coordinate of the bounding box center, normalized to image width; y is y coordinate of the bounding box center, normalized to image height; width is the width of the bounding box, normalized to image width and, finally, height is the height of the bounding box, normalized to image height.

To streamline the labeling process, we make use of segmentation masks (ground truth) provided for our MRI slices. For each slice, we first identify the largest contour in the segmentation mask, which outlines the outer boundary of LV. Based on this contour, we then create a bounding box that encapsulates the entire left ventricle complex. The OpenCV library offers functions to detect the largest contour in segmentation masks (see Figure 3).

Figure 3:

OpenCV library is used to find the outermost contour in segmentation masks.

To train the YOLOv8 model, the dataset was randomly split into 80 % for training and 20 % for validation, without grouping images by patient. Since YOLOv8 functions as a region-based detector, enforcing a per-patient split was initially considered unnecessary for this task. However, we acknowledge that multiple slices from the same patient may share morphological similarities, which could lead to an optimistic bias in the evaluation results. We include this as a potential limitation of our approach and plan to explore patient-level splitting in future work to assess its impact on generalization. In addition, no stratification by class (i.e., presence or absence of LVNC) was applied during the split, which could result in class imbalance between training and validation sets. We will consider applying stratified sampling in future studies to ensure more balanced distributions across data splits.

After training, the model is assessed on the test set, achieving an mAP50-95 score of 0.928. This score reflects high accuracy in detecting the left ventricle in short-axis MRI images (see Figure 4).

Figure 4:

Performance of the YOLOv8m model trained on the cardiology dataset.

2.3.2 Dataset cleaning with YOLOv8 models

We suspect that the segmentation models’ accuracy in this study may be limited by the quality of the cardiology dataset we are using.

Our research group, in collaboration with cardiologists from Hospitals Virgen of Arrixaca and Vall d’Hbron in Murcia and Barcelona (Spain), is developing an artificial intelligence-based application that enables the automatic diagnosis of LVNC. A complete magnetic resonance study typically consists of a series of slices, and each slice contains a series of frames. For the diagnosis of LVNC, among other things, cardiologists are particularly interested in the frames corresponding to the tele-diastolic phase (most expanded LV) and the tele-systolic phase (least expanded LV). To avoid specialists from using other applications or relying on their visual judgment to determine these phases, we have trained and introduced a YOLOv8 model into our application to detect the inner cavity of LV, allowing us to measure the area of the bounding box generated by the model. The tele-diastolic phase will have the largest area, and the tele-systolic phase the smallest.

In this line of work, over the past few months and in collaboration with several specialists, we have concluded that certain types of MRI slices are entirely discardable when diagnosing LVNC and may be affecting the training quality of convolutional models. We will classify these slices into the following categories:

1. No Ring: This category includes MRI slices where the external layer of the LV does not form a closed ring around the inner cavity. We can find an example of this type in Figure 5a.

2. Poor Quality: This category includes MRI slices with a small area, typically basal or extreme, with numerous artifacts and poor-quality LV. A slice from this category is shown in Figure 5b.

Figure 5:

Examples of slices that cause issues for the models. (a) MRI slice no ring. (b) MRI slice showing LV with small area and artefacts.

Our goal is to develop a method for automatically detecting MRI slices related to these two classes. To achieve this, we trained two YOLOv8 models:

1. Model Opening-IC-LV: We utilize this model to detect openings in the outer layer surrounding the internal cavity of the left ventricle. To train this model, approximately 400 images of the No Ring class were collected and labeled. The detections made by this model can be seen in Figure 6.

2. Model High-Quality-IC-LV: The purpose of this model is to detect a high-quality internal cavity in the left ventricle. With this model, we can determine when we are observing a well-resolved internal cavity and measure its area using a bounding box. This model is trained using a process similar to the one described in the previous section, utilizing segmentation masks.

Figure 6:

Detections made by model opening-IC-LV on slices of the No Ring class.

Thanks to the Ultralytics library [20], when we apply a YOLO model to an image, we obtain an object with several attributes, including the confidence with which the model believes it has detected what it was trained for. We can also obtain the central coordinates, as well as the width and height of the generated bounding box. In this way, with Model Opening-IC-LV and Model High-Quality-IC-LV trained for their respective purposes, we process the new cardiology dataset with both models and obtain, for each MRI slice, the following values:

1. Conf-Valid-IC: Confidence with which Model High-Quality-IC-LV believes it has detected a valid internal cavity.

2. Conf-Open-IC: Confidence with which Model Opening-IC-LV believes it has detected an opening in the internal cavity (i.e., confidence that it is a No Ring slice).

3. Area-bbox-Valid-IC: Area of the bounding box generated by Model High-Quality-IC-LV when detecting the internal cavity.

Thus, an MRI slice is represented by a triplet of values:

Now, based on these values, our objective is to find a pattern or model capable of determining when an MRI slice is discardable or not. To this end, Figure 7 shows the relationship between values Conf − Valid − IC and Conf − Open − IC for the dataset slices is shown, where the slices manually labeled as No Ring to train model Y are displayed in orange. In this way, Figure 7 shows a linear model that best divides the No Ring slices from the rest of the slices. We also establish a threshold for the confidence of a valid internal cavity, i.e., a threshold for the Conf − Valid − IC values. This confidence threshold is set at 0.88, as we observe that most slices initially classified as No Ring are above this threshold.

Figure 7:

Relationship between Conf-Valid-IC and Conf-Open-IC, and the linear model (red line) that best distinguishes No Ring and poor-quality cuts (low Conf-Valid-IC value) from the rest.

In this way, the linear model from Figure 7 provides the following plane separation:

The slices (Conf − Valid − IC, Conf − Open − IC, Area − bbox − Valid − IC) that satisfy inequality (1.1) are discarded, as they are slices with a high value of Conf − Open − IC (high confidence of being from the No Ring class), others with Conf − Valid − IC = 0 (Model High-Quality-IC-LV does not detect an internal cavity at all), or with Conf − Valid − IC lower than 0.775 (a poor-quality internal cavity). For the slices that satisfy inequality (1.2), we conduct a study on the relationship between values Conf − Valid − IC and Area − bbox − Valid − IC, as there are slices with Conf − Valid − IC lower than the 0.88 threshold, and we want to assess whether they are slices with a small area. Following the previous approach, Figure 8 shows the graphical relationship between these values and a linear model.

Figure 8:

Relationship between values Conf − Valid − IC and Area − bbox − Valid − IC of the MRI slices.

As we can observe, most MRI slices below the 0.88 threshold (Conf − Valid − IC < 0.88) are slices with a small area compared to the rest of the slices (Conf − Valid − IC ≥ 0.88). Also, to avoid being too strict when discarding slices with an Conf − Valid − IC value close to 0.88, we establish a second threshold of 0.83, thus we use the range [0.83 ≤ Conf − Valid − IC ≤ 0.88] to define a second linear model. The resulting model is the red dashed line in Figure 8. This model is defined with a negative slope to compensate for an Conf − Valid − IC value close to 0.83 with a not-so-small area (and therefore fewer artifacts), and thus the slices with a small area that are selected will have an Conf − Valid − IC value close to or greater than 0.88 (an internal cavity with acceptable resolution).

The thresholds and the two linear models used to classify slices were manually selected based on visual inspection of the plotted distributions and expert judgment in collaboration with cardiologists. While this approach has proven effective for the current dataset, we recognize the importance of automating this process. In future work, we plan to approach this as a machine learning classification problem and explore the use of automatic thresholding techniques or classifiers such as decision trees, support vector machines, or ensemble models.

Finally, thanks to the two linear models we have obtained, we are able to design a flowchart (Figure 9) that describes the cardiology dataset cleaning method. Then, we can easily translate this flowchart into an algorithm that processes MRI slices as value triplets, as described earlier, and purges the dataset.

Figure 9:

Flowchart establishing a criterion for automatically discarding slices from the dataset.

The MRI slices to be discarded according to the flowchart amount to 552 (16 % of the total), leaving 2097 images in the purged dataset related to 449 patients, for whom the target TV% is readjusted if any slices were removed. Figure 10 shows how the 552 slices to be discarded are distributed among the different patient groups.

Figure 10:

Number of MRI slices to be discarded for each patient group according to the flowchart in Figure 9.

The justification and evaluation of this assumption will be provided in the results section (Section 3), where we analyze the impact of removing these problematic slices on the overall performance of the model.

2.4.1 Combining YOLOv8 with the ViTUNeT model for LVNC

We describe our method for integrating a YOLOv8 model, trained on our cardiology dataset to detect the left ventricle, with a model using the ViTUNeT architecture. Figure 12 provides a schematic representation of this integration process. We will refer to this method as YOLOv8+ViTUNet.

Figure 12:

Integration of the YOLOv8 model with the ViTUNeT model.

First, the YOLOv8 model processes short-axis MRI images to define a Region of Interest (ROI) by drawing a bounding box around the left ventricle, directing the convolutional segmentation model to focus on this area. To ensure that the generated ROI fully contains the left ventricle, we adjust the side length of the ROI using the probability with which the YOLO model detects a valid left ventricle, according to the following formula:

where:

1. l_resize: The adjusted side length of the final square ROI.

2. w_roi, h_roi: The width and height of the original ROI, respectively.

3. prob_YOLO: The probability assigned by the YOLO model for detecting the left ventricle.

Thus, this formula expands the ROI more significantly for slices where the YOLO model is more difficult to detect.

Next, each ROI is converted to a single grayscale channel and resized to 800 × 800 pixels to standardize input dimensions. Finally, Z-score normalization is applied to normalize the pixel intensity distribution. These preprocessing steps prepare the data for training the ViTUNeT model.

We aim to achieve the following benefits with this approach of using the YOLO model to guide the convolutional model in image segmentation: firstly, to prevent the convolutional model from mistakenly identifying other structures that resemble the LV in MRI slices; secondly, to ensure uniformity in the size of the LV across the MRI images of a patient, as the left ventricle may appear smaller in the extreme slices compared to those in the center of the series.

2.4.2 Training method for ViTUNeT model

The proposed model is trained using images corresponding to unique 2D cardiac MRI slices with a resolution of 800 × 800 pixels. To normalize the values, a Z-score standardization is also applied to each original slice. During training, data augmentation is employed through random rotations of 90°, 180°, or 270° in the images with a probability of 0.25.

The loss function used is a linear combination of two components: the Lovász-Softmax loss [26] ( L L ) and the weighted binary cross-entropy loss ( L B C ) :

This choice of loss function addresses the class imbalance issue, as the trabecular area is smaller than the other areas. The RAdam optimizer [27] minimizes the loss function, with an initial learning rate of 0.005.

Additionally, to more robustly assess the segmentation quality of the models, a 5-fold stratified cross-validation process is employed during training. As described in Figure 13, the complete dataset is divided into 80 % of patients for training and 20 % for testing. Next, 5 folds are created within the training set, where each fold consists of 2 subsets: training (80 %) and validation (20 %). It is important to highlight that the division of the dataset is done “by patients,” not by slices, ensuring that a patient’s slices are not distributed across both the training and validation sets. Stratification is based on the 27.4 % threshold rule described in subsection 2.1. This ensures that the dataset is divided into five equally sized, non-overlapping folds and that the distribution of patients with LVNC and those without the disease is as balanced as possible across the different training and validation sets.

Figure 13:

Generic scheme of our 5-fold cross-validation Process.