Add tags to cvat xml (#1200)

* Extend cvat format test

* Add tags to cvat for images

* Add tags to cvat format in dm

* Add import of tags from datumaro

cvat/apps/annotation/cvat.py

@@ -124,6 +124,11 @@ def create_xml_dumper(file_object):
             self.xmlgen.startElement("cuboid", cuboid)
             self._level += 1
 
+        def open_tag(self, tag):
+            self._indent()
+            self.xmlgen.startElement("tag", tag)
+            self._level += 1
+
         def add_attribute(self, attribute):
             self._indent()
             self.xmlgen.startElement("attribute", {"name": attribute["name"]})
@@ -155,6 +160,11 @@ def create_xml_dumper(file_object):
             self._indent()
             self.xmlgen.endElement("cuboid")
 
+        def close_tag(self):
+            self._level -= 1
+            self._indent()
+            self.xmlgen.endElement("tag")
+
         def close_image(self):
             self._level -= 1
             self._indent()
@@ -268,6 +278,22 @@ def dump_as_cvat_annotation(file_object, annotations):
             else:
                 raise NotImplementedError("unknown shape type")
 
+        for tag in frame_annotation.tags:
+            tag_data = OrderedDict([
+                ("label", tag.label),
+            ])
+            if tag.group:
+                tag_data["group_id"] = str(tag.group)
+            dumper.open_tag(tag_data)
+
+            for attr in tag.attributes:
+                dumper.add_attribute(OrderedDict([
+                    ("name", attr.name),
+                    ("value", attr.value)
+                ]))
+
+            dumper.close_tag()
+
         dumper.close_image()
     dumper.close_root()
 
@@ -408,7 +434,9 @@ def load(file_object, annotations):
 
     track = None
     shape = None
+    tag = None
     image_is_opened = False
+    attributes = None
     for ev, el in context:
         if ev == 'start':
             if el.tag == 'track':
@@ -421,13 +449,22 @@ def load(file_object, annotations):
                 image_is_opened = True
                 frame_id = int(el.attrib['id'])
             elif el.tag in supported_shapes and (track is not None or image_is_opened):
+                attributes = []
                 shape = {
-                    'attributes': [],
+                    'attributes': attributes,
                     'points': [],
                 }
+            elif el.tag == 'tag' and image_is_opened:
+                attributes = []
+                tag = {
+                    'frame': frame_id,
+                    'label': el.attrib['label'],
+                    'group': int(el.attrib.get('group_id', 0)),
+                    'attributes': attributes,
+                }
         elif ev == 'end':
-            if el.tag == 'attribute' and shape is not None:
+            if el.tag == 'attribute' and attributes is not None:
-                shape['attributes'].append(annotations.Attribute(
+                attributes.append(annotations.Attribute(
                     name=el.attrib['name'],
                     value=el.text,
                 ))
@@ -484,4 +521,7 @@ def load(file_object, annotations):
                 track = None
             elif el.tag == 'image':
                 image_is_opened = False
+            elif el.tag == 'tag':
+                annotations.add_tag(annotations.Tag(**tag))
+                tag = None
             el.clear()

cvat/apps/dataset_manager/bindings.py

@@ -236,5 +236,14 @@ def import_dm_annotations(dm_dataset, cvat_task_anno):
                     points=ann.points,
                     occluded=False,
                     group=group_map.get(ann.group, 0),
-                    attributes=[],
+                    attributes=[cvat_task_anno.Attribute(name=n, value=str(v))
+                        for n, v in ann.attributes.items()],
+                ))
+            elif ann.type == datumaro.AnnotationType.label:
+                cvat_task_anno.add_shape(cvat_task_anno.Tag(
+                    frame=frame_number,
+                    label=label_cat.items[ann.label].name,
+                    group=group_map.get(ann.group, 0),
+                    attributes=[cvat_task_anno.Attribute(name=n, value=str(v))
+                        for n, v in ann.attributes.items()],
                 ))

cvat/apps/engine/tests/test_rest_api.py

@@ -2638,6 +2638,47 @@ class TaskAnnotationAPITestCase(JobAnnotationAPITestCase):
                 "occluded": False
             }]
 
