Skip to content

feat: LEAP-1840: Add KeyPoints to COCO/YOLO export #451

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 7 commits into from
May 9, 2025
Merged

feat: LEAP-1840: Add KeyPoints to COCO/YOLO export #451

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
41b7f7c
feat: LEAP-1840: Add KeyPoints to COCO export
triklozoid May 2, 2025
046eda0
Add keypoints for YOLO
triklozoid May 6, 2025
91f6ed6
Fix tests
triklozoid May 7, 2025
6e48f7f
SDK regeneration
fern-api[bot] May 7, 2025
5db25c0
Fix kp_order
triklozoid May 7, 2025
3247bbb
Fix format descriptions
triklozoid May 8, 2025
2839b9a
SDK regeneration
fern-api[bot] May 9, 2025
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
433 changes: 231 additions & 202 deletions poetry.lock
View file
Open in desktop

Large diffs are not rendered by default.

3 changes: 2 additions & 1 deletion pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ name = "label-studio-sdk"

[tool.poetry]
name = "label-studio-sdk"
version = "1.0.13.dev"
version = "1.0.13"
description = ""
readme = "README.md"
authors = []
Expand Down Expand Up @@ -45,6 +45,7 @@ jsonschema = ">=4.23.0"
lxml = ">=4.2.5"
nltk = "^3.9.1"
numpy = ">=1.26.4,<3.0.0"
opencv-python = "^4.9.0"
pandas = ">=0.24.0"
pydantic = ">= 1.9.2"
pydantic-core = "^2.18.2"
Expand Down
114 changes: 34 additions & 80 deletions src/label_studio_sdk/converter/converter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,13 +11,14 @@
from enum import Enum
from glob import glob
from shutil import copy2
from typing import Optional
from typing import Optional, List, Tuple

import ijson
import ujson as json
from PIL import Image
from label_studio_sdk.converter import brush
from label_studio_sdk.converter.audio import convert_to_asr_json_manifest
from label_studio_sdk.converter.keypoints import process_keypoints_for_coco, build_kp_order, update_categories_for_keypoints, keypoints_in_label_config, get_yolo_categories_for_keypoints
from label_studio_sdk.converter.exports import csv2
from label_studio_sdk.converter.utils import (
parse_config,
Expand All @@ -34,6 +35,7 @@
convert_annotation_to_yolo_obb,
)
from label_studio_sdk._extensions.label_studio_tools.core.utils.io import get_local_path
from label_studio_sdk.converter.exports.yolo import process_and_save_yolo_annotations

logger = logging.getLogger(__name__)

Expand Down Expand Up @@ -109,7 +111,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.COCO_WITH_IMAGES: {
"title": "COCO with Images",
Expand Down Expand Up @@ -205,6 +207,7 @@ def __init__(
self._schema = None
self.access_token = access_token
self.hostname = hostname
self.is_keypoints = None

if isinstance(config, dict):
self._schema = config
Expand Down Expand Up @@ -376,11 +379,14 @@ def _get_supported_formats(self):
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)
all_formats.remove(Format.COCO_WITH_IMAGES.name)
Expand Down Expand Up @@ -522,6 +528,9 @@ def annotation_result_from_task(self, task):
if "original_height" in r:
v["original_height"] = r["original_height"]
outputs[r["from_name"]].append(v)
if self.is_keypoints:
v['id'] = r.get('id')
v['parentID'] = r.get('parentID')

data = Converter.get_data(task, outputs, annotation)
if "agreement" in task:
Expand Down Expand Up @@ -638,6 +647,7 @@ def add_image(images, width, height, image_id, image_path):
os.makedirs(output_image_dir, exist_ok=True)
images, categories, annotations = [], [], []
categories, category_name_to_id = self._get_labels()
categories, category_name_to_id = update_categories_for_keypoints(categories, category_name_to_id, self._schema)
data_key = self._data_keys[0]
item_iterator = (
self.iter_from_dir(input_data)
Expand Down Expand Up @@ -703,9 +713,10 @@ def add_image(images, width, height, image_id, image_path):
logger.debug(f'Empty bboxes for {item["output"]}')
continue

keypoint_labels = []
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 @@ -775,11 +786,22 @@ def add_image(images, width, height, image_id, image_path):
"area": get_polygon_area(x, y),
}
)
elif "keypointlabels" in label:
keypoint_labels.append(label)
else:
raise ValueError("Unknown label type")

