diff --git a/label_studio_converter/converter.py b/label_studio_converter/converter.py index f652a28b..af371dbd 100644 --- a/label_studio_converter/converter.py +++ b/label_studio_converter/converter.py @@ -57,6 +57,7 @@ class Format(Enum): ASR_MANIFEST = 10 YOLO = 11 CSV_OLD = 12 + YOLO_OBB = 13 def __str__(self): return self.name @@ -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.', @@ -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( @@ -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') @@ -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. @@ -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') @@ -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( diff --git a/label_studio_converter/main.py b/label_studio_converter/main.py index 2eba0678..b126aed6 100644 --- a/label_studio_converter/main.py +++ b/label_studio_converter/main.py @@ -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() diff --git a/tests/test_export_yolo.py b/tests/test_export_yolo.py index 13904bb3..41da4389 100644 --- a/tests/test_export_yolo.py +++ b/tests/test_export_yolo.py @@ -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"""