Skip to content

Commit

Permalink
feat: Add YOLO_OBB format model
Browse files Browse the repository at this point in the history
  • Loading branch information
ftapajos committed Mar 21, 2024
1 parent 5171df4 commit f71572d
Show file tree
Hide file tree
Showing 3 changed files with 110 additions and 15 deletions.
70 changes: 55 additions & 15 deletions label_studio_converter/converter.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@ class Format(Enum):
ASR_MANIFEST = 10
YOLO = 11
CSV_OLD = 12
YOLO_OBB = 13

def __str__(self):
return self.name
Expand Down Expand Up @@ -121,6 +122,14 @@ class Converter(object):
'link': 'https://labelstud.io/guide/export.html#YOLO',
'tags': ['image segmentation', 'object detection'],
},
Format.YOLO_OBB: {
'title': 'YOLOv8-OBB',
'description': 'Popular TXT format is created for each image file. Each txt file contains annotations for '
'the corresponding image fileThe YOLO OBB format designates bounding boxes by their four corner points '
'with coordinates normalized between 0 and 1',
'link': 'https://labelstud.io/guide/export.html#YOLO',
'tags': ['image segmentation', 'object detection'],
},
Format.BRUSH_TO_NUMPY: {
'title': 'Brush labels to NumPy',
'description': 'Export your brush labels as NumPy 2d arrays. Each label outputs as one image.',
Expand Down Expand Up @@ -215,7 +224,7 @@ def convert(self, input_data, output_data, format, is_dir=True, **kwargs):
self.convert_to_coco(
input_data, output_data, output_image_dir=image_dir, is_dir=is_dir
)
elif format == Format.YOLO:
elif format == Format.YOLO or format == Format.YOLO_OBB:
image_dir = kwargs.get('image_dir')
label_dir = kwargs.get('label_dir')
self.convert_to_yolo(
Expand All @@ -224,6 +233,7 @@ def convert(self, input_data, output_data, format, is_dir=True, **kwargs):
output_image_dir=image_dir,
output_label_dir=label_dir,
is_dir=is_dir,
is_obb=(format == Format.YOLO_OBB)
)
elif format == Format.VOC:
image_dir = kwargs.get('image_dir')
Expand Down Expand Up @@ -727,6 +737,7 @@ def convert_to_yolo(
output_label_dir=None,
is_dir=True,
split_labelers=False,
is_obb=False,
):
"""Convert data in a specific format to the YOLO format.
Expand All @@ -744,8 +755,13 @@ def convert_to_yolo(
A boolean indicating whether `input_data` is a directory (True) or a JSON file (False).
split_labelers : bool, optional
A boolean indicating whether to create a dedicated subfolder for each labeler in the output label directory.
is_obb : bool, optional
A boolean indicating whether the format is obb
"""
self._check_format(Format.YOLO)
if is_obb:
self._check_format(Format.YOLO_OBB)
else:
self._check_format(Format.YOLO)
ensure_dir(output_dir)
notes_file = os.path.join(output_dir, 'notes.json')
class_file = os.path.join(output_dir, 'classes.txt')
Expand Down Expand Up @@ -851,20 +867,44 @@ def convert_to_yolo(
or 'rectangle' in label
or 'labels' in label
):
xywh = self.rotated_rectangle(label)
if xywh is None:
continue
if is_obb:
x1=label["x"]/100
y1=label["y"]/100
w=label["width"]/100
h=label["height"]/100
beta = math.pi * (
label["rotation"] / 180
) if "rotation" in label else 0.0

# Compute the vectors between points
v12 = (h*math.sin(beta), h*math.cos(beta))
v23 = (w * math.cos(beta), - w*math.sin(beta))

annotations.append(
[
category_id,
x1, y1,
x1 + v12[0], y1 + v12[1],
x1 + v12[0] + v23[0], y1 + v12[1] + v23[1],
x1 + v23[0], y1 + v23[1]
]
)

x, y, w, h = xywh
annotations.append(
[
category_id,
(x + w / 2) / 100,
(y + h / 2) / 100,
w / 100,
h / 100,
]
)
else:
xywh = self.rotated_rectangle(label)
if xywh is None:
continue

x, y, w, h = xywh
annotations.append(
[
category_id,
(x + w / 2) / 100,
(y + h / 2) / 100,
w / 100,
h / 100,
]
)
elif "polygonlabels" in label or 'polygon' in label:
points_abs = [(x / 100, y / 100) for x, y in label["points"]]
annotations.append(
Expand Down
3 changes: 3 additions & 0 deletions label_studio_converter/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -151,6 +151,9 @@ def export(args):
)
elif args.format == Format.YOLO:
c.convert_to_yolo(args.input, args.output, is_dir=not args.heartex_format)
elif args.format == Format.YOLO_OBB:
c.convert_to_yolo(args.input, args.output,
is_dir=not args.heartex_format, is_obb=True)
else:
raise FormatNotSupportedError()

Expand Down
52 changes: 52 additions & 0 deletions tests/test_export_yolo.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,58 @@ def create_temp_folder():
shutil.rmtree(temp_dir)


def test_convert_to_yolo_obb(create_temp_folder):
"""Check converstion label_studio json exported file to yolo with multiple labelers"""

# Generates a temporary folder and return the absolute path
# The temporary folder contains all the data generate by the following function
# For debugging replace create_temp_folder with "./tmp"
tmp_folder = create_temp_folder

output_dir = tmp_folder
output_image_dir = os.path.join(output_dir, "tmp_image")
output_label_dir = os.path.join(output_dir, "tmp_label")
project_dir = "."

converter = Converter(LABEL_CONFIG_PATH, project_dir)
converter.convert_to_yolo(
INPUT_JSON_PATH,
output_dir,
output_image_dir=output_image_dir,
output_label_dir=output_label_dir,
is_dir=False,
split_labelers=True,
is_obb=True,
)

abs_path_label_dir = os.path.abspath(output_label_dir)
expected_paths = [
os.path.join(abs_path_label_dir, "1", "image1.txt"),
os.path.join(abs_path_label_dir, "1", "image2.txt"),
os.path.join(abs_path_label_dir, "2", "image1.txt"),
]
generated_paths = get_os_walk(abs_path_label_dir)
# Check all files and subfolders have been generated.
assert check_equal_list_of_strings(
expected_paths, generated_paths
), f"Generated file: \n {generated_paths} \n does not match expected ones: \n {expected_paths}"
# Check all the annotations have been converted to yolo
for file in expected_paths:
with open(file) as f:
lines = f.readlines()
for line in lines:
split_line = line.split(" ")
assert len(split_line) == 9, "OBB lines should have 9 parameters"
assert int(split_line[0]) == float(
split_line[0]
), f"Label should be an integer. Not {split_line[0]}"
for number in split_line:
float(number)
assert (
len(lines) == 2
), f"Expect different number of annotations in file {file}."


def test_convert_to_yolo(create_temp_folder):
"""Check converstion label_studio json exported file to yolo with multiple labelers"""

Expand Down

0 comments on commit f71572d

Please sign in to comment.