if os.getenv("LABEL_STUDIO_FORCE_ANNOTATOR_EXPORT"):
annotations[-1].update({"annotator": get_annotator(item)})
if keypoint_labels:
kp_order = build_kp_order(self._schema)
annotations.append(process_keypoints_for_coco(
keypoint_labels,
kp_order,
annotation_id=len(annotations),
image_id=image_id,
category_name_to_id=category_name_to_id,
))

with io.open(output_file, mode="w", encoding="utf8") as fout:
json.dump(
Expand Down Expand Up @@ -846,7 +868,14 @@ def convert_to_yolo(
else:
output_label_dir = os.path.join(output_dir, "labels")
os.makedirs(output_label_dir, exist_ok=True)
categories, category_name_to_id = self._get_labels()
is_keypoints = keypoints_in_label_config(self._schema)

if is_keypoints:
# we use this attribute to add id and parentID to annotation data
self.is_keypoints = True
categories, category_name_to_id = get_yolo_categories_for_keypoints(self._schema)
else:
categories, category_name_to_id = self._get_labels()
data_key = self._data_keys[0]
item_iterator = (
self.iter_from_dir(input_data)
Expand Down Expand Up @@ -923,82 +952,7 @@ def convert_to_yolo(
pass
continue

annotations = []
for label in labels:
category_name = None
category_names = [] # considering multi-label
for key in ["rectanglelabels", "polygonlabels", "labels"]:
if key in label and len(label[key]) > 0:
# change to save multi-label
for category_name in label[key]:
category_names.append(category_name)

if len(category_names) == 0:
logger.debug(
"Unknown label type or labels are empty: " + str(label)
)
continue

for category_name in category_names:
if category_name not in category_name_to_id:
category_id = len(categories)
category_name_to_id[category_name] = category_id
categories.append({"id": category_id, "name": category_name})
category_id = category_name_to_id[category_name]

if (
"rectanglelabels" in label
or "rectangle" in label
or "labels" in label
):
# yolo obb
if is_obb:
obb_annotation = convert_annotation_to_yolo_obb(label)
if obb_annotation is None:
continue

top_left, top_right, bottom_right, bottom_left = (
obb_annotation
)
x1, y1 = top_left
x2, y2 = top_right
x3, y3 = bottom_right
x4, y4 = bottom_left
annotations.append(
[category_id, x1, y1, x2, y2, x3, y3, x4, y4]
)

# simple yolo
else:
annotation = convert_annotation_to_yolo(label)
if annotation is None:
continue

(
x,
y,
w,
h,
) = annotation
annotations.append([category_id, x, y, w, h])

elif "polygonlabels" in label or "polygon" in label:
if not ('points' in label):
continue
points_abs = [(x / 100, y / 100) for x, y in label["points"]]
annotations.append(
[category_id]
+ [coord for point in points_abs for coord in point]
)
else:
raise ValueError(f"Unknown label type {label}")
with open(label_path, "w") as f:
for annotation in annotations:
for idx, l in enumerate(annotation):
if idx == len(annotation) - 1:
f.write(f"{l}\n")
else:
f.write(f"{l} ")
categories, category_name_to_id = process_and_save_yolo_annotations(labels, label_path, category_name_to_id, categories, is_obb, is_keypoints, self._schema)
with open(class_file, "w", encoding="utf8") as f:
for c in categories:
f.write(c["name"] + "\n")
Expand Down
149 changes: 149 additions & 0 deletions src/label_studio_sdk/converter/exports/yolo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,149 @@
import logging
from label_studio_sdk.converter.utils import convert_annotation_to_yolo, convert_annotation_to_yolo_obb
from label_studio_sdk.converter.keypoints import build_kp_order

logger = logging.getLogger(__name__)

def process_keypoints_for_yolo(labels, label_path,
category_name_to_id, categories,
is_obb, kp_order):
class_map = {c['name']: c['id'] for c in categories}

rectangles = {}
for item in labels:
if item['type'].lower() == 'rectanglelabels':
bbox_id = item['id']
cls_name = item['rectanglelabels'][0]
cls_idx = class_map.get(cls_name)
if cls_idx is None:
continue

x = item['x'] / 100.0
y = item['y'] / 100.0
width = item['width'] / 100.0
height = item['height'] / 100.0
x_c = x + width / 2.0
y_c = y + height / 2.0

rectangles[bbox_id] = {
'class_idx': cls_idx,
'x_center': x_c,
'y_center': y_c,
'width': width,
'height': height,
'kp_dict': {}
}

for item in labels:
if item['type'].lower() == 'keypointlabels':
parent_id = item.get('parentID')
if parent_id not in rectangles:
continue
label_name = item['keypointlabels'][0]
kp_x = item['x'] / 100.0
kp_y = item['y'] / 100.0
rectangles[parent_id]['kp_dict'][label_name] = (kp_x, kp_y, 2) # 2 = visible

lines = []
for rect in rectangles.values():
base = [
rect['class_idx'],
rect['x_center'],
rect['y_center'],
rect['width'],
rect['height']
]
keypoints = []
for k in kp_order:
keypoints.extend(rect['kp_dict'].get(k, (0.0, 0.0, 0)))
line = ' '.join(map(str, base + keypoints))
lines.append(line)

with open(label_path, 'w', encoding='utf-8') as f:
f.write('\n'.join(lines))


def process_and_save_yolo_annotations(labels, label_path, category_name_to_id, categories, is_obb, is_keypoints, label_config):
if is_keypoints:
kp_order = build_kp_order(label_config)
process_keypoints_for_yolo(labels, label_path, category_name_to_id, categories, is_obb, kp_order)
return categories, category_name_to_id

annotations = []
for label in labels:
category_name = None
category_names = [] # considering multi-label
for key in ["rectanglelabels", "polygonlabels", "labels"]:
if key in label and len(label[key]) > 0:
# change to save multi-label
for category_name in label[key]:
category_names.append(category_name)

if len(category_names) == 0:
logger.debug(
"Unknown label type or labels are empty: " + str(label)
)
continue

for category_name in category_names:
if category_name not in category_name_to_id:
category_id = len(categories)
category_name_to_id[category_name] = category_id
categories.append({"id": category_id, "name": category_name})
category_id = category_name_to_id[category_name]

if (
"rectanglelabels" in label
or "rectangle" in label
or "labels" in label
):
# yolo obb
if is_obb:
obb_annotation = convert_annotation_to_yolo_obb(label)
if obb_annotation is None:
continue

top_left, top_right, bottom_right, bottom_left = (
obb_annotation
)
x1, y1 = top_left
x2, y2 = top_right
x3, y3 = bottom_right
x4, y4 = bottom_left
annotations.append(
[category_id, x1, y1, x2, y2, x3, y3, x4, y4]
)

# simple yolo
else:
annotation = convert_annotation_to_yolo(label)
if annotation is None:
continue

(
x,
y,
w,
h,
) = annotation
annotations.append([category_id, x, y, w, h])

elif "polygonlabels" in label or "polygon" in label:
if not ('points' in label):
continue
points_abs = [(x / 100, y / 100) for x, y in label["points"]]
annotations.append(
[category_id]
+ [coord for point in points_abs for coord in point]
)
else:
raise ValueError(f"Unknown label type {label}")
with open(label_path, "w") as f:
for annotation in annotations:
for idx, l in enumerate(annotation):
if idx == len(annotation) - 1:
f.write(f"{l}\n")
else:
f.write(f"{l} ")

return categories, category_name_to_id
Loading
Loading