Skip to content
This repository has been archived by the owner on Aug 21, 2024. It is now read-only.

feat: Add YOLO_OBB #281

Closed
wants to merge 1 commit into from
Closed

feat: Add YOLO_OBB #281

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 62 additions & 14 deletions label_studio_converter/converter.py
View file
Open in desktop
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 '
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

replace the description to this:

'description': 'Popular TXT format is created for each image file. Each txt file contains annotations for '
'the corresponding image file. The YOLO OBB format designates bounding boxes by their four corner points '
'with coordinates normalized between 0 and 1, so it is possible to export rotated objects.',

'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,52 @@ 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:
base_y = label["original_height"]
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would replace it to call get_rotated_corner_points from my code, it's a bit more clear and computationally optimized.

base_x = label["original_width"]

x1=label["x"]*base_x
y1=label["y"]*base_y
w=label["width"]*base_x
h=label["height"]*base_y
beta = math.pi * (
label["rotation"] / 180
) if "rotation" in label else 0.0

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

X = [
(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,
X= [
(P[0]/base_x/100, P[1]/base_y/100,) for P in X
]
)

annotations.append(
[category_id] + list(sum(X, ()))
)

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
View file
Open in desktop
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
View file
Open in desktop
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
Loading