Add skeletons in datumaro format (#47)

- Added support for import and export of skeletons in Datumaro format
- Added tests

CHANGELOG.md

@@ -43,6 +43,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   (<https://github.com/cvat-ai/datumaro/pull/39>)
 - An option to specify scale factor in `resize` transform
   (<https://github.com/cvat-ai/datumaro/pull/46>)
+- Skeleton support in datumaro format
+  (<https://github.com/cvat-ai/datumaro/pull/47>)
 
 ### Changed
 - `env.detect_dataset()` now returns a list of detected formats at all recursion levels

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:
         """

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)
 

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

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
+        ]
+      }
+    }
+  ]
+}