From 8bd98b06736d59bef640c1da771ba17c0b0bdbc5 Mon Sep 17 00:00:00 2001 From: makseq-ubnt Date: Wed, 8 Feb 2023 02:35:09 +0000 Subject: [PATCH] feat: Keypoints in COCO --- label_studio_converter/converter.py | 38 ++++++++++++++++++++++++++--- 1 file changed, 34 insertions(+), 4 deletions(-) diff --git a/label_studio_converter/converter.py b/label_studio_converter/converter.py index 6488698e..2f87809b 100644 --- a/label_studio_converter/converter.py +++ b/label_studio_converter/converter.py @@ -93,7 +93,7 @@ class Converter(object): 'description': 'Popular machine learning format used by the COCO dataset for object detection and image ' 'segmentation tasks with polygons and rectangles.', 'link': 'https://labelstud.io/guide/export.html#COCO', - 'tags': ['image segmentation', 'object detection'] + 'tags': ['image segmentation', 'object detection', 'keypoints'] }, Format.VOC: { 'title': 'Pascal VOC XML', @@ -243,9 +243,14 @@ def _get_supported_formats(self): 'PolygonLabels' in output_tag_types) or 'Rectangle' in output_tag_types and 'Labels' in output_tag_types or 'PolygonLabels' in output_tag_types and 'Labels' in output_tag_types): - - all_formats.remove(Format.COCO.name) all_formats.remove(Format.YOLO.name) + if not ('Image' in input_tag_types and ('RectangleLabels' in output_tag_types or + 'PolygonLabels' in output_tag_types or + 'KeyPointLabels' in output_tag_types) or + 'Rectangle' in output_tag_types and 'Labels' in output_tag_types or + 'PolygonLabels' in output_tag_types and 'Labels' in output_tag_types or + 'KeyPointLabels' in output_tag_types and 'Labels' in output_tag_types): + all_formats.remove(Format.COCO.name) if not ('Image' in input_tag_types and ('BrushLabels' in output_tag_types or 'brushlabels' in output_tag_types or 'Brush' in output_tag_types and 'Labels' in output_tag_types)): all_formats.remove(Format.BRUSH_TO_NUMPY.name) @@ -539,7 +544,7 @@ def add_image(images, width, height, image_id, image_path): for label in labels: category_name = None - for key in ['rectanglelabels', 'polygonlabels', 'labels']: + for key in ['rectanglelabels', 'polygonlabels', 'keypointlabels', 'labels']: if key in label and len(label[key]) > 0: category_name = label[key][0] break @@ -592,6 +597,31 @@ def add_image(images, width, height, image_id, image_path): 'iscrowd': 0, 'area': get_polygon_area(x, y) }) + elif "keypointlabels" in label: + # write only 1 point, probably it's not correct for keypoints in coco format, + # because we need to concatenate all points from keypointslabel together and + # write them as one whole array, more details: https://cocodataset.org/#format-data + visibility_flag = 2 + if 'flag:0' in label['keypointlabels']: + visibility_flag = 0 + if 'flag:1' in label['keypointlabels']: + visibility_flag = 1 + if 'flag:2' in label['keypointlabels']: + visibility_flag = 2 + x, y = label['x'] / 100 * width, label['y'] * height + points = [(x, y, visibility_flag)] + + annotations.append({ + 'id': annotation_id, + 'image_id': image_id, + 'category_id': category_id, + 'num_keypoints': len([points]), + 'keypoints': points, + 'bbox': get_polygon_bounding_box([x], [y]), + 'ignore': 0, + 'iscrowd': 0, + 'area': get_polygon_area([x], [y]) + }) else: raise ValueError("Unknown label type")