From 025d9740ab24e7ca1ccc62eb6ab2f6d43c6fb8d3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fl=C3=A1vio=20Tapaj=C3=B3s?= Date: Tue, 19 Mar 2024 14:53:29 -0300 Subject: [PATCH] WIP: Add YOLO_OBB Author could not test this commit. Any help is appreciated ref: https://docs.ultralytics.com/datasets/obb/#supported-obb-dataset-formats ref: #280 --- label_studio_converter/converter.py | 69 ++++++++++++++++++++++------- label_studio_converter/main.py | 3 ++ 2 files changed, 57 insertions(+), 15 deletions(-) diff --git a/label_studio_converter/converter.py b/label_studio_converter/converter.py index f652a28b..858f6e48 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,43 @@ 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 = (w * math.cos(beta), - w*math.sen(beta)) + v14 = (-h*math.sen(beta), - h*math.cos(beta)) + annotations.append( + [ + category_id, + x1, y1, + x1 + v12[0], y1 + v12[1], + x1 + v12[0] + v14[0], y1 + v12[1] + v14[1], + x1 + v14[0], y1 + v14[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()