[Datumaro] CVAT format import (#974)
* Add label-specific attributes * Add CVAT format import * Register CVAT format * Add little more logs * Little refactoring for tests * Cvat format checks * Add missing check * Refactor datumaro format * Little refactoring * Regularize dataset importer logic * Fix project import issue * Refactor coco extractor * Refactor tests * Codacymain
zhiltsov-max
6 years ago
committed by
Nikita Manovich
parent
8edfe0dcb4
commit
36b1e9c1ce
@ -0,0 +1,286 @@
|
||||
|
||||
# Copyright (C) 2019 Intel Corporation
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
from collections import OrderedDict
|
||||
import os.path as osp
|
||||
import xml.etree as ET
|
||||
|
||||
from datumaro.components.extractor import (Extractor, DatasetItem,
|
||||
DEFAULT_SUBSET_NAME, AnnotationType,
|
||||
PointsObject, PolygonObject, PolyLineObject, BboxObject,
|
||||
LabelCategories
|
||||
)
|
||||
from datumaro.components.formats.cvat import CvatPath
|
||||
from datumaro.util.image import lazy_image
|
||||
|
||||
|
||||
class CvatExtractor(Extractor):
|
||||
_SUPPORTED_SHAPES = ('box', 'polygon', 'polyline', 'points')
|
||||
|
||||
def __init__(self, path):
|
||||
super().__init__()
|
||||
|
||||
assert osp.isfile(path)
|
||||
rootpath = path.rsplit(CvatPath.ANNOTATIONS_DIR, maxsplit=1)[0]
|
||||
self._path = rootpath
|
||||
|
||||
subset = osp.splitext(osp.basename(path))[0]
|
||||
if subset == DEFAULT_SUBSET_NAME:
|
||||
subset = None
|
||||
self._subset = subset
|
||||
|
||||
items, categories = self._parse(path)
|
||||
self._items = self._load_items(items)
|
||||
self._categories = categories
|
||||
|
||||
def categories(self):
|
||||
return self._categories
|
||||
|
||||
def __iter__(self):
|
||||
for item in self._items.values():
|
||||
yield item
|
||||
|
||||
def __len__(self):
|
||||
return len(self._items)
|
||||
|
||||
def subsets(self):
|
||||
if self._subset:
|
||||
return [self._subset]
|
||||
return None
|
||||
|
||||
def get_subset(self, name):
|
||||
if name != self._subset:
|
||||
return None
|
||||
return self
|
||||
|
||||
@classmethod
|
||||
def _parse(cls, path):
|
||||
context = ET.ElementTree.iterparse(path, events=("start", "end"))
|
||||
context = iter(context)
|
||||
|
||||
categories = cls._parse_meta(context)
|
||||
|
||||
items = OrderedDict()
|
||||
|
||||
track = None
|
||||
shape = None
|
||||
image = None
|
||||
for ev, el in context:
|
||||
if ev == 'start':
|
||||
if el.tag == 'track':
|
||||
track = {
|
||||
'id': el.attrib.get('id'),
|
||||
'label': el.attrib.get('label'),
|
||||
'group': int(el.attrib.get('group_id', 0)),
|
||||
}
|
||||
elif el.tag == 'image':
|
||||
image = {
|
||||
'name': el.attrib.get('name'),
|
||||
'frame': el.attrib['id'],
|
||||
}
|
||||
elif el.tag in cls._SUPPORTED_SHAPES and (track or image):
|
||||
shape = {
|
||||
'type': None,
|
||||
'attributes': {},
|
||||
}
|
||||
if track:
|
||||
shape.update(track)
|
||||
if image:
|
||||
shape.update(image)
|
||||
elif ev == 'end':
|
||||
if el.tag == 'attribute' and shape is not None:
|
||||
shape['attributes'][el.attrib['name']] = el.text
|
||||
elif el.tag in cls._SUPPORTED_SHAPES:
|
||||
if track is not None:
|
||||
shape['frame'] = el.attrib['frame']
|
||||
shape['outside'] = (el.attrib.get('outside') == '1')
|
||||
shape['keyframe'] = (el.attrib.get('keyframe') == '1')
|
||||
if image is not None:
|
||||
shape['label'] = el.attrib.get('label')
|
||||
shape['group'] = int(el.attrib.get('group_id', 0))
|
||||
|
||||
shape['type'] = el.tag
|
||||
shape['occluded'] = (el.attrib.get('occluded') == '1')
|
||||
shape['z_order'] = int(el.attrib.get('z_order', 0))
|
||||
|
||||
if el.tag == 'box':
|
||||
shape['points'] = list(map(float, [
|
||||
el.attrib['xtl'], el.attrib['ytl'],
|
||||
el.attrib['xbr'], el.attrib['ybr'],
|
||||
]))
|
||||
else:
|
||||
shape['points'] = []
|
||||
for pair in el.attrib['points'].split(';'):
|
||||
shape['points'].extend(map(float, pair.split(',')))
|
||||
|
||||
frame_desc = items.get(shape['frame'], {
|
||||
'name': shape.get('name'),
|
||||
'annotations': [],
|
||||
})
|
||||
frame_desc['annotations'].append(
|
||||
cls._parse_ann(shape, categories))
|
||||
items[shape['frame']] = frame_desc
|
||||
shape = None
|
||||
|
||||
elif el.tag == 'track':
|
||||
track = None
|
||||
elif el.tag == 'image':
|
||||
image = None
|
||||
el.clear()
|
||||
|
||||
return items, categories
|
||||
|
||||
@staticmethod
|
||||
def _parse_meta(context):
|
||||
ev, el = next(context)
|
||||
if not (ev == 'start' and el.tag == 'annotations'):
|
||||
raise Exception("Unexpected token ")
|
||||
|
||||
categories = {}
|
||||
|
||||
has_z_order = False
|
||||
mode = 'annotation'
|
||||
labels = OrderedDict()
|
||||
label = None
|
||||
|
||||
# Recursive descent parser
|
||||
el = None
|
||||
states = ['annotations']
|
||||
def accepted(expected_state, tag, next_state=None):
|
||||
state = states[-1]
|
||||
if state == expected_state and el is not None and el.tag == tag:
|
||||
if not next_state:
|
||||
next_state = tag
|
||||
states.append(next_state)
|
||||
return True
|
||||
return False
|
||||
def consumed(expected_state, tag):
|
||||
state = states[-1]
|
||||
if state == expected_state and el is not None and el.tag == tag:
|
||||
states.pop()
|
||||
return True
|
||||
return False
|
||||
|
||||
for ev, el in context:
|
||||
if ev == 'start':
|
||||
if accepted('annotations', 'meta'): pass
|
||||
elif accepted('meta', 'task'): pass
|
||||
elif accepted('task', 'z_order'): pass
|
||||
elif accepted('task', 'labels'): pass
|
||||
elif accepted('labels', 'label'):
|
||||
label = { 'name': None, 'attributes': set() }
|
||||
elif accepted('label', 'name', next_state='label_name'): pass
|
||||
elif accepted('label', 'attributes'): pass
|
||||
elif accepted('attributes', 'attribute'): pass
|
||||
elif accepted('attribute', 'name', next_state='attr_name'): pass
|
||||
elif accepted('annotations', 'image') or \
|
||||
accepted('annotations', 'track') or \
|
||||
accepted('annotations', 'tag'):
|
||||
break
|
||||
else:
|
||||
pass
|
||||
elif ev == 'end':
|
||||
if consumed('meta', 'meta'):
|
||||
break
|
||||
elif consumed('task', 'task'): pass
|
||||
elif consumed('z_order', 'z_order'):
|
||||
has_z_order = (el.text == 'True')
|
||||
elif consumed('label_name', 'name'):
|
||||
label['name'] = el.text
|
||||
elif consumed('attr_name', 'name'):
|
||||
label['attributes'].add(el.text)
|
||||
elif consumed('attribute', 'attribute'): pass
|
||||
elif consumed('attributes', 'attributes'): pass
|
||||
elif consumed('label', 'label'):
|
||||
labels[label['name']] = label['attributes']
|
||||
label = None
|
||||
elif consumed('labels', 'labels'): pass
|
||||
else:
|
||||
pass
|
||||
|
||||
assert len(states) == 1 and states[0] == 'annotations', \
|
||||
"Expected 'meta' section in the annotation file, path: %s" % states
|
||||
|
||||
common_attrs = ['occluded']
|
||||
if has_z_order:
|
||||
common_attrs.append('z_order')
|
||||
if mode == 'interpolation':
|
||||
common_attrs.append('keyframe')
|
||||
common_attrs.append('outside')
|
||||
|
||||
label_cat = LabelCategories(attributes=common_attrs)
|
||||
for label, attrs in labels.items():
|
||||
label_cat.add(label, attributes=attrs)
|
||||
|
||||
categories[AnnotationType.label] = label_cat
|
||||
|
||||
return categories
|
||||
|
||||
@classmethod
|
||||
def _parse_ann(cls, ann, categories):
|
||||
ann_id = ann.get('id')
|
||||
ann_type = ann['type']
|
||||
|
||||
attributes = ann.get('attributes', {})
|
||||
if 'occluded' in categories[AnnotationType.label].attributes:
|
||||
attributes['occluded'] = ann.get('occluded', False)
|
||||
if 'z_order' in categories[AnnotationType.label].attributes:
|
||||
attributes['z_order'] = ann.get('z_order', 0)
|
||||
if 'outside' in categories[AnnotationType.label].attributes:
|
||||
attributes['outside'] = ann.get('outside', False)
|
||||
if 'keyframe' in categories[AnnotationType.label].attributes:
|
||||
attributes['keyframe'] = ann.get('keyframe', False)
|
||||
|
||||
group = ann.get('group')
|
||||
if group == 0:
|
||||
group = None
|
||||
|
||||
label = ann.get('label')
|
||||
label_id = categories[AnnotationType.label].find(label)[0]
|
||||
|
||||
points = ann.get('points', [])
|
||||
|
||||
if ann_type == 'polyline':
|
||||
return PolyLineObject(points, label=label_id,
|
||||
id=ann_id, attributes=attributes, group=group)
|
||||
|
||||
elif ann_type == 'polygon':
|
||||
return PolygonObject(points, label=label_id,
|
||||
id=ann_id, attributes=attributes, group=group)
|
||||
|
||||
elif ann_type == 'points':
|
||||
return PointsObject(points, label=label_id,
|
||||
id=ann_id, attributes=attributes, group=group)
|
||||
|
||||
elif ann_type == 'box':
|
||||
x, y = points[0], points[1]
|
||||
w, h = points[2] - x, points[3] - y
|
||||
return BboxObject(x, y, w, h, label=label_id,
|
||||
id=ann_id, attributes=attributes, group=group)
|
||||
|
||||
else:
|
||||
raise NotImplementedError("Unknown annotation type '%s'" % ann_type)
|
||||
|
||||
def _load_items(self, parsed):
|
||||
for item_id, item_desc in parsed.items():
|
||||
file_name = item_desc.get('name')
|
||||
if not file_name:
|
||||
file_name = item_id
|
||||
file_name += CvatPath.IMAGE_EXT
|
||||
image = self._find_image(file_name)
|
||||
|
||||
parsed[item_id] = DatasetItem(id=item_id, subset=self._subset,
|
||||
image=image, annotations=item_desc.get('annotations', None))
|
||||
return parsed
|
||||
|
||||
def _find_image(self, file_name):
|
||||
images_dir = osp.join(self._path, CvatPath.IMAGES_DIR)
|
||||
search_paths = [
|
||||
osp.join(images_dir, file_name),
|
||||
osp.join(images_dir, self._subset or DEFAULT_SUBSET_NAME, file_name),
|
||||
]
|
||||
for image_path in search_paths:
|
||||
if osp.exists(image_path):
|
||||
return lazy_image(image_path)
|
||||
@ -0,0 +1,10 @@
|
||||
|
||||
# Copyright (C) 2019 Intel Corporation
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
class CvatPath:
|
||||
IMAGES_DIR = 'images'
|
||||
ANNOTATIONS_DIR = 'annotations'
|
||||
|
||||
IMAGE_EXT = '.jpg'
|
||||
@ -0,0 +1,46 @@
|
||||
|
||||
# Copyright (C) 2019 Intel Corporation
|
||||
#
|
||||
# SPDX-License-Identifier: MIT
|
||||
|
||||
from glob import glob
|
||||
import logging as log
|
||||
import os.path as osp
|
||||
|
||||
from datumaro.components.formats.cvat import CvatPath
|
||||
|
||||
|
||||
class CvatImporter:
|
||||
EXTRACTOR_NAME = 'cvat'
|
||||
|
||||
def __call__(self, path, **extra_params):
|
||||
from datumaro.components.project import Project # cyclic import
|
||||
project = Project()
|
||||
|
||||
if path.endswith('.xml') and osp.isfile(path):
|
||||
subset_paths = [path]
|
||||
else:
|
||||
subset_paths = glob(osp.join(path, '*.xml'))
|
||||
|
||||
if osp.basename(osp.normpath(path)) != CvatPath.ANNOTATIONS_DIR:
|
||||
path = osp.join(path, CvatPath.ANNOTATIONS_DIR)
|
||||
subset_paths += glob(osp.join(path, '*.xml'))
|
||||
|
||||
if len(subset_paths) == 0:
|
||||
raise Exception("Failed to find 'cvat' dataset at '%s'" % path)
|
||||
|
||||
for subset_path in subset_paths:
|
||||
if not osp.isfile(subset_path):
|
||||
continue
|
||||
|
||||
log.info("Found a dataset at '%s'" % subset_path)
|
||||
|
||||
subset_name = osp.splitext(osp.basename(subset_path))[0]
|
||||
|
||||
project.add_source(subset_name, {
|
||||
'url': subset_path,
|
||||
'format': self.EXTRACTOR_NAME,
|
||||
'options': extra_params,
|
||||
})
|
||||
|
||||
return project
|
||||
@ -0,0 +1,148 @@
|
||||
import numpy as np
|
||||
import os
|
||||
import os.path as osp
|
||||
from xml.etree import ElementTree as ET
|
||||
|
||||
from unittest import TestCase
|
||||
|
||||
from datumaro.components.extractor import (Extractor, DatasetItem,
|
||||
AnnotationType, PointsObject, PolygonObject, PolyLineObject, BboxObject,
|
||||
LabelCategories,
|
||||
)
|
||||
from datumaro.components.importers.cvat import CvatImporter
|
||||
import datumaro.components.formats.cvat as Cvat
|
||||
from datumaro.util.image import save_image
|
||||
from datumaro.util.test_utils import TestDir
|
||||
|
||||
|
||||
class CvatExtractorTest(TestCase):
|
||||
@staticmethod
|
||||
def generate_dummy_cvat(path):
|
||||
images_dir = osp.join(path, Cvat.CvatPath.IMAGES_DIR)
|
||||
anno_dir = osp.join(path, Cvat.CvatPath.ANNOTATIONS_DIR)
|
||||
|
||||
os.makedirs(images_dir)
|
||||
os.makedirs(anno_dir)
|
||||
|
||||
root_elem = ET.Element('annotations')
|
||||
ET.SubElement(root_elem, 'version').text = '1.1'
|
||||
|
||||
meta_elem = ET.SubElement(root_elem, 'meta')
|
||||
task_elem = ET.SubElement(meta_elem, 'task')
|
||||
ET.SubElement(task_elem, 'z_order').text = 'True'
|
||||
ET.SubElement(task_elem, 'mode').text = 'interpolation'
|
||||
|
||||
labels_elem = ET.SubElement(task_elem, 'labels')
|
||||
|
||||
label1_elem = ET.SubElement(labels_elem, 'label')
|
||||
ET.SubElement(label1_elem, 'name').text = 'label1'
|
||||
label1_attrs_elem = ET.SubElement(label1_elem, 'attributes')
|
||||
|
||||
label1_a1_elem = ET.SubElement(label1_attrs_elem, 'attribute')
|
||||
ET.SubElement(label1_a1_elem, 'name').text = 'a1'
|
||||
ET.SubElement(label1_a1_elem, 'input_type').text = 'checkbox'
|
||||
ET.SubElement(label1_a1_elem, 'default_value').text = 'false'
|
||||
ET.SubElement(label1_a1_elem, 'values').text = 'false\ntrue'
|
||||
|
||||
label1_a2_elem = ET.SubElement(label1_attrs_elem, 'attribute')
|
||||
ET.SubElement(label1_a2_elem, 'name').text = 'a2'
|
||||
ET.SubElement(label1_a2_elem, 'input_type').text = 'radio'
|
||||
ET.SubElement(label1_a2_elem, 'default_value').text = 'v1'
|
||||
ET.SubElement(label1_a2_elem, 'values').text = 'v1\nv2\nv3'
|
||||
|
||||
label2_elem = ET.SubElement(labels_elem, 'label')
|
||||
ET.SubElement(label2_elem, 'name').text = 'label2'
|
||||
|
||||
# item 1
|
||||
save_image(osp.join(images_dir, 'img0.jpg'), np.ones((8, 8, 3)))
|
||||
item1_elem = ET.SubElement(root_elem, 'image')
|
||||
item1_elem.attrib.update({
|
||||
'id': '0', 'name': 'img0', 'width': '8', 'height': '8'
|
||||
})
|
||||
|
||||
item1_ann1_elem = ET.SubElement(item1_elem, 'box')
|
||||
item1_ann1_elem.attrib.update({
|
||||
'label': 'label1', 'occluded': '1', 'z_order': '1',
|
||||
'xtl': '0', 'ytl': '2', 'xbr': '4', 'ybr': '4'
|
||||
})
|
||||
item1_ann1_a1_elem = ET.SubElement(item1_ann1_elem, 'attribute')
|
||||
item1_ann1_a1_elem.attrib['name'] = 'a1'
|
||||
item1_ann1_a1_elem.text = 'true'
|
||||
item1_ann1_a2_elem = ET.SubElement(item1_ann1_elem, 'attribute')
|
||||
item1_ann1_a2_elem.attrib['name'] = 'a2'
|
||||
item1_ann1_a2_elem.text = 'v3'
|
||||
|
||||
item1_ann2_elem = ET.SubElement(item1_elem, 'polyline')
|
||||
item1_ann2_elem.attrib.update({
|
||||
'label': '', 'points': '1.0,2;3,4;5,6;7,8'
|
||||
})
|
||||
|
||||
# item 2
|
||||
save_image(osp.join(images_dir, 'img1.jpg'), np.ones((10, 10, 3)))
|
||||
item2_elem = ET.SubElement(root_elem, 'image')
|
||||
item2_elem.attrib.update({
|
||||
'id': '1', 'name': 'img1', 'width': '8', 'height': '8'
|
||||
})
|
||||
|
||||
item2_ann1_elem = ET.SubElement(item2_elem, 'polygon')
|
||||
item2_ann1_elem.attrib.update({
|
||||
'label': '', 'points': '1,2;3,4;6,5', 'z_order': '1',
|
||||
})
|
||||
|
||||
item2_ann2_elem = ET.SubElement(item2_elem, 'points')
|
||||
item2_ann2_elem.attrib.update({
|
||||
'label': 'label2', 'points': '1,2;3,4;5,6', 'z_order': '2',
|
||||
})
|
||||
|
||||
with open(osp.join(anno_dir, 'train.xml'), 'w') as f:
|
||||
f.write(ET.tostring(root_elem, encoding='unicode'))
|
||||
|
||||
def test_can_load(self):
|
||||
class TestExtractor(Extractor):
|
||||
def __iter__(self):
|
||||
return iter([
|
||||
DatasetItem(id=1, subset='train', image=np.ones((8, 8, 3)),
|
||||
annotations=[
|
||||
BboxObject(0, 2, 4, 2, label=0,
|
||||
attributes={
|
||||
'occluded': True, 'z_order': 1,
|
||||
'a1': 'true', 'a2': 'v3'
|
||||
}),
|
||||
PolyLineObject([1, 2, 3, 4, 5, 6, 7, 8],
|
||||
attributes={'occluded': False, 'z_order': 0}),
|
||||
]),
|
||||
DatasetItem(id=2, subset='train', image=np.ones((10, 10, 3)),
|
||||
annotations=[
|
||||
PolygonObject([1, 2, 3, 4, 6, 5],
|
||||
attributes={'occluded': False, 'z_order': 1}),
|
||||
PointsObject([1, 2, 3, 4, 5, 6], label=1,
|
||||
attributes={'occluded': False, 'z_order': 2}),
|
||||
]),
|
||||
])
|
||||
|
||||
def categories(self):
|
||||
label_categories = LabelCategories()
|
||||
for i in range(10):
|
||||
label_categories.add('label_' + str(i))
|
||||
return {
|
||||
AnnotationType.label: label_categories,
|
||||
}
|
||||
|
||||
with TestDir() as test_dir:
|
||||
self.generate_dummy_cvat(test_dir.path)
|
||||
source_dataset = TestExtractor()
|
||||
|
||||
parsed_dataset = CvatImporter()(test_dir.path).make_dataset()
|
||||
|
||||
self.assertListEqual(
|
||||
sorted(source_dataset.subsets()),
|
||||
sorted(parsed_dataset.subsets()),
|
||||
)
|
||||
self.assertEqual(len(source_dataset), len(parsed_dataset))
|
||||
for subset_name in source_dataset.subsets():
|
||||
source_subset = source_dataset.get_subset(subset_name)
|
||||
parsed_subset = parsed_dataset.get_subset(subset_name)
|
||||
for item_a, item_b in zip(source_subset, parsed_subset):
|
||||
self.assertEqual(len(item_a.annotations), len(item_b.annotations))
|
||||
for ann_a, ann_b in zip(item_a.annotations, item_b.annotations):
|
||||
self.assertEqual(ann_a, ann_b)
|
||||
Loading…
Reference in New Issue