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

feat: Keypoints in COCO #197

Closed
wants to merge 2 commits into from
Closed

feat: Keypoints in COCO #197

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8bd98b0
feat: Keypoints in COCO
makseq Feb 8, 2023
0fb4954
Merge branch 'master' into feat/keypoints-coco
makseq Mar 21, 2023
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
69 changes: 46 additions & 23 deletions label_studio_converter/converter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,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',
Expand Down Expand Up @@ -289,8 +289,10 @@ def _get_supported_formats(self):
input_tag_types.add(input_tag['type'])

all_formats = [f.name for f in Format]

if not ('Text' in input_tag_types and 'Labels' in output_tag_types):
all_formats.remove(Format.CONLL2003.name)

if not (
'Image' in input_tag_types
and (
Expand All @@ -300,30 +302,26 @@ def _get_supported_formats(self):
)
):
all_formats.remove(Format.VOC.name)
if not (
'Image' in input_tag_types
and (
'RectangleLabels' in output_tag_types
or '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)

if not ('Image' in input_tag_types and ('RectangleLabels' in output_tag_types or
'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.YOLO.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
)
):

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)
all_formats.remove(Format.BRUSH_TO_PNG.name)

if not (
('Audio' in input_tag_types or 'AudioPlus' in input_tag_types)
and 'TextArea' in output_tag_types
Expand Down Expand Up @@ -652,7 +650,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
Expand Down Expand Up @@ -715,6 +713,31 @@ def add_image(images, width, height, image_id, image_path):
'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")

Expand Down