Fix project import with skeletons

CHANGELOG.md

@@ -25,6 +25,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 non-ascii paths while adding files from "Connected file share" (issue #4428)
 - Removed unnecessary volumes defined in docker-compose.serverless.yml
 (<https://github.com/openvinotoolkit/cvat/pull/4659>)
+- Project import with skeletons (<https://github.com/opencv/cvat/pull/4867>)
 
 ### Security
 - TDB

cvat/apps/dataset_manager/bindings.py

@@ -524,9 +524,7 @@ def _import_track(self, track, parent_label_id=None):
             ]
             shape['attributes'] = [self._import_attribute(label_id, attrib, mutable=True)
                 for attrib in shape['attributes']
-                if self._get_mutable_attribute_id(label_id, attrib.name) or (
-                    self.soft_attribute_import and attrib.name not in CVAT_INTERNAL_ATTRIBUTES
-                )
+                if self._get_mutable_attribute_id(label_id, attrib.name)
             ]
             shape['points'] = list(map(float, shape['points']))
 
@@ -1101,9 +1099,9 @@ def _load_user_info(meta: dict):
     def _read_cvat_anno(self, cvat_frame_anno: Union[ProjectData.Frame, TaskData.Frame], labels: list):
         categories = self.categories()
         label_cat = categories[dm.AnnotationType.label]
-        def map_label(name): return label_cat.find(name)[0]
+        def map_label(name, parent=''): return label_cat.find(name, parent)[0]
         label_attrs = {
-            label['name']: label['attributes']
+            label.get('parent', '') + label['name']: label['attributes']
             for _, label in labels
         }
 
@@ -1198,9 +1196,9 @@ def _make_image(i, **kwargs):
     def _read_cvat_anno(self, cvat_frame_anno: TaskData.Frame, labels: list):
         categories = self.categories()
         label_cat = categories[dm.AnnotationType.label]
-        def map_label(name, parent=""): return label_cat.find(name, parent)[0]
+        def map_label(name, parent=''): return label_cat.find(name, parent)[0]
         label_attrs = {
-            label.get("parent", "") + label['name']: label['attributes']
+            label.get('parent', '') + label['name']: label['attributes']
             for _, label in labels
         }
 
@@ -1562,6 +1560,7 @@ def import_dm_annotations(dm_dataset: dm.Dataset, instance_data: Union[TaskData,
                     for n, v in ann.attributes.items()
                 ]
 
+                points = []
                 if ann.type in shapes:
                     points = []
                     if ann.type == dm.AnnotationType.cuboid_3d:
@@ -1647,6 +1646,11 @@ def import_dm_annotations(dm_dataset: dm.Dataset, instance_data: Union[TaskData,
 
                         if ann.type == dm.AnnotationType.skeleton:
                             for element in ann.elements:
+                                element_keyframe = dm.util.cast(element.attributes.get('keyframe', None), bool) is True
+                                element_outside = dm.util.cast(element.attributes.pop('outside', None), bool) is True
+                                if not element_keyframe and not element_outside:
+                                    continue
+
                                 if element.label not in tracks[track_id]['elements']:
                                     tracks[track_id]['elements'][element.label] = instance_data.Track(
                                         label=label_cat.items[element.label].name,
@@ -1659,7 +1663,6 @@ def import_dm_annotations(dm_dataset: dm.Dataset, instance_data: Union[TaskData,
                                     for n, v in element.attributes.items()
                                 ]
                                 element_occluded = dm.util.cast(element.attributes.pop('occluded', None), bool) is True
-                                element_outside = dm.util.cast(element.attributes.pop('outside', None), bool) is True
                                 element_source = element.attributes.pop('source').lower() \
                                     if element.attributes.get('source', '').lower() in {'auto', 'manual'} else 'manual'
                                 tracks[track_id]['elements'][element.label].shapes.append(instance_data.TrackedShape(

cvat/apps/dataset_manager/formats/cvat.py

@@ -14,7 +14,7 @@
 
 from datumaro.components.annotation import (AnnotationType, Bbox, Label,
                                             LabelCategories, Points, Polygon,
-                                            PolyLine)
+                                            PolyLine, Skeleton)
 from datumaro.components.dataset import Dataset, DatasetItem
 from datumaro.components.extractor import (DEFAULT_SUBSET_NAME, Extractor,
                                            Importer)
@@ -119,23 +119,34 @@ def _parse(cls, path):
         items = OrderedDict()
 
         track = None
+        track_element = None
+        track_shapes = None
         shape = None
+        shape_element = None
         tag = None
         attributes = None
+        element_attributes = None
         image = None
         subset = None
         for ev, el in context:
             if ev == 'start':
                 if el.tag == 'track':
-                    frame_size = tasks_info[int(el.attrib.get('task_id'))]['frame_size'] if el.attrib.get('task_id') else tuple(tasks_info.values())[0]['frame_size']
-                    track = {
-                        'id': el.attrib['id'],
-                        'label': el.attrib.get('label'),
-                        'group': int(el.attrib.get('group_id', 0)),
-                        'height': frame_size[0],
-                        'width': frame_size[1],
-                    }
-                    subset = el.attrib.get('subset')
+                    if track:
+                        track_element = {
+                            'id': el.attrib['id'],
+                            'label': el.attrib.get('label'),
+                        }
+                    else:
+                        frame_size = tasks_info[int(el.attrib.get('task_id'))]['frame_size'] if el.attrib.get('task_id') else tuple(tasks_info.values())[0]['frame_size']
+                        track = {
+                            'id': el.attrib['id'],
+                            'label': el.attrib.get('label'),
+                            'group': int(el.attrib.get('group_id', 0)),
+                            'height': frame_size[0],
+                            'width': frame_size[1],
+                        }
+                        subset = el.attrib.get('subset')
+                        track_shapes = {}
                 elif el.tag == 'image':
                     image = {
                         'name': el.attrib.get('name'),
@@ -145,16 +156,28 @@ def _parse(cls, path):
                     }
                     subset = el.attrib.get('subset')
                 elif el.tag in cls._SUPPORTED_SHAPES and (track or image):
-                    attributes = {}
-                    shape = {
-                        'type': None,
-                        'attributes': attributes,
-                    }
-                    if track:
-                        shape.update(track)
-                        shape['track_id'] = int(track['id'])
-                    if image:
-                        shape.update(image)
+                    if shape and shape['type'] == 'skeleton':
+                        element_attributes = {}
+                        shape_element = {
+                            'type': 'rectangle' if el.tag == 'box' else el.tag,
+                            'attributes': element_attributes,
+                        }
+                        shape_element.update(image)
+                    else:
+                        attributes = {}
+                        shape = {
+                            'type': 'rectangle' if el.tag == 'box' else el.tag,
+                            'attributes': attributes,
+                        }
+                        shape['elements'] = []
+                        if track_element:
+                            shape.update(track_element)
+                            shape['track_id'] = int(track_element['id'])
+                        elif track:
+                            shape.update(track)
+                            shape['track_id'] = int(track['id'])
+                        if image:
+                            shape.update(image)
                 elif el.tag == 'tag' and image:
                     attributes = {}
                     tag = {
@@ -165,7 +188,19 @@ def _parse(cls, path):
                     }
                     subset = el.attrib.get('subset')
             elif ev == 'end':
-                if el.tag == 'attribute' and attributes is not None:
+                if el.tag == 'attribute' and element_attributes is not None and shape_element is not None:
+                    attr_value = el.text or ''
+                    attr_type = attribute_types.get(el.attrib['name'])
+                    if el.text in ['true', 'false']:
+                        attr_value = attr_value == 'true'
+                    elif attr_type is not None and attr_type != 'text':
+                        try:
+                            attr_value = float(attr_value)
+                        except ValueError:
+                            pass
+                    element_attributes[el.attrib['name']] = attr_value
+
+                if el.tag == 'attribute' and attributes is not None and shape_element is None:
                     attr_value = el.text or ''
                     attr_type = attribute_types.get(el.attrib['name'])
                     if el.text in ['true', 'false']:
@@ -176,6 +211,37 @@ def _parse(cls, path):
                         except ValueError:
                             pass
                     attributes[el.attrib['name']] = attr_value
+
+                elif el.tag in cls._SUPPORTED_SHAPES and shape["type"] == "skeleton" and el.tag != "skeleton":
+                    shape_element['label'] = el.attrib.get('label')
+                    shape_element['group'] = int(el.attrib.get('group_id', 0))
+
+                    shape_element['type'] = el.tag
+                    shape_element['z_order'] = int(el.attrib.get('z_order', 0))
+
+                    if el.tag == 'box':
+                        shape_element['points'] = list(map(float, [
+                            el.attrib['xtl'], el.attrib['ytl'],
+                            el.attrib['xbr'], el.attrib['ybr'],
+                        ]))
+                    else:
+                        shape_element['points'] = []
+                        for pair in el.attrib['points'].split(';'):
+                            shape_element['points'].extend(map(float, pair.split(',')))
+
+                    if el.tag == 'points' and el.attrib.get('occluded') == '1':
+                        shape_element['visibility'] = [Points.Visibility.hidden] * (len(shape_element['points']) // 2)
+                    else:
+                        shape_element['occluded'] = (el.attrib.get('occluded') == '1')
+
+                    if el.tag == 'points' and el.attrib.get('outside') == '1':
+                        shape_element['visibility'] = [Points.Visibility.absent] * (len(shape_element['points']) // 2)
+                    else:
+                        shape_element['outside'] = (el.attrib.get('outside') == '1')
+
+                    shape['elements'].append(shape_element)
+                    shape_element = None
+
                 elif el.tag in cls._SUPPORTED_SHAPES:
                     if track is not None:
                         shape['frame'] = el.attrib['frame']
@@ -194,15 +260,22 @@ def _parse(cls, path):
                             el.attrib['xtl'], el.attrib['ytl'],
                             el.attrib['xbr'], el.attrib['ybr'],
                         ]))
+                    elif el.tag == 'skeleton':
+                        shape['points'] = []
                     else:
                         shape['points'] = []
                         for pair in el.attrib['points'].split(';'):
                             shape['points'].extend(map(float, pair.split(',')))
+                    if track_element:
+                        track_shapes[shape['frame']]['elements'].append(shape)
+                    elif track:
+                        track_shapes[shape['frame']] = shape
+                    else:
+                        frame_desc = items.get((subset, shape['frame']), {'annotations': []})
+                        frame_desc['annotations'].append(
+                            cls._parse_shape_ann(shape, categories))
+                        items[(subset, shape['frame'])] = frame_desc
 
-                    frame_desc = items.get((subset, shape['frame']), {'annotations': []})
-                    frame_desc['annotations'].append(
-                        cls._parse_shape_ann(shape, categories))
-                    items[(subset, shape['frame'])] = frame_desc
                     shape = None
 
                 elif el.tag == 'tag':
@@ -212,7 +285,15 @@ def _parse(cls, path):
                     items[(subset, tag['frame'])] = frame_desc
                     tag = None
                 elif el.tag == 'track':
-                    track = None
+                    if track_element:
+                        track_element = None
+                    else:
+                        for track_shape in track_shapes.values():
+                            frame_desc = items.get((subset, track_shape['frame']), {'annotations': []})
+                            frame_desc['annotations'].append(
+                                cls._parse_shape_ann(track_shape, categories))
+                            items[(subset, track_shape['frame'])] = frame_desc
+                        track = None
                 elif el.tag == 'image':
                     frame_desc = items.get((subset, image['frame']), {'annotations': []})
                     frame_desc.update({
@@ -377,7 +458,8 @@ def _parse_shape_ann(cls, ann, categories):
                 id=ann_id, attributes=attributes, group=group)
 
         elif ann_type == 'points':
-            return Points(points, label=label_id, z_order=z_order,
+            visibility = ann.get('visibility', None)
+            return Points(points, visibility, label=label_id, z_order=z_order,
                 id=ann_id, attributes=attributes, group=group)
 
         elif ann_type == 'box':
@@ -386,6 +468,14 @@ def _parse_shape_ann(cls, ann, categories):
             return Bbox(x, y, w, h, label=label_id, z_order=z_order,
                 id=ann_id, attributes=attributes, group=group)
 
+        elif ann_type == 'skeleton':
+            elements = []
+            for element in ann.get('elements', []):
+                elements.append(cls._parse_shape_ann(element, categories))
+
+            return Skeleton(elements, label=label_id, z_order=z_order,
+                id=ann_id, attributes=attributes, group=group)
+
         else:
             raise NotImplementedError("Unknown annotation type '%s'" % ann_type)
 
@@ -963,7 +1053,10 @@ def dump_track(idx, track):
                 elements=[],
             ) for element in shape.elements]
         }
-        if isinstance(annotations, ProjectData): track['task_id'] = shape.task_id
+        if isinstance(annotations, ProjectData):
+            track['task_id'] = shape.task_id
+            for element in track['elements']:
+                element.task_id = shape.task_id
         dump_track(counter, annotations.Track(**track))
         counter += 1