+            polygon_shapes_with_attrs = [{
+                "frame": 2,
+                "label_id": task["labels"][0]["id"],
+                "group": 1,
+                "attributes": [
+                    {
+                        "spec_id": task["labels"][0]["attributes"][0]["id"],
+                        "value": task["labels"][0]["attributes"][0]["values"][1]
+                    },
+                    {
+                        "spec_id": task["labels"][0]["attributes"][1]["id"],
+                        "value": task["labels"][0]["attributes"][1]["default_value"]
+                    }
+                ],
+                "points": [20.0, 0.1, 10, 3.22, 4, 7, 10, 30, 1, 2, 4.44, 5.55],
+                "type": "polygon",
+                "occluded": True
+            }]
+
+            tags_wo_attrs = [{
+                "frame": 2,
+                "label_id": task["labels"][1]["id"],
+                "group": 3,
+                "attributes": []
+            }]
+            tags_with_attrs = [{
+                "frame": 1,
+                "label_id": task["labels"][0]["id"],
+                "group": 0,
+                "attributes": [
+                    {
+                        "spec_id": task["labels"][0]["attributes"][0]["id"],
+                        "value": task["labels"][0]["attributes"][0]["values"][1]
+                    },
+                    {
+                        "spec_id": task["labels"][0]["attributes"][1]["id"],
+                        "value": task["labels"][0]["attributes"][1]["default_value"]
+                    }
+                ],
+            }]
+
             annotations = {
                     "version": 0,
                     "tags": [],
@@ -2648,7 +2689,9 @@ class TaskAnnotationAPITestCase(JobAnnotationAPITestCase):
                 annotations["tracks"] = rectangle_tracks_with_attrs + rectangle_tracks_wo_attrs
 
             elif annotation_format == "CVAT XML 1.1 for images":
-                annotations["shapes"] = rectangle_shapes_with_attrs + rectangle_shapes_wo_attrs
+                annotations["shapes"] = rectangle_shapes_with_attrs + rectangle_shapes_wo_attrs \
+                    + polygon_shapes_wo_attrs + polygon_shapes_with_attrs
+                annotations["tags"] = tags_with_attrs + tags_wo_attrs
 
             elif annotation_format == "PASCAL VOC ZIP 1.1" or \
                  annotation_format == "YOLO ZIP 1.1" or \

datumaro/datumaro/plugins/cvat_format/converter.py

@@ -113,6 +113,11 @@ class XmlAnnotationWriter:
         self.xmlgen.startElement('points', points)
         self._level += 1
 
+    def open_tag(self, tag):
+        self._indent()
+        self.xmlgen.startElement("tag", tag)
+        self._level += 1
+
     def add_attribute(self, attribute):
         self._indent()
         self.xmlgen.startElement('attribute', {'name': attribute['name']})
@@ -136,6 +141,9 @@ class XmlAnnotationWriter:
     def close_points(self):
         self._close_element('points')
 
+    def close_tag(self):
+        self._close_element('tag')
+
     def close_image(self):
         self._close_element('image')
 
@@ -201,6 +209,8 @@ class _SubsetWriter:
             if ann.type in {AnnotationType.points, AnnotationType.polyline,
                     AnnotationType.polygon, AnnotationType.bbox}:
                 self._write_shape(ann)
+            elif ann.type == AnnotationType.label:
+                self._write_tag(ann)
             else:
                 continue
 
@@ -303,6 +313,28 @@ class _SubsetWriter:
         else:
             raise NotImplementedError("unknown shape type")
 
+    def _write_tag(self, label):
+        if label.label is None:
+            return
+
+        tag_data = OrderedDict([
+            ('label', self._get_label(label.label).name),
+        ])
+        if label.group:
+            tag_data['group_id'] = str(label.group)
+        self._writer.open_tag(tag_data)
+
+        for attr_name, attr_value in label.attributes.items():
+            if isinstance(attr_value, bool):
+                attr_value = 'true' if attr_value else 'false'
+            if attr_name in self._get_label(label.label).attributes:
+                self._writer.add_attribute(OrderedDict([
+                    ("name", str(attr_name)),
+                    ("value", str(attr_value)),
+                ]))
+
+        self._writer.close_tag()
+
 class _Converter:
     def __init__(self, extractor, save_dir, save_images=False):
         self._extractor = extractor

datumaro/datumaro/plugins/cvat_format/extractor.py

@@ -9,7 +9,7 @@ import xml.etree as ET
 
 from datumaro.components.extractor import (SourceExtractor,
     DEFAULT_SUBSET_NAME, DatasetItem,
-    AnnotationType, Points, Polygon, PolyLine, Bbox,
+    AnnotationType, Points, Polygon, PolyLine, Bbox, Label,
     LabelCategories
 )
 from datumaro.util.image import Image
@@ -73,6 +73,8 @@ class CvatExtractor(SourceExtractor):
 
         track = None
         shape = None
+        tag = None
+        attributes = None
         image = None
         for ev, el in context:
             if ev == 'start':
@@ -92,16 +94,25 @@ class CvatExtractor(SourceExtractor):
                         'height': el.attrib.get('height'),
                     }
                 elif el.tag in cls._SUPPORTED_SHAPES and (track or image):
