Skip to content

Commit

Permalink
fix problem with annotation not in image
Browse files Browse the repository at this point in the history
  • Loading branch information
phinik committed Jan 15, 2024
1 parent 0dbd5af commit 759f4f0
Showing 1 changed file with 24 additions and 21 deletions.
45 changes: 24 additions & 21 deletions yoeo/scripts/createYOEOLabelsFromTORSO-21.py
Original file line number Diff line number Diff line change
Expand Up @@ -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'],
}
Expand Down Expand Up @@ -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']}"

Expand All @@ -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}'")

Expand Down

0 comments on commit 759f4f0

Please sign in to comment.