diff --git a/CHANGELOG.md b/CHANGELOG.md index 25bb23324a..3ba0b70a60 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -43,6 +43,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 () - An option to specify scale factor in `resize` transform () +- Skeleton support in datumaro format + () ### Changed - `env.detect_dataset()` now returns a list of detected formats at all recursion levels diff --git a/datumaro/components/annotation.py b/datumaro/components/annotation.py index a64759ba3b..ee8bd73c1c 100644 --- a/datumaro/components/annotation.py +++ b/datumaro/components/annotation.py @@ -708,9 +708,11 @@ class Category: @classmethod def from_iterable( cls, - iterable: Union[ - Tuple[int, List[str]], - Tuple[int, List[str], Set[Tuple[int, int]]], + iterable: Iterable[ + Union[ + Tuple[int, List[str]], + Tuple[int, List[str], Set[Tuple[int, int]]], + ], ], ) -> PointsCategories: """ diff --git a/datumaro/plugins/datumaro_format/converter.py b/datumaro/plugins/datumaro_format/converter.py index 9ce8eece99..a058d38ae0 100644 --- a/datumaro/plugins/datumaro_format/converter.py +++ b/datumaro/plugins/datumaro_format/converter.py @@ -16,6 +16,7 @@ from datumaro.components.annotation import ( Annotation, + AnnotationType, Bbox, Caption, Cuboid3d, @@ -28,11 +29,12 @@ Polygon, PolyLine, RleMask, + Skeleton, _Shape, ) from datumaro.components.converter import Converter from datumaro.components.dataset import ItemStatus -from datumaro.components.extractor import DEFAULT_SUBSET_NAME, DatasetItem +from datumaro.components.extractor import DEFAULT_SUBSET_NAME, CategoriesInfo, DatasetItem from datumaro.components.media import Image, MediaElement, PointCloud from datumaro.util import cast, dump_json_file @@ -143,11 +145,13 @@ def add_item(self, item: DatasetItem): converted_ann = self._convert_caption_object(ann) elif isinstance(ann, Cuboid3d): converted_ann = self._convert_cuboid_3d_object(ann) + elif isinstance(ann, Skeleton): + converted_ann = self._convert_skeleton_object(ann) else: raise NotImplementedError() annotations.append(converted_ann) - def add_categories(self, categories): + def add_categories(self, categories: CategoriesInfo): for ann_type, desc in categories.items(): if isinstance(desc, LabelCategories): converted_desc = self._convert_label_categories(desc) @@ -162,7 +166,7 @@ def add_categories(self, categories): def write(self, ann_file): dump_json_file(ann_file, self._data) - def _convert_annotation(self, obj): + def _convert_annotation(self, obj: Annotation) -> dict: assert isinstance(obj, Annotation) ann_json = { @@ -266,6 +270,37 @@ def _convert_cuboid_3d_object(self, obj): ) return converted + def _convert_skeleton_object(self, obj: Skeleton) -> dict: + label_ordering = [ + item["labels"] + for item in self.categories[AnnotationType.points.name]["items"] + if item["label_id"] == obj.label + ][0] + + def get_label_position(label_id): + return label_ordering.index( + self.categories[AnnotationType.label.name]["labels"][label_id]["name"] + ) + + points = [0.0, 0.0, Points.Visibility.absent.value] * len(label_ordering) + points_attributes = [{}] * len(label_ordering) + for element in obj.elements: + assert len(element.points) == 2 and len(element.visibility) == 1 + position = get_label_position(element.label) + points[position * 3 : position * 3 + 3] = [ + element.points[0], + element.points[1], + element.visibility[0].value, + ] + points_attributes[position] = element.attributes + + return dict( + self._convert_annotation(obj), + label_id=cast(obj.label, int), + points=points, + points_attributes=points_attributes, + ) + def _convert_attribute_categories(self, attributes): return sorted(attributes) diff --git a/datumaro/plugins/datumaro_format/extractor.py b/datumaro/plugins/datumaro_format/extractor.py index 62f781169d..1afd07171c 100644 --- a/datumaro/plugins/datumaro_format/extractor.py +++ b/datumaro/plugins/datumaro_format/extractor.py @@ -3,6 +3,7 @@ # SPDX-License-Identifier: MIT import os.path as osp +from typing import List, Union from datumaro.components.annotation import ( AnnotationType, @@ -17,12 +18,13 @@ Polygon, PolyLine, RleMask, + Skeleton, ) -from datumaro.components.errors import DatasetImportError +from datumaro.components.errors import DatasetImportError, InvalidAnnotationError from datumaro.components.extractor import DatasetItem, Importer, SourceExtractor from datumaro.components.format_detection import FormatDetectionContext from datumaro.components.media import Image, MediaElement, PointCloud -from datumaro.util import parse_json, parse_json_file +from datumaro.util import parse_json, parse_json_file, take_by from .format import DatumaroPath @@ -150,8 +152,7 @@ def _load_items(self, parsed): return items - @staticmethod - def _load_annotations(item): + def _load_annotations(self, item: dict): parsed = item["annotations"] loaded = [] @@ -251,11 +252,54 @@ def _load_annotations(item): ) ) + elif ann_type == AnnotationType.skeleton: + loaded.append( + Skeleton( + elements=self._load_skeleton_elements_annotations(ann, label_id, points), + label=label_id, + id=ann_id, + attributes=attributes, + group=group, + z_order=z_order, + ) + ) + else: raise NotImplementedError() return loaded + def _load_skeleton_elements_annotations( + self, ann: dict, label_id: int, points: List[Union[float, int]] + ) -> List[Points]: + if len(points) % 3 != 0: + raise InvalidAnnotationError( + f"Points have invalid value count {len(points)}, " + "which is not divisible by 3. Expected (x, y, visibility) triplets." + ) + points_attributes = ann.get("points_attributes") + if len(points) != len(points_attributes) * 3: + raise InvalidAnnotationError( + f"Points and Points_attributes lengths ({len(points)}, {len(points_attributes)}) do not match, " + "for each triplet (x, y, visibility) in points there should be one dict in points_attributes." + ) + + label_category = self._categories[AnnotationType.label] + sub_labels = self._categories[AnnotationType.points].items[label_id].labels + return [ + Points( + points=[x, y], + visibility=[v], + label=label_category.find( + name=sub_label, parent=label_category.items[label_id].name + )[0], + attributes=attrs, + ) + for (x, y, v), sub_label, attrs in zip( + take_by(points, 3), sub_labels, points_attributes + ) + ] + class DatumaroImporter(Importer): @classmethod diff --git a/tests/assets/datumaro_dataset/with_skeleton/annotations/default.json b/tests/assets/datumaro_dataset/with_skeleton/annotations/default.json new file mode 100644 index 0000000000..ae9e8780e8 --- /dev/null +++ b/tests/assets/datumaro_dataset/with_skeleton/annotations/default.json @@ -0,0 +1,100 @@ +{ + "info": {}, + "categories": { + "label": { + "labels": [ + { + "name": "skeleton-label", + "parent": "", + "attributes": [] + }, + { + "name": "point1-label", + "parent": "skeleton-label", + "attributes": [] + }, + { + "name": "point3-label", + "parent": "skeleton-label", + "attributes": [] + }, + { + "name": "point2-label", + "parent": "skeleton-label", + "attributes": [] + } + ], + "attributes": [ + "occluded", + "point2-attribute-text", + "point3-attribute-checkbox" + ] + }, + "points": { + "items": [ + { + "label_id": 0, + "labels": [ + "point1-label", + "point3-label", + "point2-label" + ], + "joints": [ + [ + 2, + 3 + ], + [ + 1, + 2 + ] + ] + } + ] + } + }, + "items": [ + { + "id": "100", + "annotations": [ + { + "id": 0, + "type": "skeleton", + "attributes": { + "occluded": false, + "keyframe": false + }, + "group": 0, + "label_id": 0, + "points": [ + 0.9, + 3.53, + 2, + 2.45, + 7.6, + 2, + 5.2, + 2.5, + 2 + ], + "points_attributes": [ + {}, + { + "point3-attribute-checkbox": true + }, + { + "point2-attribute-text": "some text" + } + ] + } + ], + "image": { + "path": "100.jpg", + "size": [ + 10, + 6 + ] + } + } + ] +} \ No newline at end of file diff --git a/tests/assets/datumaro_dataset/with_skeleton/images/default/100.jpg b/tests/assets/datumaro_dataset/with_skeleton/images/default/100.jpg new file mode 100644 index 0000000000..4409e5e19b Binary files /dev/null and b/tests/assets/datumaro_dataset/with_skeleton/images/default/100.jpg differ diff --git a/tests/test_datumaro_format.py b/tests/test_datumaro_format.py index 6089b2242e..667b14d1da 100644 --- a/tests/test_datumaro_format.py +++ b/tests/test_datumaro_format.py @@ -19,6 +19,7 @@ PointsCategories, Polygon, PolyLine, + Skeleton, ) from datumaro.components.environment import Environment from datumaro.components.extractor import DatasetItem @@ -250,6 +251,67 @@ def test_can_import_pcd_dataset(self): compare_datasets_strict(self, expected, Dataset.load(dataset_path)) + @mark_requirement(Requirements.DATUM_GENERAL_REQ) + def test_can_import_skeleton_dataset(self): + dataset_path = str(Path(__file__).parent / "assets" / "datumaro_dataset" / "with_skeleton") + + label_categories = LabelCategories.from_iterable( + [ + ("skeleton-label",), + ("point1-label", "skeleton-label"), + ("point3-label", "skeleton-label"), + ("point2-label", "skeleton-label"), + ] + ) + label_categories.attributes = { + "point3-attribute-checkbox", + "occluded", + "point2-attribute-text", + } + points_categories = PointsCategories.from_iterable( + [ + (0, ["point1-label", "point3-label", "point2-label"], {(2, 3), (1, 2)}), + ] + ) + + expected = Dataset.from_iterable( + [ + DatasetItem( + id="100", + subset="default", + media=Image(data=np.ones((10, 6, 3))), + annotations=[ + Skeleton( + [ + Points( + [0.9, 3.53], + label=1, + attributes={}, + ), + Points( + [2.45, 7.6], + label=2, + attributes={"point3-attribute-checkbox": True}, + ), + Points( + [5.2, 2.5], + label=3, + attributes={"point2-attribute-text": "some text"}, + ), + ], + label=0, + attributes={"occluded": False, "keyframe": False}, + ), + ], + ), + ], + categories={ + AnnotationType.label: label_categories, + AnnotationType.points: points_categories, + }, + ) + compare_datasets_strict(self, expected, Dataset.load(dataset_path)) + @mark_requirement(Requirements.DATUM_GENERAL_REQ) def test_can_detect(self): with TestDir() as test_dir: @@ -519,3 +581,79 @@ def test_can_save_and_load_with_pointcloud(self): compare=None, dimension=Dimensions.dim_3d, ) + + @mark_requirement(Requirements.DATUM_GENERAL_REQ) + def test_can_save_and_load_with_skeleton(self): + categories = { + AnnotationType.label: LabelCategories.from_iterable( + [ + ("point2", "skeleton"), + ("point1", "skeleton"), + ("skeleton",), + ("point3", "skeleton"), + ("point4", "skeleton"), + ] + ), + AnnotationType.points: PointsCategories.from_iterable( + [ + (2, ["point2", "point1", "point3", "point4"], {(2, 1), (0, 3)}), + ] + ), + } + source_elements = [ + Points([1, 1], label=0, attributes={"occluded": False, "outside": True}), + Points( + [2, 2], + label=1, + attributes={"occluded": False, "outside": False, "custom": "value"}, + ), + Points([4, 4], label=4, attributes={"occluded": False, "outside": False}), + ] + source_dataset = Dataset.from_iterable( + [ + DatasetItem( + id="img1", + subset="train", + media=Image(data=np.ones((10, 10, 3))), + annotations=[ + Skeleton( + source_elements, + label=2, + attributes={"occluded": False}, + ) + ], + ), + ], + categories=categories, + ) + target_elements = [ + source_elements[0], + source_elements[1], + Points([0, 0], label=3, visibility=[Points.Visibility.absent.value]), + source_elements[2], + ] + target_dataset = Dataset.from_iterable( + [ + DatasetItem( + id="img1", + subset="train", + media=Image(data=np.ones((10, 10, 3))), + annotations=[ + Skeleton( + target_elements, + label=2, + attributes={"occluded": False}, + ) + ], + ), + ], + categories=categories, + ) + + with TestDir() as test_dir: + self._test_save_and_load( + source_dataset, + partial(DatumaroConverter.convert, save_media=True), + test_dir, + target_dataset=target_dataset, + )