Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add tags to cvat xml #1200

Merged
merged 4 commits into from
Mar 1, 2020
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 43 additions & 3 deletions cvat/apps/annotation/cvat.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]})
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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()

Expand Down Expand Up @@ -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':
Expand All @@ -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,
))
Expand Down Expand Up @@ -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()
11 changes: 10 additions & 1 deletion cvat/apps/dataset_manager/bindings.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()],
))
45 changes: 44 additions & 1 deletion cvat/apps/engine/tests/test_rest_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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": [],
Expand All @@ -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 \
Expand Down
32 changes: 32 additions & 0 deletions datumaro/datumaro/plugins/cvat_format/converter.py
Original file line number Diff line number Diff line change
Expand Up @@ -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']})
Expand All @@ -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')

Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down
37 changes: 31 additions & 6 deletions datumaro/datumaro/plugins/cvat_format/extractor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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':
Expand All @@ -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'
Expand All @@ -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']
Expand All @@ -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':
Expand Down Expand Up @@ -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']

Expand Down Expand Up @@ -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')
Expand Down
6 changes: 5 additions & 1 deletion datumaro/tests/test_cvat_format.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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',
Expand Down Expand Up @@ -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',
Expand Down