From 759f4f025221a1e4d112bfa9930c9b01efc82281 Mon Sep 17 00:00:00 2001 From: Philipp Donn <30521025+phinik@users.noreply.github.com> Date: Mon, 15 Jan 2024 16:18:45 +0100 Subject: [PATCH] fix problem with annotation not in image --- yoeo/scripts/createYOEOLabelsFromTORSO-21.py | 45 +++++++++++--------- 1 file changed, 24 insertions(+), 21 deletions(-) diff --git a/yoeo/scripts/createYOEOLabelsFromTORSO-21.py b/yoeo/scripts/createYOEOLabelsFromTORSO-21.py index ffb0371..443e4b5 100755 --- a/yoeo/scripts/createYOEOLabelsFromTORSO-21.py +++ b/yoeo/scripts/createYOEOLabelsFromTORSO-21.py @@ -34,7 +34,7 @@ def range_limited_float_type_0_to_1(arg): # Available classes for YOEO CLASSES = { - 'bb_classes': ['ball', 'goalpost', 'robot'] if args.robots_with_team_colors else ['ball', 'goalpost', 'robot_blue', 'robot_red', 'robot_unknown'], + 'bb_classes': ['ball', 'goalpost', 'robot'] if not args.robots_with_team_colors else ['ball', 'goalpost', 'robot_blue', 'robot_red', 'robot_unknown'], 'segmentation_classes': ['background', 'lines', 'field'], 'skip_classes': ['obstacle', 'L-Intersection', 'X-Intersection', 'T-Intersection'], } @@ -122,8 +122,12 @@ def range_limited_float_type_0_to_1(arg): annotations = [] for annotation in image_data['annotations']: - class_name = annotation['type'] + # Skip annotations that are not in the image + if not annotation['in_image']: + continue + # Derive the class name of the current annotation + class_name = annotation['type'] if args.robots_with_team_colors and class_name == 'robot': class_name += f"_{annotation['color']}" @@ -134,25 +138,24 @@ def range_limited_float_type_0_to_1(arg): (args.skip_concealed and annotation.get('concealed', False))): continue elif class_name in CLASSES['bb_classes']: # Handle bounding boxes - if annotation['in_image']: # If annotation is not in image, do nothing - min_x = min(map(lambda x: x[0], annotation['vector'])) - max_x = max(map(lambda x: x[0], annotation['vector'])) - min_y = min(map(lambda x: x[1], annotation['vector'])) - max_y = max(map(lambda x: x[1], annotation['vector'])) - - annotation_width = max_x - min_x - annotation_height = max_y - min_y - relative_annotation_width = annotation_width / img_width - relative_annotation_height = annotation_height / img_height - - center_x = min_x + (annotation_width / 2) - center_y = min_y + (annotation_height / 2) - relative_center_x = center_x / img_width - relative_center_y = center_y / img_height - - # Derive classID from index in predefined classes - classID = CLASSES['bb_classes'].index(class_name) - annotations.append(f"{classID} {relative_center_x} {relative_center_y} {relative_annotation_width} {relative_annotation_height}") + min_x = min(map(lambda x: x[0], annotation['vector'])) + max_x = max(map(lambda x: x[0], annotation['vector'])) + min_y = min(map(lambda x: x[1], annotation['vector'])) + max_y = max(map(lambda x: x[1], annotation['vector'])) + + annotation_width = max_x - min_x + annotation_height = max_y - min_y + relative_annotation_width = annotation_width / img_width + relative_annotation_height = annotation_height / img_height + + center_x = min_x + (annotation_width / 2) + center_y = min_y + (annotation_height / 2) + relative_center_x = center_x / img_width + relative_center_y = center_y / img_height + + # Derive classID from index in predefined classes + classID = CLASSES['bb_classes'].index(class_name) + annotations.append(f"{classID} {relative_center_x} {relative_center_y} {relative_annotation_width} {relative_annotation_height}") else: print(f"The annotation type '{class_name}' is not supported. Image: '{img_name_with_extension}'")