From 6b52269dbe13f673f70e5670545eb35716290424 Mon Sep 17 00:00:00 2001 From: luuzk <44912180+luuzk@users.noreply.github.com> Date: Fri, 21 Feb 2025 11:50:02 +0100 Subject: [PATCH 1/2] Update dataset_utils.py --- yolo/utils/dataset_utils.py | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/yolo/utils/dataset_utils.py b/yolo/utils/dataset_utils.py index dd9a66ab..ce7c170f 100644 --- a/yolo/utils/dataset_utils.py +++ b/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) From 85bf2f24aeff67212c0c87aa1aadd96ae72b5eac Mon Sep 17 00:00:00 2001 From: luuzk <44912180+luuzk@users.noreply.github.com> Date: Fri, 21 Feb 2025 11:52:27 +0100 Subject: [PATCH 2/2] Add YOLO detection to segmentation format conversion --- yolo/tools/data_loader.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/yolo/tools/data_loader.py b/yolo/tools/data_loader.py index c44f00c6..b1eefc96 100644 --- a/yolo/tools/data_loader.py +++ b/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 = []