You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
255 lines
11 KiB
Python
255 lines
11 KiB
Python
import numpy as np
|
|
import os.path as osp
|
|
|
|
from unittest import TestCase
|
|
|
|
from datumaro.components.extractor import (Extractor, DatasetItem,
|
|
AnnotationType, Points, Polygon, PolyLine, Bbox, Label,
|
|
LabelCategories,
|
|
)
|
|
from datumaro.plugins.cvat_format.importer import CvatImporter
|
|
from datumaro.plugins.cvat_format.converter import CvatConverter
|
|
from datumaro.util.image import Image
|
|
from datumaro.util.test_utils import TestDir, compare_datasets
|
|
|
|
|
|
DUMMY_IMAGE_DATASET_DIR = osp.join(osp.dirname(__file__),
|
|
'assets', 'cvat_dataset', 'for_images')
|
|
|
|
DUMMY_VIDEO_DATASET_DIR = osp.join(osp.dirname(__file__),
|
|
'assets', 'cvat_dataset', 'for_video')
|
|
|
|
class CvatImporterTest(TestCase):
|
|
def test_can_detect_image(self):
|
|
self.assertTrue(CvatImporter.detect(DUMMY_IMAGE_DATASET_DIR))
|
|
|
|
def test_can_detect_video(self):
|
|
self.assertTrue(CvatImporter.detect(DUMMY_VIDEO_DATASET_DIR))
|
|
|
|
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)),
|
|
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],
|
|
attributes={'occluded': False}),
|
|
]),
|
|
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}),
|
|
Points([1, 2, 3, 4, 5, 6], label=1, z_order=2,
|
|
attributes={'occluded': False}),
|
|
]),
|
|
])
|
|
|
|
def categories(self):
|
|
label_categories = LabelCategories()
|
|
label_categories.add('label1', attributes={'a1', 'a2'})
|
|
label_categories.add('label2')
|
|
return { AnnotationType.label: label_categories }
|
|
|
|
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,
|
|
}),
|
|
]),
|
|
])
|
|
|
|
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 }
|
|
|
|
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,
|
|
target_dataset=None, importer_args=None):
|
|
converter(source_dataset, test_dir)
|
|
|
|
if importer_args is None:
|
|
importer_args = {}
|
|
parsed_dataset = CvatImporter()(test_dir, **importer_args).make_dataset()
|
|
|
|
if target_dataset is None:
|
|
target_dataset = source_dataset
|
|
|
|
compare_datasets(self, expected=target_dataset, actual=parsed_dataset)
|
|
|
|
def test_can_save_and_load(self):
|
|
label_categories = LabelCategories()
|
|
for i in range(10):
|
|
label_categories.add(str(i))
|
|
label_categories.items[2].attributes.update(['a1', 'a2'])
|
|
label_categories.attributes.update(['occluded'])
|
|
|
|
class SrcExtractor(Extractor):
|
|
def __iter__(self):
|
|
return iter([
|
|
DatasetItem(id=0, subset='s1', image=np.zeros((5, 10, 3)),
|
|
annotations=[
|
|
Polygon([0, 0, 4, 0, 4, 4],
|
|
label=1, group=4,
|
|
attributes={ 'occluded': True }),
|
|
Points([1, 1, 3, 2, 2, 3],
|
|
label=2,
|
|
attributes={ 'a1': 'x', 'a2': 42,
|
|
'unknown': 'bar' }),
|
|
Label(1),
|
|
Label(2, attributes={ 'a1': 'y', 'a2': 44 }),
|
|
]
|
|
),
|
|
DatasetItem(id=1, subset='s1',
|
|
annotations=[
|
|
PolyLine([0, 0, 4, 0, 4, 4],
|
|
label=3, id=4, group=4),
|
|
Bbox(5, 0, 1, 9,
|
|
label=3, id=4, group=4),
|
|
]
|
|
),
|
|
|
|
DatasetItem(id=2, subset='s2', image=np.ones((5, 10, 3)),
|
|
annotations=[
|
|
Polygon([0, 0, 4, 0, 4, 4], z_order=1,
|
|
label=3, group=4,
|
|
attributes={ 'occluded': False }),
|
|
PolyLine([5, 0, 9, 0, 5, 5]), # will be skipped as no label
|
|
]
|
|
),
|
|
|
|
DatasetItem(id=3, subset='s3', image=Image(
|
|
path='3.jpg', size=(2, 4))),
|
|
])
|
|
|
|
def categories(self):
|
|
return { AnnotationType.label: label_categories }
|
|
|
|
class DstExtractor(Extractor):
|
|
def __iter__(self):
|
|
return iter([
|
|
DatasetItem(id=0, subset='s1', image=np.zeros((5, 10, 3)),
|
|
annotations=[
|
|
Polygon([0, 0, 4, 0, 4, 4],
|
|
label=1, group=4,
|
|
attributes={ 'occluded': True }),
|
|
Points([1, 1, 3, 2, 2, 3],
|
|
label=2,
|
|
attributes={ 'occluded': False,
|
|
'a1': 'x', 'a2': 42 }),
|
|
Label(1),
|
|
Label(2, attributes={ 'a1': 'y', 'a2': 44 }),
|
|
]
|
|
),
|
|
DatasetItem(id=1, subset='s1',
|
|
annotations=[
|
|
PolyLine([0, 0, 4, 0, 4, 4],
|
|
label=3, group=4,
|
|
attributes={ 'occluded': False }),
|
|
Bbox(5, 0, 1, 9,
|
|
label=3, group=4,
|
|
attributes={ 'occluded': False }),
|
|
]
|
|
),
|
|
|
|
DatasetItem(id=2, subset='s2', image=np.ones((5, 10, 3)),
|
|
annotations=[
|
|
Polygon([0, 0, 4, 0, 4, 4], z_order=1,
|
|
label=3, group=4,
|
|
attributes={ 'occluded': False }),
|
|
]
|
|
),
|
|
|
|
DatasetItem(id=3, subset='s3', image=Image(
|
|
path='3.jpg', size=(2, 4))),
|
|
])
|
|
|
|
def categories(self):
|
|
return { AnnotationType.label: label_categories }
|
|
|
|
with TestDir() as test_dir:
|
|
self._test_save_and_load(SrcExtractor(),
|
|
CvatConverter(save_images=True), test_dir,
|
|
target_dataset=DstExtractor())
|