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 os
import os.path as osp
from xml.etree import ElementTree as ET
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.converter import CvatConverter
from datumaro.plugins.cvat_format.format import CvatPath
from datumaro.util.image import save_image, Image
from datumaro.util.image import Image
from datumaro.util.test_utils import TestDir, compare_datasets
def generate_dummy_cvat(path):
images_dir = osp.join(path, CvatPath.IMAGES_DIR)
os.makedirs(images_dir)
DUMMY_IMAGE_DATASET_DIR = osp.join(osp.dirname(__file__),
'assets', 'cvat_dataset', 'for_images')
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': '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'))
DUMMY_VIDEO_DATASET_DIR = osp.join(osp.dirname(__file__),
'assets', 'cvat_dataset', 'for_video')
class CvatImporterTest(TestCase):
def test_can_detect(self):
with TestDir() as test_dir:
generate_dummy_cvat(test_dir)
def test_can_detect_image(self):
self.assertTrue(CvatImporter.detect(DUMMY_IMAGE_DATASET_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(self):
class TestExtractor(Extractor):
def test_can_load_image(self):
class DstExtractor(Extractor):
def __iter__(self):
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=[
Bbox(0, 2, 4, 2, label=0, z_order=1,
attributes={
'occluded': True,
'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}),
]),
DatasetItem(id=1, subset='train', image=np.ones((10, 10, 3)),
DatasetItem(id=1, subset='train',
image=np.ones((10, 10, 3)),
annotations=[
Polygon([1, 2, 3, 4, 6, 5], z_order=1,
attributes={'occluded': False}),
@ -128,18 +55,94 @@ class CvatExtractorTest(TestCase):
label_categories = LabelCategories()
label_categories.add('label1', attributes={'a1', 'a2'})
label_categories.add('label2')
return {
AnnotationType.label: label_categories,
}
return { AnnotationType.label: label_categories }
with TestDir() as test_dir:
generate_dummy_cvat(test_dir)
source_dataset = TestExtractor()
parsed_dataset = CvatImporter()(DUMMY_IMAGE_DATASET_DIR).make_dataset()
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):
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],
label=1, group=4,
attributes={ 'occluded': True }),
Polygon([5, 0, 9, 0, 5, 5],
label=2, group=4,
attributes={ 'unknown': 'bar' }),
Points([1, 1, 3, 2, 2, 3],
label=2,
attributes={ 'a1': 'x', 'a2': 42 }),
attributes={ 'a1': 'x', 'a2': 42,
'unknown': 'bar' }),
Label(1),
Label(2, attributes={ 'a1': 'y', 'a2': 44 }),
]
@ -213,9 +214,6 @@ class CvatConverterTest(TestCase):
Polygon([0, 0, 4, 0, 4, 4],
label=1, group=4,
attributes={ 'occluded': True }),
Polygon([5, 0, 9, 0, 5, 5],
label=2, group=4,
attributes={ 'occluded': False }),
Points([1, 1, 3, 2, 2, 3],
label=2,
attributes={ 'occluded': False,

Loading…
Cancel
Save