From b580a77cf0e9bc4de5b9d8ddd1f49151c2072e59 Mon Sep 17 00:00:00 2001 From: Philipp Donn <30521025+phinik@users.noreply.github.com> Date: Mon, 15 Jan 2024 15:10:59 +0100 Subject: [PATCH 1/3] added team color support --- yoeo/scripts/createYOEOLabelsFromTORSO-21.py | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/yoeo/scripts/createYOEOLabelsFromTORSO-21.py b/yoeo/scripts/createYOEOLabelsFromTORSO-21.py index ff74d22..fd1123f 100755 --- a/yoeo/scripts/createYOEOLabelsFromTORSO-21.py +++ b/yoeo/scripts/createYOEOLabelsFromTORSO-21.py @@ -11,8 +11,9 @@ # Available classes for YOEO CLASSES = { 'bb_classes': ['ball', 'goalpost', 'robot'], + 'bb_classes_with_robot_colors': ['ball', 'goalpost', 'robot_blue', 'robot_red', 'robot_unknown'], 'segmentation_classes': ['background', 'lines', 'field'], - 'skip_classes': ['obstacle', 'L-Intersection', 'X-Intersection', 'T-Intersection'] + 'skip_classes': ['obstacle', 'L-Intersection', 'X-Intersection', 'T-Intersection'], } @@ -37,6 +38,7 @@ def range_limited_float_type_0_to_1(arg): parser.add_argument("--skip-blurred", action="store_true", help="Skip blurred labels") parser.add_argument("--skip-concealed", action="store_true", help="Skip concealed labels") parser.add_argument("--skip-classes", nargs="+", default=[], help="These bounding box classes will be skipped") +parser.add_argument("--robots-with-team-colors", action="store_true", help="The robot class will be subdivided into subclasses, one for each team color (currently either 'blue', 'red' or 'unknown').") args = parser.parse_args() # Remove skipped classes from CLASSES list @@ -145,7 +147,19 @@ def range_limited_float_type_0_to_1(arg): relative_center_x = center_x / img_width relative_center_y = center_y / img_height - classID = CLASSES['bb_classes'].index(annotation['type']) # Derive classID from index in predefined classes + # Derive classID from index in predefined classes + if not args.robots_with_team_colors: + classID = CLASSES['bb_classes'].index(annotation['type']) + else: + class_name = annotation['type'] + + # If the annotation contains a robot, the team color has to be appended to the annotation type + # to get the full class name. + if class_name == 'robot': + class_name += f"_{annotation['color']}" + + classID = CLASSES['bb_classes_with_robot_colors'].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 '{annotation['type']}' is not supported. Image: '{img_name_with_extension}'") @@ -170,7 +184,7 @@ def range_limited_float_type_0_to_1(arg): # The names file contains the class names of bb detections and segmentations names_path = os.path.join(destination_dir, "yoeo_names.yaml") names = { - 'detection': CLASSES['bb_classes'], + 'detection': CLASSES['bb_classes'] if not args.robots_with_team_colors else CLASSES['bb_classes_with_robot_colors'], 'segmentation': CLASSES["segmentation_classes"], } with open(names_path, "w") as names_file: From 0dbd5af0a536103a540bab870d9e34a1548ea2ca Mon Sep 17 00:00:00 2001 From: Philipp Donn <30521025+phinik@users.noreply.github.com> Date: Mon, 15 Jan 2024 15:35:26 +0100 Subject: [PATCH 2/3] fix skip classes with team color --- yoeo/scripts/createYOEOLabelsFromTORSO-21.py | 44 ++++++++------------ 1 file changed, 18 insertions(+), 26 deletions(-) diff --git a/yoeo/scripts/createYOEOLabelsFromTORSO-21.py b/yoeo/scripts/createYOEOLabelsFromTORSO-21.py index fd1123f..ffb0371 100755 --- a/yoeo/scripts/createYOEOLabelsFromTORSO-21.py +++ b/yoeo/scripts/createYOEOLabelsFromTORSO-21.py @@ -8,15 +8,6 @@ from tqdm import tqdm -# Available classes for YOEO -CLASSES = { - 'bb_classes': ['ball', 'goalpost', 'robot'], - 'bb_classes_with_robot_colors': ['ball', 'goalpost', 'robot_blue', 'robot_red', 'robot_unknown'], - 'segmentation_classes': ['background', 'lines', 'field'], - 'skip_classes': ['obstacle', 'L-Intersection', 'X-Intersection', 'T-Intersection'], - } - - def range_limited_float_type_0_to_1(arg): """Type function for argparse - a float within some predefined bounds Derived from 'https://stackoverflow.com/questions/55324449/how-to-specify-a-minimum-or-maximum-float-value-with-argparse/55410582#55410582'. @@ -41,6 +32,13 @@ def range_limited_float_type_0_to_1(arg): parser.add_argument("--robots-with-team-colors", action="store_true", help="The robot class will be subdivided into subclasses, one for each team color (currently either 'blue', 'red' or 'unknown').") args = parser.parse_args() +# 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'], + 'segmentation_classes': ['background', 'lines', 'field'], + 'skip_classes': ['obstacle', 'L-Intersection', 'X-Intersection', 'T-Intersection'], + } + # Remove skipped classes from CLASSES list for skip_class in args.skip_classes: if skip_class in CLASSES['bb_classes']: @@ -124,13 +122,18 @@ def range_limited_float_type_0_to_1(arg): annotations = [] for annotation in image_data['annotations']: + class_name = annotation['type'] + + if args.robots_with_team_colors and class_name == 'robot': + class_name += f"_{annotation['color']}" + # Skip annotations, if is not a bounding box or should be skipped or is blurred or concealed and user chooses to skip them - if (annotation['type'] in CLASSES['segmentation_classes'] or # Handled by segmentations - annotation['type'] in CLASSES['skip_classes'] or # Skip this annotation class + if (class_name in CLASSES['segmentation_classes'] or # Handled by segmentations + class_name in CLASSES['skip_classes'] or # Skip this annotation class (args.skip_blurred and annotation.get('blurred', False)) or (args.skip_concealed and annotation.get('concealed', False))): continue - elif annotation['type'] in CLASSES['bb_classes']: # Handle bounding boxes + 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'])) @@ -148,21 +151,10 @@ def range_limited_float_type_0_to_1(arg): relative_center_y = center_y / img_height # Derive classID from index in predefined classes - if not args.robots_with_team_colors: - classID = CLASSES['bb_classes'].index(annotation['type']) - else: - class_name = annotation['type'] - - # If the annotation contains a robot, the team color has to be appended to the annotation type - # to get the full class name. - if class_name == 'robot': - class_name += f"_{annotation['color']}" - - classID = CLASSES['bb_classes_with_robot_colors'].index(class_name) - + 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 '{annotation['type']}' is not supported. Image: '{img_name_with_extension}'") + print(f"The annotation type '{class_name}' is not supported. Image: '{img_name_with_extension}'") # Store bounding box annotations in .txt file with open(os.path.join(labels_dir, img_name_without_extension + ".txt"), "w") as output: @@ -184,7 +176,7 @@ def range_limited_float_type_0_to_1(arg): # The names file contains the class names of bb detections and segmentations names_path = os.path.join(destination_dir, "yoeo_names.yaml") names = { - 'detection': CLASSES['bb_classes'] if not args.robots_with_team_colors else CLASSES['bb_classes_with_robot_colors'], + 'detection': CLASSES['bb_classes'], 'segmentation': CLASSES["segmentation_classes"], } with open(names_path, "w") as names_file: 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 3/3] 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}'")