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.
647 lines
25 KiB
Python
647 lines
25 KiB
Python
import json
|
|
import numpy as np
|
|
import os
|
|
import os.path as osp
|
|
|
|
from unittest import TestCase
|
|
|
|
from datumaro.components.project import Project
|
|
from datumaro.components.extractor import (Extractor, DatasetItem,
|
|
AnnotationType, Label, Mask, Points, Polygon, Bbox, Caption,
|
|
LabelCategories, PointsCategories
|
|
)
|
|
from datumaro.plugins.coco_format.converter import (
|
|
CocoConverter,
|
|
CocoImageInfoConverter,
|
|
CocoCaptionsConverter,
|
|
CocoInstancesConverter,
|
|
CocoPersonKeypointsConverter,
|
|
CocoLabelsConverter,
|
|
)
|
|
from datumaro.plugins.coco_format.importer import CocoImporter
|
|
from datumaro.util.image import save_image, Image
|
|
from datumaro.util.test_utils import TestDir, compare_datasets
|
|
|
|
|
|
class CocoImporterTest(TestCase):
|
|
@staticmethod
|
|
def generate_annotation():
|
|
annotation = {
|
|
'licenses': [],
|
|
'info': {},
|
|
'categories': [],
|
|
'images': [],
|
|
'annotations': [],
|
|
}
|
|
annotation['licenses'].append({
|
|
'name': '',
|
|
'id': 0,
|
|
'url': '',
|
|
})
|
|
annotation['info'] = {
|
|
'contributor': '',
|
|
'date_created': '',
|
|
'description': '',
|
|
'url': '',
|
|
'version': '',
|
|
'year': '',
|
|
}
|
|
annotation['licenses'].append({
|
|
'name': '',
|
|
'id': 0,
|
|
'url': '',
|
|
})
|
|
annotation['categories'].append({
|
|
'id': 1,
|
|
'name': 'TEST',
|
|
'supercategory': '',
|
|
})
|
|
annotation['images'].append({
|
|
"id": 1,
|
|
"width": 5,
|
|
"height": 10,
|
|
"file_name": '000000000001.jpg',
|
|
"license": 0,
|
|
"flickr_url": '',
|
|
"coco_url": '',
|
|
"date_captured": 0,
|
|
})
|
|
annotation['annotations'].append({
|
|
"id": 1,
|
|
"image_id": 1,
|
|
"category_id": 1,
|
|
"segmentation": [[0, 0, 1, 0, 1, 2, 0, 2]],
|
|
"area": 2,
|
|
"bbox": [0, 0, 1, 2],
|
|
"iscrowd": 0,
|
|
})
|
|
annotation['annotations'].append({
|
|
"id": 2,
|
|
"image_id": 1,
|
|
"category_id": 1,
|
|
"segmentation": {
|
|
"counts": [
|
|
0, 10,
|
|
5, 5,
|
|
5, 5,
|
|
0, 10,
|
|
10, 0],
|
|
"size": [10, 5]},
|
|
"area": 30,
|
|
"bbox": [0, 0, 10, 4],
|
|
"iscrowd": 1,
|
|
})
|
|
return annotation
|
|
|
|
def COCO_dataset_generate(self, path):
|
|
img_dir = osp.join(path, 'images', 'val')
|
|
ann_dir = osp.join(path, 'annotations')
|
|
os.makedirs(img_dir)
|
|
os.makedirs(ann_dir)
|
|
|
|
image = np.ones((10, 5, 3))
|
|
save_image(osp.join(img_dir, '000000000001.jpg'), image)
|
|
|
|
annotation = self.generate_annotation()
|
|
|
|
with open(osp.join(ann_dir, 'instances_val.json'), 'w') as outfile:
|
|
json.dump(annotation, outfile)
|
|
|
|
def test_can_import(self):
|
|
class DstExtractor(Extractor):
|
|
def __iter__(self):
|
|
return iter([
|
|
DatasetItem(id=1, image=np.ones((10, 5, 3)), subset='val',
|
|
annotations=[
|
|
Polygon([0, 0, 1, 0, 1, 2, 0, 2], label=0,
|
|
id=1, group=1, attributes={'is_crowd': False}),
|
|
Mask(np.array(
|
|
[[1, 0, 0, 1, 0]] * 5 +
|
|
[[1, 1, 1, 1, 0]] * 5
|
|
), label=0,
|
|
id=2, group=2, attributes={'is_crowd': True}),
|
|
]
|
|
),
|
|
])
|
|
|
|
def categories(self):
|
|
label_cat = LabelCategories()
|
|
label_cat.add('TEST')
|
|
return { AnnotationType.label: label_cat }
|
|
|
|
with TestDir() as test_dir:
|
|
self.COCO_dataset_generate(test_dir)
|
|
|
|
dataset = Project.import_from(test_dir, 'coco').make_dataset()
|
|
|
|
compare_datasets(self, DstExtractor(), dataset)
|
|
|
|
def test_can_detect(self):
|
|
with TestDir() as test_dir:
|
|
self.COCO_dataset_generate(test_dir)
|
|
|
|
self.assertTrue(CocoImporter.detect(test_dir))
|
|
|
|
class CocoConverterTest(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 = CocoImporter()(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_captions(self):
|
|
class TestExtractor(Extractor):
|
|
def __iter__(self):
|
|
return iter([
|
|
DatasetItem(id=1, subset='train',
|
|
annotations=[
|
|
Caption('hello', id=1, group=1),
|
|
Caption('world', id=2, group=2),
|
|
]),
|
|
DatasetItem(id=2, subset='train',
|
|
annotations=[
|
|
Caption('test', id=3, group=3),
|
|
]),
|
|
|
|
DatasetItem(id=3, subset='val',
|
|
annotations=[
|
|
Caption('word', id=1, group=1),
|
|
]
|
|
),
|
|
])
|
|
|
|
with TestDir() as test_dir:
|
|
self._test_save_and_load(TestExtractor(),
|
|
CocoCaptionsConverter(), test_dir)
|
|
|
|
def test_can_save_and_load_instances(self):
|
|
label_categories = LabelCategories()
|
|
for i in range(10):
|
|
label_categories.add(str(i))
|
|
categories = { AnnotationType.label: label_categories }
|
|
|
|
class TestExtractor(Extractor):
|
|
def __iter__(self):
|
|
return iter([
|
|
DatasetItem(id=1, subset='train', image=np.ones((4, 4, 3)),
|
|
annotations=[
|
|
# Bbox + single polygon
|
|
Bbox(0, 1, 2, 2,
|
|
label=2, group=1, id=1,
|
|
attributes={ 'is_crowd': False }),
|
|
Polygon([0, 1, 2, 1, 2, 3, 0, 3],
|
|
attributes={ 'is_crowd': False },
|
|
label=2, group=1, id=1),
|
|
]),
|
|
DatasetItem(id=2, subset='train', image=np.ones((4, 4, 3)),
|
|
annotations=[
|
|
# Mask + bbox
|
|
Mask(np.array([
|
|
[0, 1, 0, 0],
|
|
[0, 1, 0, 0],
|
|
[0, 1, 1, 1],
|
|
[0, 0, 0, 0]],
|
|
),
|
|
attributes={ 'is_crowd': True },
|
|
label=4, group=3, id=3),
|
|
Bbox(1, 0, 2, 2, label=4, group=3, id=3,
|
|
attributes={ 'is_crowd': True }),
|
|
]),
|
|
|
|
DatasetItem(id=3, subset='val', image=np.ones((4, 4, 3)),
|
|
annotations=[
|
|
# Bbox + mask
|
|
Bbox(0, 1, 2, 2, label=4, group=3, id=3,
|
|
attributes={ 'is_crowd': True }),
|
|
Mask(np.array([
|
|
[0, 0, 0, 0],
|
|
[1, 1, 1, 0],
|
|
[1, 1, 0, 0],
|
|
[0, 0, 0, 0]],
|
|
),
|
|
attributes={ 'is_crowd': True },
|
|
label=4, group=3, id=3),
|
|
]),
|
|
])
|
|
|
|
def categories(self):
|
|
return categories
|
|
|
|
class DstExtractor(Extractor):
|
|
def __iter__(self):
|
|
return iter([
|
|
DatasetItem(id=1, subset='train', image=np.ones((4, 4, 3)),
|
|
annotations=[
|
|
Polygon([0, 1, 2, 1, 2, 3, 0, 3],
|
|
attributes={ 'is_crowd': False },
|
|
label=2, group=1, id=1),
|
|
]),
|
|
DatasetItem(id=2, subset='train', image=np.ones((4, 4, 3)),
|
|
annotations=[
|
|
Mask(np.array([
|
|
[0, 1, 0, 0],
|
|
[0, 1, 0, 0],
|
|
[0, 1, 1, 1],
|
|
[0, 0, 0, 0]],
|
|
),
|
|
attributes={ 'is_crowd': True },
|
|
label=4, group=3, id=3),
|
|
]),
|
|
|
|
DatasetItem(id=3, subset='val', image=np.ones((4, 4, 3)),
|
|
annotations=[
|
|
Mask(np.array([
|
|
[0, 0, 0, 0],
|
|
[1, 1, 1, 0],
|
|
[1, 1, 0, 0],
|
|
[0, 0, 0, 0]],
|
|
),
|
|
attributes={ 'is_crowd': True },
|
|
label=4, group=3, id=3),
|
|
]),
|
|
])
|
|
|
|
def categories(self):
|
|
return categories
|
|
|
|
with TestDir() as test_dir:
|
|
self._test_save_and_load(TestExtractor(),
|
|
CocoInstancesConverter(), test_dir,
|
|
target_dataset=DstExtractor())
|
|
|
|
def test_can_merge_polygons_on_loading(self):
|
|
label_categories = LabelCategories()
|
|
for i in range(10):
|
|
label_categories.add(str(i))
|
|
categories = { AnnotationType.label: label_categories }
|
|
|
|
class TestExtractor(Extractor):
|
|
def __iter__(self):
|
|
return iter([
|
|
DatasetItem(id=1, image=np.zeros((6, 10, 3)),
|
|
annotations=[
|
|
Polygon([0, 0, 4, 0, 4, 4],
|
|
label=3, id=4, group=4),
|
|
Polygon([5, 0, 9, 0, 5, 5],
|
|
label=3, id=4, group=4),
|
|
]
|
|
),
|
|
])
|
|
|
|
def categories(self):
|
|
return categories
|
|
|
|
class TargetExtractor(TestExtractor):
|
|
def __iter__(self):
|
|
items = list(super().__iter__())
|
|
items[0]._annotations = [
|
|
Mask(np.array([
|
|
[0, 1, 1, 1, 0, 1, 1, 1, 1, 0],
|
|
[0, 0, 1, 1, 0, 1, 1, 1, 0, 0],
|
|
[0, 0, 0, 1, 0, 1, 1, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]],
|
|
# only internal fragment (without the border),
|
|
# but not everywhere...
|
|
),
|
|
label=3, id=4, group=4,
|
|
attributes={ 'is_crowd': False }),
|
|
]
|
|
return iter(items)
|
|
|
|
with TestDir() as test_dir:
|
|
self._test_save_and_load(TestExtractor(),
|
|
CocoInstancesConverter(), test_dir,
|
|
importer_args={'merge_instance_polygons': True},
|
|
target_dataset=TargetExtractor())
|
|
|
|
def test_can_crop_covered_segments(self):
|
|
label_categories = LabelCategories()
|
|
for i in range(10):
|
|
label_categories.add(str(i))
|
|
|
|
class SrcTestExtractor(Extractor):
|
|
def __iter__(self):
|
|
return iter([
|
|
DatasetItem(id=1, image=np.zeros((5, 5, 3)),
|
|
annotations=[
|
|
Mask(np.array([
|
|
[0, 0, 1, 1, 1],
|
|
[0, 0, 1, 1, 1],
|
|
[1, 1, 0, 1, 1],
|
|
[1, 1, 1, 0, 0],
|
|
[1, 1, 1, 0, 0]],
|
|
),
|
|
label=2, id=1, z_order=0),
|
|
Polygon([1, 1, 4, 1, 4, 4, 1, 4],
|
|
label=1, id=2, z_order=1),
|
|
]
|
|
),
|
|
])
|
|
|
|
def categories(self):
|
|
return { AnnotationType.label: label_categories }
|
|
|
|
class DstTestExtractor(Extractor):
|
|
def __iter__(self):
|
|
return iter([
|
|
DatasetItem(id=1, image=np.zeros((5, 5, 3)),
|
|
annotations=[
|
|
Mask(np.array([
|
|
[0, 0, 1, 1, 1],
|
|
[0, 0, 0, 0, 1],
|
|
[1, 0, 0, 0, 1],
|
|
[1, 0, 0, 0, 0],
|
|
[1, 1, 1, 0, 0]],
|
|
),
|
|
attributes={ 'is_crowd': True },
|
|
label=2, id=1, group=1),
|
|
|
|
Polygon([1, 1, 4, 1, 4, 4, 1, 4],
|
|
label=1, id=2, group=2,
|
|
attributes={ 'is_crowd': False }),
|
|
]
|
|
),
|
|
])
|
|
|
|
def categories(self):
|
|
return { AnnotationType.label: label_categories }
|
|
|
|
with TestDir() as test_dir:
|
|
self._test_save_and_load(SrcTestExtractor(),
|
|
CocoInstancesConverter(crop_covered=True), test_dir,
|
|
target_dataset=DstTestExtractor())
|
|
|
|
def test_can_convert_polygons_to_mask(self):
|
|
label_categories = LabelCategories()
|
|
for i in range(10):
|
|
label_categories.add(str(i))
|
|
|
|
class SrcTestExtractor(Extractor):
|
|
def __iter__(self):
|
|
return iter([
|
|
DatasetItem(id=1, image=np.zeros((6, 10, 3)),
|
|
annotations=[
|
|
Polygon([0, 0, 4, 0, 4, 4],
|
|
label=3, id=4, group=4),
|
|
Polygon([5, 0, 9, 0, 5, 5],
|
|
label=3, id=4, group=4),
|
|
]
|
|
),
|
|
])
|
|
|
|
def categories(self):
|
|
return { AnnotationType.label: label_categories }
|
|
|
|
class DstTestExtractor(Extractor):
|
|
def __iter__(self):
|
|
return iter([
|
|
DatasetItem(id=1, image=np.zeros((6, 10, 3)),
|
|
annotations=[
|
|
Mask(np.array([
|
|
[0, 1, 1, 1, 0, 1, 1, 1, 1, 0],
|
|
[0, 0, 1, 1, 0, 1, 1, 1, 0, 0],
|
|
[0, 0, 0, 1, 0, 1, 1, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]],
|
|
# only internal fragment (without the border),
|
|
# but not everywhere...
|
|
),
|
|
attributes={ 'is_crowd': True },
|
|
label=3, id=4, group=4),
|
|
]
|
|
),
|
|
])
|
|
|
|
def categories(self):
|
|
return { AnnotationType.label: label_categories }
|
|
|
|
with TestDir() as test_dir:
|
|
self._test_save_and_load(SrcTestExtractor(),
|
|
CocoInstancesConverter(segmentation_mode='mask'), test_dir,
|
|
target_dataset=DstTestExtractor())
|
|
|
|
def test_can_convert_masks_to_polygons(self):
|
|
label_categories = LabelCategories()
|
|
for i in range(10):
|
|
label_categories.add(str(i))
|
|
|
|
class SrcTestExtractor(Extractor):
|
|
def __iter__(self):
|
|
return iter([
|
|
DatasetItem(id=1, image=np.zeros((5, 10, 3)),
|
|
annotations=[
|
|
Mask(np.array([
|
|
[0, 1, 1, 1, 0, 1, 1, 1, 1, 0],
|
|
[0, 0, 1, 1, 0, 1, 1, 1, 0, 0],
|
|
[0, 0, 0, 1, 0, 1, 1, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 1, 0, 0, 0, 0],
|
|
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
|
|
]),
|
|
label=3, id=4, group=4),
|
|
]
|
|
),
|
|
])
|
|
|
|
def categories(self):
|
|
return { AnnotationType.label: label_categories }
|
|
|
|
class DstTestExtractor(Extractor):
|
|
def __iter__(self):
|
|
return iter([
|
|
DatasetItem(id=1, image=np.zeros((5, 10, 3)),
|
|
annotations=[
|
|
Polygon(
|
|
[3.0, 2.5, 1.0, 0.0, 3.5, 0.0, 3.0, 2.5],
|
|
label=3, id=4, group=4,
|
|
attributes={ 'is_crowd': False }),
|
|
Polygon(
|
|
[5.0, 3.5, 4.5, 0.0, 8.0, 0.0, 5.0, 3.5],
|
|
label=3, id=4, group=4,
|
|
attributes={ 'is_crowd': False }),
|
|
]
|
|
),
|
|
])
|
|
|
|
def categories(self):
|
|
return { AnnotationType.label: label_categories }
|
|
|
|
with TestDir() as test_dir:
|
|
self._test_save_and_load(SrcTestExtractor(),
|
|
CocoInstancesConverter(segmentation_mode='polygons'), test_dir,
|
|
target_dataset=DstTestExtractor())
|
|
|
|
def test_can_save_and_load_images(self):
|
|
class TestExtractor(Extractor):
|
|
def __iter__(self):
|
|
return iter([
|
|
DatasetItem(id=1, subset='train'),
|
|
DatasetItem(id=2, subset='train'),
|
|
|
|
DatasetItem(id=2, subset='val'),
|
|
DatasetItem(id=3, subset='val'),
|
|
DatasetItem(id=4, subset='val'),
|
|
|
|
DatasetItem(id=5, subset='test'),
|
|
])
|
|
|
|
with TestDir() as test_dir:
|
|
self._test_save_and_load(TestExtractor(),
|
|
CocoImageInfoConverter(), test_dir)
|
|
|
|
def test_can_save_and_load_labels(self):
|
|
class TestExtractor(Extractor):
|
|
def __iter__(self):
|
|
return iter([
|
|
DatasetItem(id=1, subset='train',
|
|
annotations=[
|
|
Label(4, id=1, group=1),
|
|
Label(9, id=2, group=2),
|
|
]),
|
|
DatasetItem(id=2, subset='train',
|
|
annotations=[
|
|
Label(4, id=4, group=4),
|
|
]),
|
|
|
|
DatasetItem(id=3, subset='val',
|
|
annotations=[
|
|
Label(2, id=1, group=1),
|
|
]),
|
|
])
|
|
|
|
def categories(self):
|
|
label_categories = LabelCategories()
|
|
for i in range(10):
|
|
label_categories.add(str(i))
|
|
return {
|
|
AnnotationType.label: label_categories,
|
|
}
|
|
|
|
with TestDir() as test_dir:
|
|
self._test_save_and_load(TestExtractor(),
|
|
CocoLabelsConverter(), test_dir)
|
|
|
|
def test_can_save_and_load_keypoints(self):
|
|
label_categories = LabelCategories()
|
|
points_categories = PointsCategories()
|
|
for i in range(10):
|
|
label_categories.add(str(i))
|
|
points_categories.add(i, joints=[[0, 1], [1, 2]])
|
|
categories = {
|
|
AnnotationType.label: label_categories,
|
|
AnnotationType.points: points_categories,
|
|
}
|
|
|
|
class TestExtractor(Extractor):
|
|
def __iter__(self):
|
|
return iter([
|
|
DatasetItem(id=1, subset='train', image=np.zeros((5, 5, 3)),
|
|
annotations=[
|
|
# Full instance annotations: polygon + keypoints
|
|
Points([0, 0, 0, 2, 4, 1], [0, 1, 2],
|
|
label=3, group=1, id=1),
|
|
Polygon([0, 0, 4, 0, 4, 4],
|
|
label=3, group=1, id=1),
|
|
|
|
# Full instance annotations: bbox + keypoints
|
|
Points([1, 2, 3, 4, 2, 3], group=2, id=2),
|
|
Bbox(1, 2, 2, 2, group=2, id=2),
|
|
]),
|
|
DatasetItem(id=2, subset='train', image=np.zeros((5, 4, 3)),
|
|
annotations=[
|
|
# Solitary keypoints
|
|
Points([1, 2, 0, 2, 4, 1], label=5, id=3),
|
|
|
|
# Some other solitary annotations (bug #1387)
|
|
Polygon([0, 0, 4, 0, 4, 4], label=3, id=4),
|
|
]),
|
|
|
|
DatasetItem(id=3, subset='val',
|
|
annotations=[
|
|
# Solitary keypoints with no label
|
|
Points([0, 0, 1, 2, 3, 4], [0, 1, 2], id=3),
|
|
]),
|
|
])
|
|
|
|
def categories(self):
|
|
return categories
|
|
|
|
class DstTestExtractor(TestExtractor):
|
|
def __iter__(self):
|
|
return iter([
|
|
DatasetItem(id=1, subset='train', image=np.zeros((5, 5, 3)),
|
|
annotations=[
|
|
Points([0, 0, 0, 2, 4, 1], [0, 1, 2],
|
|
label=3, group=1, id=1,
|
|
attributes={'is_crowd': False}),
|
|
Polygon([0, 0, 4, 0, 4, 4],
|
|
label=3, group=1, id=1,
|
|
attributes={'is_crowd': False}),
|
|
|
|
Points([1, 2, 3, 4, 2, 3],
|
|
group=2, id=2,
|
|
attributes={'is_crowd': False}),
|
|
Polygon([1, 2, 3, 2, 3, 4, 1, 4],
|
|
group=2, id=2,
|
|
attributes={'is_crowd': False}),
|
|
]),
|
|
DatasetItem(id=2, subset='train',
|
|
annotations=[
|
|
Points([1, 2, 0, 2, 4, 1],
|
|
label=5, group=3, id=3,
|
|
attributes={'is_crowd': False}),
|
|
Polygon([0, 1, 4, 1, 4, 2, 0, 2],
|
|
label=5, group=3, id=3,
|
|
attributes={'is_crowd': False}),
|
|
]),
|
|
|
|
DatasetItem(id=3, subset='val',
|
|
annotations=[
|
|
Points([0, 0, 1, 2, 3, 4], [0, 1, 2],
|
|
group=3, id=3,
|
|
attributes={'is_crowd': False}),
|
|
Polygon([1, 2, 3, 2, 3, 4, 1, 4],
|
|
group=3, id=3,
|
|
attributes={'is_crowd': False}),
|
|
]),
|
|
])
|
|
|
|
with TestDir() as test_dir:
|
|
self._test_save_and_load(TestExtractor(),
|
|
CocoPersonKeypointsConverter(), test_dir,
|
|
target_dataset=DstTestExtractor())
|
|
|
|
def test_can_save_dataset_with_no_subsets(self):
|
|
class TestExtractor(Extractor):
|
|
def __iter__(self):
|
|
return iter([
|
|
DatasetItem(id=1),
|
|
DatasetItem(id=2),
|
|
])
|
|
|
|
def categories(self):
|
|
return { AnnotationType.label: LabelCategories() }
|
|
|
|
with TestDir() as test_dir:
|
|
self._test_save_and_load(TestExtractor(),
|
|
CocoConverter(), test_dir)
|
|
|
|
def test_can_save_dataset_with_image_info(self):
|
|
class TestExtractor(Extractor):
|
|
def __iter__(self):
|
|
return iter([
|
|
DatasetItem(id=1, image=Image(path='1.jpg', size=(10, 15))),
|
|
])
|
|
|
|
with TestDir() as test_dir:
|
|
self._test_save_and_load(TestExtractor(),
|
|
CocoConverter(tasks='image_info'), test_dir) |