🩹 [Fix] Labels in YOLO detection format

yolo/tools/data_loader.py

@@ -16,6 +16,7 @@
 from yolo.tools.data_augmentation import AugmentationComposer
 from yolo.tools.dataset_preparation import prepare_dataset
 from yolo.utils.dataset_utils import (
+    convert_bboxes,
     create_image_metadata,
     locate_label_paths,
     scale_segmentation,
@@ -104,7 +105,8 @@ def filter_data(self, dataset_path: Path, phase_name: str, sort_image: bool = Fa
                 if not label_path.is_file():
                     continue
                 with open(label_path, "r") as file:
-                    image_seg_annotations = [list(map(float, line.strip().split())) for line in file]
+                    annotations = [list(map(float, line.strip().split())) for line in file]
+                    image_seg_annotations = convert_bboxes(annotations)
             else:
                 image_seg_annotations = []
 

yolo/utils/dataset_utils.py

@@ -116,6 +116,35 @@ def scale_segmentation(
     return seg_array_with_cat
 
 
+def convert_bboxes(
+        annotations: list[list[float]],
+) -> list[list[float]]:
+    """
+    Converts annotations in YOLO detection format (class_id, cx, cy, w, h) or YOLO segmentation format \
+        (class_id, x1, y1, x2, y2, ..., xn, yn) to YOLO segmentation format.
+
+    Args:
+        annotations (list[list[float]]): List of annotations in any YOLO format.
+
+    Returns:
+        list[list[float]]: List of annotations in any YOLO segmentation format.
+    """
+    segmentation_data = []
+
+    for anno in annotations:
+        # YOLO segmentation format
+        if len(anno) > 5:
+            segmentation_data.append(anno)
+            continue
+
+        # YOLO detection format
+        category_id, cx, cy, w, h = anno
+        x1, y1, x2, y2 = cx - w / 2, cy - h / 2, cx + w / 2, cy + h / 2
+        segmentation_data.append([category_id, x1, y1, x2, y1, x2, y2, x1, y2])
+
+    return segmentation_data
+
+
 def tensorlize(data):
     try:
         img_paths, bboxes, img_ratios = zip(*data)