diff --git a/cvat/apps/annotation/cvat.py b/cvat/apps/annotation/cvat.py index 5e7b3fe3738..ff0b58dc6e8 100644 --- a/cvat/apps/annotation/cvat.py +++ b/cvat/apps/annotation/cvat.py @@ -124,6 +124,11 @@ def open_cuboid(self, cuboid): 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 close_cuboid(self): 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: - shape['attributes'].append(annotations.Attribute( + if el.tag == 'attribute' and attributes is not None: + 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() diff --git a/cvat/apps/dataset_manager/bindings.py b/cvat/apps/dataset_manager/bindings.py index afb0a0b71f1..f2b88723672 100644 --- a/cvat/apps/dataset_manager/bindings.py +++ b/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()], )) \ No newline at end of file diff --git a/cvat/apps/engine/tests/test_rest_api.py b/cvat/apps/engine/tests/test_rest_api.py index 680c44ed15a..b88070d8dfe 100644 --- a/cvat/apps/engine/tests/test_rest_api.py +++ b/cvat/apps/engine/tests/test_rest_api.py @@ -2638,6 +2638,47 @@ def _get_initial_annotation(annotation_format): "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 @@ def _get_initial_annotation(annotation_format): 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.0" or \ annotation_format == "YOLO ZIP 1.0" or \ diff --git a/datumaro/datumaro/plugins/cvat_format/converter.py b/datumaro/datumaro/plugins/cvat_format/converter.py index d18699d1d09..a64addad02b 100644 --- a/datumaro/datumaro/plugins/cvat_format/converter.py +++ b/datumaro/datumaro/plugins/cvat_format/converter.py @@ -113,6 +113,11 @@ def open_points(self, points): 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 @@ def close_polyline(self): 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 @@ def _write_item(self, item, index): 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 @@ def _write_shape(self, shape): 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 diff --git a/datumaro/datumaro/plugins/cvat_format/extractor.py b/datumaro/datumaro/plugins/cvat_format/extractor.py index 1f36f70cc77..407c551d014 100644 --- a/datumaro/datumaro/plugins/cvat_format/extractor.py +++ b/datumaro/datumaro/plugins/cvat_format/extractor.py @@ -9,7 +9,7 @@ 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 @@ def _parse(cls, path): track = None shape = None + tag = None + attributes = None image = None for ev, el in context: if ev == 'start': @@ -92,16 +94,25 @@ def _parse(cls, path): '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 @@ def _parse(cls, path): 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 @@ def _parse(cls, path): 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 @@ def consumed(expected_state, tag): 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 @@ def _parse_ann(cls, ann, categories): 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') diff --git a/datumaro/tests/test_cvat_format.py b/datumaro/tests/test_cvat_format.py index 17061217845..cc45bee921c 100644 --- a/datumaro/tests/test_cvat_format.py +++ b/datumaro/tests/test_cvat_format.py @@ -6,7 +6,7 @@ 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 @@ def __iter__(self): 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 @@ def __iter__(self): 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',