Add assets importing tests for cvat format (#1652)

main
zhiltsov-max 6 years ago committed by GitHub
parent 441ce14093
commit 696a47b3c5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -1,7 +1,5 @@
import numpy as np import numpy as np
import os
import os.path as osp import os.path as osp
from xml.etree import ElementTree as ET
from unittest import TestCase from unittest import TestCase
@ -11,111 +9,40 @@ from datumaro.components.extractor import (Extractor, DatasetItem,
) )
from datumaro.plugins.cvat_format.importer import CvatImporter from datumaro.plugins.cvat_format.importer import CvatImporter
from datumaro.plugins.cvat_format.converter import CvatConverter from datumaro.plugins.cvat_format.converter import CvatConverter
from datumaro.plugins.cvat_format.format import CvatPath from datumaro.util.image import Image
from datumaro.util.image import save_image, Image
from datumaro.util.test_utils import TestDir, compare_datasets from datumaro.util.test_utils import TestDir, compare_datasets
def generate_dummy_cvat(path): DUMMY_IMAGE_DATASET_DIR = osp.join(osp.dirname(__file__),
images_dir = osp.join(path, CvatPath.IMAGES_DIR) 'assets', 'cvat_dataset', 'for_images')
os.makedirs(images_dir)
root_elem = ET.Element('annotations') DUMMY_VIDEO_DATASET_DIR = osp.join(osp.dirname(__file__),
ET.SubElement(root_elem, 'version').text = '1.1' 'assets', 'cvat_dataset', 'for_video')
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': '10', 'height': '10'
})
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(path, 'train.xml'), 'w') as f:
f.write(ET.tostring(root_elem, encoding='unicode'))
class CvatImporterTest(TestCase): class CvatImporterTest(TestCase):
def test_can_detect(self): def test_can_detect_image(self):
with TestDir() as test_dir: self.assertTrue(CvatImporter.detect(DUMMY_IMAGE_DATASET_DIR))
generate_dummy_cvat(test_dir)
self.assertTrue(CvatImporter.detect(test_dir)) def test_can_detect_video(self):
self.assertTrue(CvatImporter.detect(DUMMY_VIDEO_DATASET_DIR))
class CvatExtractorTest(TestCase): def test_can_load_image(self):
def test_can_load(self): class DstExtractor(Extractor):
class TestExtractor(Extractor):
def __iter__(self): def __iter__(self):
return iter([ return iter([
DatasetItem(id=0, subset='train', image=np.ones((8, 8, 3)), DatasetItem(id=0, subset='train',
image=np.ones((8, 8, 3)),
annotations=[ annotations=[
Bbox(0, 2, 4, 2, label=0, z_order=1, Bbox(0, 2, 4, 2, label=0, z_order=1,
attributes={ attributes={
'occluded': True, 'occluded': True,
'a1': True, 'a2': 'v3' 'a1': True, 'a2': 'v3'
}), }),
PolyLine([1, 2, 3, 4, 5, 6, 7, 8], z_order=0, PolyLine([1, 2, 3, 4, 5, 6, 7, 8],
attributes={'occluded': False}), attributes={'occluded': False}),
]), ]),
DatasetItem(id=1, subset='train', image=np.ones((10, 10, 3)), DatasetItem(id=1, subset='train',
image=np.ones((10, 10, 3)),
annotations=[ annotations=[
Polygon([1, 2, 3, 4, 6, 5], z_order=1, Polygon([1, 2, 3, 4, 6, 5], z_order=1,
attributes={'occluded': False}), attributes={'occluded': False}),
@ -128,18 +55,94 @@ class CvatExtractorTest(TestCase):
label_categories = LabelCategories() label_categories = LabelCategories()
label_categories.add('label1', attributes={'a1', 'a2'}) label_categories.add('label1', attributes={'a1', 'a2'})
label_categories.add('label2') label_categories.add('label2')
return { return { AnnotationType.label: label_categories }
AnnotationType.label: label_categories,
}
with TestDir() as test_dir: parsed_dataset = CvatImporter()(DUMMY_IMAGE_DATASET_DIR).make_dataset()
generate_dummy_cvat(test_dir)
source_dataset = TestExtractor() compare_datasets(self, DstExtractor(), parsed_dataset)
def test_can_load_video(self):
class DstExtractor(Extractor):
def __iter__(self):
return iter([
DatasetItem(id=10, subset='annotations',
image=np.ones((20, 25, 3)),
annotations=[
Bbox(3, 4, 7, 1, label=2,
id=0,
attributes={
'occluded': True,
'outside': False, 'keyframe': True,
'track_id': 0
}),
Points([21.95, 8.00, 2.55, 15.09, 2.23, 3.16],
label=0,
id=1,
attributes={
'occluded': False,
'outside': False, 'keyframe': True,
'track_id': 1, 'hgl': 'hgkf',
}),
]),
DatasetItem(id=13, subset='annotations',
image=np.ones((20, 25, 3)),
annotations=[
Bbox(7, 6, 7, 2, label=2,
id=0,
attributes={
'occluded': False,
'outside': True, 'keyframe': True,
'track_id': 0
}),
Points([21.95, 8.00, 9.55, 15.09, 5.23, 1.16],
label=0,
id=1,
attributes={
'occluded': False,
'outside': True, 'keyframe': True,
'track_id': 1, 'hgl': 'jk',
}),
PolyLine([7.85, 13.88, 3.50, 6.67, 15.90, 2.00, 13.31, 7.21],
label=2,
id=2,
attributes={
'occluded': False,
'outside': False, 'keyframe': True,
'track_id': 2,
}),
]),
DatasetItem(id=16, subset='annotations',
image=Image(path='frame_0000016.png',
size=(20, 25)), # no image in the dataset files
annotations=[
Bbox(8, 7, 6, 10, label=2,
id=0,
attributes={
'occluded': False,
'outside': True, 'keyframe': True,
'track_id': 0
}),
PolyLine([7.85, 13.88, 3.50, 6.67, 15.90, 2.00, 13.31, 7.21],
label=2,
id=2,
attributes={
'occluded': False,
'outside': True, 'keyframe': True,
'track_id': 2,
}),
]),
])
parsed_dataset = CvatImporter()(test_dir).make_dataset() def categories(self):
label_categories = LabelCategories()
label_categories.add('klhg', attributes={'hgl'})
label_categories.add('z U k')
label_categories.add('II')
return { AnnotationType.label: label_categories }
compare_datasets(self, source_dataset, parsed_dataset) parsed_dataset = CvatImporter()(DUMMY_VIDEO_DATASET_DIR).make_dataset()
compare_datasets(self, DstExtractor(), parsed_dataset)
class CvatConverterTest(TestCase): class CvatConverterTest(TestCase):
def _test_save_and_load(self, source_dataset, converter, test_dir, def _test_save_and_load(self, source_dataset, converter, test_dir,
@ -170,12 +173,10 @@ class CvatConverterTest(TestCase):
Polygon([0, 0, 4, 0, 4, 4], Polygon([0, 0, 4, 0, 4, 4],
label=1, group=4, label=1, group=4,
attributes={ 'occluded': True }), attributes={ 'occluded': True }),
Polygon([5, 0, 9, 0, 5, 5],
label=2, group=4,
attributes={ 'unknown': 'bar' }),
Points([1, 1, 3, 2, 2, 3], Points([1, 1, 3, 2, 2, 3],
label=2, label=2,
attributes={ 'a1': 'x', 'a2': 42 }), attributes={ 'a1': 'x', 'a2': 42,
'unknown': 'bar' }),
Label(1), Label(1),
Label(2, attributes={ 'a1': 'y', 'a2': 44 }), Label(2, attributes={ 'a1': 'y', 'a2': 44 }),
] ]
@ -213,9 +214,6 @@ class CvatConverterTest(TestCase):
Polygon([0, 0, 4, 0, 4, 4], Polygon([0, 0, 4, 0, 4, 4],
label=1, group=4, label=1, group=4,
attributes={ 'occluded': True }), attributes={ 'occluded': True }),
Polygon([5, 0, 9, 0, 5, 5],
label=2, group=4,
attributes={ 'occluded': False }),
Points([1, 1, 3, 2, 2, 3], Points([1, 1, 3, 2, 2, 3],
label=2, label=2,
attributes={ 'occluded': False, attributes={ 'occluded': False,

Loading…
Cancel
Save