+                    attributes = {}
                     shape = {
                         'type': None,
-                        'attributes': {},
+                        'attributes': attributes,
                     }
                     if track:
                         shape.update(track)
                     if image:
                         shape.update(image)
+                elif el.tag == 'tag' and image:
+                    attributes = {}
+                    tag = {
+                        'frame': image['frame'],
+                        'attributes': attributes,
+                        'group': int(el.attrib.get('group_id', 0)),
+                        'label': el.attrib['label'],
+                    }
             elif ev == 'end':
-                if el.tag == 'attribute' and shape is not None:
+                if el.tag == 'attribute' and attributes is not None:
                     attr_value = el.text
                     if el.text in ['true', 'false']:
                         attr_value = attr_value == 'true'
@@ -110,7 +121,7 @@ class CvatExtractor(SourceExtractor):
                             attr_value = float(attr_value)
                         except Exception:
                             pass
-                    shape['attributes'][el.attrib['name']] = attr_value
+                    attributes[el.attrib['name']] = attr_value
                 elif el.tag in cls._SUPPORTED_SHAPES:
                     if track is not None:
                         shape['frame'] = el.attrib['frame']
@@ -136,10 +147,16 @@ class CvatExtractor(SourceExtractor):
 
                     frame_desc = items.get(shape['frame'], {'annotations': []})
                     frame_desc['annotations'].append(
-                        cls._parse_ann(shape, categories))
+                        cls._parse_shape_ann(shape, categories))
                     items[shape['frame']] = frame_desc
                     shape = None
 
+                elif el.tag == 'tag':
+                    frame_desc = items.get(tag['frame'], {'annotations': []})
+                    frame_desc['annotations'].append(
+                        cls._parse_tag_ann(tag, categories))
+                    items[tag['frame']] = frame_desc
+                    tag = None
                 elif el.tag == 'track':
                     track = None
                 elif el.tag == 'image':
@@ -252,7 +269,7 @@ class CvatExtractor(SourceExtractor):
         return categories, frame_size
 
     @classmethod
-    def _parse_ann(cls, ann, categories):
+    def _parse_shape_ann(cls, ann, categories):
         ann_id = ann.get('id')
         ann_type = ann['type']
 
@@ -294,6 +311,14 @@ class CvatExtractor(SourceExtractor):
         else:
             raise NotImplementedError("Unknown annotation type '%s'" % ann_type)
 
+    @classmethod
+    def _parse_tag_ann(cls, ann, categories):
+        label = ann.get('label')
+        label_id = categories[AnnotationType.label].find(label)[0]
+        group = ann.get('group')
+        attributes = ann.get('attributes')
+        return Label(label_id, attributes=attributes, group=group)
+
     def _load_items(self, parsed):
         for frame_id, item_desc in parsed.items():
             filename = item_desc.get('name')

datumaro/tests/test_cvat_format.py

@@ -6,7 +6,7 @@ from xml.etree import ElementTree as ET
 from unittest import TestCase
 
 from datumaro.components.extractor import (Extractor, DatasetItem,
-    AnnotationType, Points, Polygon, PolyLine, Bbox,
+    AnnotationType, Points, Polygon, PolyLine, Bbox, Label,
     LabelCategories,
 )
 from datumaro.plugins.cvat_format.importer import CvatImporter
@@ -173,6 +173,8 @@ class CvatConverterTest(TestCase):
                             Points([1, 1, 3, 2, 2, 3],
                                 label=2,
                                 attributes={ 'a1': 'x', 'a2': 42 }),
+                            Label(1),
+                            Label(2, attributes={ 'a1': 'y', 'a2': 44 }),
                         ]
                     ),
                     DatasetItem(id=1, subset='s1',
@@ -215,6 +217,8 @@ class CvatConverterTest(TestCase):
                                 label=2,
                                 attributes={ 'z_order': 0, 'occluded': False,
                                     'a1': 'x', 'a2': 42 }),
+                            Label(1),
+                            Label(2, attributes={ 'a1': 'y', 'a2': 44 }),
                         ]
                     ),
                     DatasetItem(id=1, subset='s1',