def test_annotation_comparison(self):
a = Dataset.from_iterable([
DatasetItem(id=1, annotations=[
Caption('hello'), # unmatched
Caption('world', group=5),
Label(2, attributes={ 'x': 1, 'y': '2', }),
Bbox(1, 2, 3, 4, label=4, z_order=1, attributes={
'score': 1.0,
}),
Bbox(5, 6, 7, 8, group=5),
Points([1, 2, 2, 0, 1, 1], label=0, z_order=4),
Mask(label=3, z_order=2, image=np.ones((2, 3))),
]),
], categories=['a', 'b', 'c', 'd'])
b = Dataset.from_iterable([
DatasetItem(id=1, annotations=[
Caption('world', group=5),
Label(2, attributes={ 'x': 1, 'y': '2', }),
Bbox(1, 2, 3, 4, label=4, z_order=1, attributes={
'score': 1.0,
}),
Bbox(5, 6, 7, 8, group=5),
Bbox(5, 6, 7, 8, group=5), # unmatched
Points([1, 2, 2, 0, 1, 1], label=0, z_order=4),
Mask(label=3, z_order=2, image=np.ones((2, 3))),
]),
], categories=['a', 'b', 'c', 'd'])
comp = ExactComparator()
matched, unmatched, _, _, errors = comp.compare_datasets(a, b)
self.assertEqual(6, len(matched), matched)
self.assertEqual(2, len(unmatched), unmatched)
self.assertEqual(0, len(errors), errors)
def test_can_save_and_load_captions(self):
expected_dataset = Dataset.from_iterable([
DatasetItem(id=1,
subset='train',
annotations=[
Caption('hello', id=1, group=1),
Caption('world', id=2, group=2),
],
attributes={'id': 1}),
DatasetItem(id=2,
subset='train',
annotations=[
Caption('test', id=3, group=3),
],
attributes={'id': 2}),
DatasetItem(id=3,
subset='val',
annotations=[
Caption('word', id=1, group=1),
],
attributes={'id': 1}),
])
with TestDir() as test_dir:
self._test_save_and_load(expected_dataset,
CocoCaptionsConverter.convert, test_dir)
def test_can_import_captions(self):
expected_dataset = Dataset.from_iterable([
DatasetItem(id=1,
subset='train',
annotations=[
Caption('hello', id=1, group=1),
Caption('world', id=2, group=2),
],
attributes={'id': 1}),
DatasetItem(id=2,
subset='train',
annotations=[
Caption('test', id=3, group=3),
],
attributes={'id': 2}),
DatasetItem(id=3,
subset='val',
annotations=[
Caption('word', id=1, group=1),
],
attributes={'id': 1}),
])
dataset = Dataset.import_from(
osp.join(DUMMY_DATASET_DIR, 'coco_captions'), 'coco')
compare_datasets(self, expected_dataset, dataset)
def __iter__(self):
return iter([
DatasetItem(id=100,
subset='train',
image=np.ones((10, 6, 3)),
annotations=[
Caption('hello', id=1),
Caption('world', id=2, group=5),
Label(2, id=3, attributes={
'x': 1,
'y': '2',
}),
Bbox(1,
2,
3,
4,
label=4,
id=4,
z_order=1,
attributes={
'score': 1.0,
}),
Bbox(5, 6, 7, 8, id=5, group=5),
Points([1, 2, 2, 0, 1, 1],
label=0,
id=5,
z_order=4),
Mask(label=3,
id=5,
z_order=2,
image=np.ones((2, 3))),
]),
DatasetItem(id=21,
subset='train',
annotations=[
Caption('test'),
Label(2),
Bbox(1, 2, 3, 4, 5, id=42, group=42)
]),
DatasetItem(id=2,
subset='val',
annotations=[
PolyLine([1, 2, 3, 4, 5, 6, 7, 8],
id=11,
z_order=1),
Polygon([1, 2, 3, 4, 5, 6, 7, 8],
id=12,
z_order=4),
]),
DatasetItem(id=42,
subset='test',
attributes={
'a1': 5,
'a2': '42'
}),
DatasetItem(id=42),
DatasetItem(id=43, image=Image(path='1/b/c.qq', size=(2, 4))),
])
def test_dataset(self):
label_categories = LabelCategories()
for i in range(5):
label_categories.add('cat' + str(i))
mask_categories = MaskCategories(
generate_colormap(len(label_categories.items)))
points_categories = PointsCategories()
for index, _ in enumerate(label_categories.items):
points_categories.add(index, ['cat1', 'cat2'], joints=[[0, 1]])
return Dataset.from_iterable([
DatasetItem(id=100, subset='train', image=np.ones((10, 6, 3)),
annotations=[
Caption('hello', id=1),
Caption('world', id=2, group=5),
Label(2, id=3, attributes={
'x': 1,
'y': '2',
}),
Bbox(1, 2, 3, 4, label=4, id=4, z_order=1, attributes={
'score': 1.0,
}),
Bbox(5, 6, 7, 8, id=5, group=5),
Points([1, 2, 2, 0, 1, 1], label=0, id=5, z_order=4),
Mask(label=3, id=5, z_order=2, image=np.ones((2, 3))),
]),
DatasetItem(id=21, subset='train',
annotations=[
Caption('test'),
Label(2),
Bbox(1, 2, 3, 4, label=5, id=42, group=42)
]),
DatasetItem(id=2, subset='val',
annotations=[
PolyLine([1, 2, 3, 4, 5, 6, 7, 8], id=11, z_order=1),
Polygon([1, 2, 3, 4, 5, 6, 7, 8], id=12, z_order=4),
]),
DatasetItem(id=42, subset='test',
attributes={'a1': 5, 'a2': '42'}),
DatasetItem(id=42),
DatasetItem(id=43, image=Image(path='1/b/c.qq', size=(2, 4))),
], categories={
AnnotationType.label: label_categories,
AnnotationType.mask: mask_categories,
AnnotationType.points: points_categories,
})
def test_can_save_and_load_captions(self):
expected_dataset = Dataset.from_iterable([
DatasetItem(id=1, subset='train',
annotations=[
Caption('caption_0'),
]),
DatasetItem(id=2, subset='train',
annotations=[
Caption('caption_1'),
]),
])
with TestDir() as test_dir:
self._test_save_and_load(expected_dataset,
IcdarWordRecognitionConverter.convert, test_dir)
def transform_item(self, item):
annotations = item.annotations
anns = [p for p in annotations if 'text' in p.attributes]
for ann in anns:
annotations.append(Caption(ann.attributes['text']))
annotations.remove(ann)
return item.wrap(annotations=annotations)
def _load_recognition_items(self):
items = {}
with open(self._path, encoding='utf-8') as f:
for line in f:
line = line.strip()
objects = line.split(', ')
if len(objects) == 2:
image = objects[0]
captions = objects[1].split()
else:
image = objects[0][:-1]
captions = []
item_id = image[:-len(IcdarPath.IMAGE_EXT)]
image_path = osp.join(osp.dirname(self._path),
IcdarPath.IMAGES_DIR, image)
if item_id not in items:
items[item_id] = DatasetItem(item_id,
subset=self._subset,
image=image_path)
annotations = items[item_id].annotations
for caption in captions:
if caption[0] == '\"' and caption[-1] == '\"':
caption = caption[1:-1]
annotations.append(Caption(caption))
return items
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),
]
),
])
def _load_annotations(item):
parsed = item['annotations']
loaded = []
for ann in parsed:
ann_id = ann.get('id')
ann_type = AnnotationType[ann['type']]
attributes = ann.get('attributes')
group = ann.get('group')
label_id = ann.get('label_id')
z_order = ann.get('z_order')
points = ann.get('points')
if ann_type == AnnotationType.label:
loaded.append(Label(label=label_id,
id=ann_id, attributes=attributes, group=group))
elif ann_type == AnnotationType.mask:
rle = ann['rle']
rle['counts'] = rle['counts'].encode('ascii')
loaded.append(RleMask(rle=rle, label=label_id,
id=ann_id, attributes=attributes, group=group,
z_order=z_order))
elif ann_type == AnnotationType.polyline:
loaded.append(PolyLine(points, label=label_id,
id=ann_id, attributes=attributes, group=group,
z_order=z_order))
elif ann_type == AnnotationType.polygon:
loaded.append(Polygon(points, label=label_id,
id=ann_id, attributes=attributes, group=group,
z_order=z_order))
elif ann_type == AnnotationType.bbox:
x, y, w, h = ann['bbox']
loaded.append(Bbox(x, y, w, h, label=label_id,
id=ann_id, attributes=attributes, group=group,
z_order=z_order))
elif ann_type == AnnotationType.points:
loaded.append(Points(points, label=label_id,
id=ann_id, attributes=attributes, group=group,
z_order=z_order))
elif ann_type == AnnotationType.caption:
caption = ann.get('caption')
loaded.append(Caption(caption,
id=ann_id, attributes=attributes, group=group))
else:
raise NotImplementedError()
return loaded
def test_item_representations():
item = DatasetItem(id=1, subset='subset', path=['a', 'b'],
image=np.ones((5, 4, 3)),
annotations=[
Label(0, attributes={'a1': 1, 'a2': '2'}, id=1, group=2),
Caption('hello', id=1),
Caption('world', group=5),
Label(2, id=3, attributes={ 'x': 1, 'y': '2' }),
Bbox(1, 2, 3, 4, label=4, id=4, attributes={ 'a': 1.0 }),
Bbox(5, 6, 7, 8, id=5, group=5),
Points([1, 2, 2, 0, 1, 1], label=0, id=5),
Mask(id=5, image=np.ones((3, 2))),
Mask(label=3, id=5, image=np.ones((2, 3))),
PolyLine([1, 2, 3, 4, 5, 6, 7, 8], id=11),
Polygon([1, 2, 3, 4, 5, 6, 7, 8]),
]
)
encoded = DatasetItemEncoder.encode(item)
DatasetItemEncoder.to_string(encoded)
def test_can_import_captions(self):
expected_dataset = Dataset.from_iterable([
DatasetItem(id='word_1', subset='train',
image=np.ones((10, 15, 3)),
annotations=[
Caption('PROPER'),
]
),
DatasetItem(id='word_2', subset='train',
image=np.ones((10, 15, 3)),
annotations=[
Caption("Canon"),
]
),
])
dataset = Dataset.import_from(
osp.join(DUMMY_DATASET_DIR, 'word_recognition'), 'icdar')
compare_datasets(self, expected_dataset, dataset)
def test_can_save_and_load_captions(self):
expected_dataset = Dataset.from_iterable([
DatasetItem(id='a/b/1',
subset='train',
image=np.ones((10, 15, 3)),
annotations=[
Caption('caption 0'),
]),
DatasetItem(id=2,
subset='train',
image=np.ones((10, 15, 3)),
annotations=[
Caption('caption_1'),
]),
])
with TestDir() as test_dir:
self._test_save_and_load(
expected_dataset,
partial(IcdarWordRecognitionConverter.convert,
save_images=True), test_dir, 'icdar_word_recognition')
def _load_recognition_items(self):
items = {}
with open(self._path, encoding='utf-8') as f:
for line in f:
line = line.strip()
objects = line.split(', ')
if len(objects) == 2:
image = objects[0]
objects = objects[1].split('\"')
if 1 < len(objects):
if len(objects) % 2:
captions = [
objects[2 * i + 1]
for i in range(int(len(objects) / 2))
]
else:
raise Exception("Line %s: unexpected number "
"of quotes in filename" % line)
else:
captions = objects[0].split()
else:
image = objects[0][:-1]
captions = []
item_id = osp.splitext(image)[0]
image_path = osp.join(osp.dirname(self._path),
IcdarPath.IMAGES_DIR, image)
if item_id not in items:
items[item_id] = DatasetItem(item_id,
subset=self._subset,
image=image_path)
annotations = items[item_id].annotations
for caption in captions:
annotations.append(Caption(caption))
return items
def _load_annotations(self, ann, image_info=None):
parsed_annotations = []
ann_id = ann.get('id')
attributes = {}
if 'score' in ann:
attributes['score'] = ann['score']
group = ann_id # make sure all tasks' annotations are merged
if self._task in [CocoTask.instances, CocoTask.person_keypoints]:
x, y, w, h = ann['bbox']
label_id = self._get_label_id(ann)
is_crowd = bool(ann['iscrowd'])
attributes['is_crowd'] = is_crowd
if self._task is CocoTask.person_keypoints:
keypoints = ann['keypoints']
points = [p for i, p in enumerate(keypoints) if i % 3 != 2]
visibility = keypoints[2::3]
parsed_annotations.append(
Points(points,
visibility,
label=label_id,
id=ann_id,
attributes=attributes,
group=group))
segmentation = ann.get('segmentation')
if segmentation and segmentation != [[]]:
rle = None
if isinstance(segmentation, list):
if not self._merge_instance_polygons:
# polygon - a single object can consist of multiple parts
for polygon_points in segmentation:
parsed_annotations.append(
Polygon(points=polygon_points,
label=label_id,
id=ann_id,
attributes=attributes,
group=group))
else:
# merge all parts into a single mask RLE
img_h = image_info['height']
img_w = image_info['width']
rles = mask_utils.frPyObjects(segmentation, img_h,
img_w)
rle = mask_utils.merge(rles)
elif isinstance(segmentation['counts'], list):
# uncompressed RLE
img_h = image_info['height']
img_w = image_info['width']
mask_h, mask_w = segmentation['size']
if img_h == mask_h and img_w == mask_w:
rle = mask_utils.frPyObjects([segmentation], mask_h,
mask_w)[0]
else:
log.warning(
"item #%s: mask #%s "
"does not match image size: %s vs. %s. "
"Skipping this annotation.", image_info['id'],
ann_id, (mask_h, mask_w), (img_h, img_w))
else:
# compressed RLE
rle = segmentation
if rle is not None:
parsed_annotations.append(
RleMask(rle=rle,
label=label_id,
id=ann_id,
attributes=attributes,
group=group))
else:
parsed_annotations.append(
Bbox(x,
y,
w,
h,
label=label_id,
id=ann_id,
attributes=attributes,
group=group))
elif self._task is CocoTask.labels:
label_id = self._get_label_id(ann)
parsed_annotations.append(
Label(label=label_id,
id=ann_id,
attributes=attributes,
group=group))
elif self._task is CocoTask.captions:
caption = ann['caption']
parsed_annotations.append(
Caption(caption, id=ann_id, attributes=attributes,
group=group))
else:
raise NotImplementedError()
return parsed_annotations
def test_stats(self):
dataset = Dataset.from_iterable(
[
DatasetItem(
id=1,
image=np.ones((5, 5, 3)),
annotations=[
Caption('hello'),
Caption('world'),
Label(2, attributes={
'x': 1,
'y': '2',
}),
Bbox(1, 2, 2, 2, label=2, attributes={
'score': 0.5,
}),
Bbox(5,
6,
2,
2,
attributes={
'x': 1,
'y': '3',
'occluded': True,
}),
Points([1, 2, 2, 0, 1, 1], label=0),
Mask(label=3,
image=np.array([
[0, 0, 1, 1, 1],
[0, 0, 1, 1, 1],
[0, 0, 1, 1, 1],
[0, 0, 0, 0, 0],
[0, 0, 0, 0, 0],
])),
]),
DatasetItem(
id=2,
image=np.ones((2, 4, 3)),
annotations=[
Label(2, attributes={
'x': 2,
'y': '2',
}),
Bbox(1, 2, 2, 2, label=3, attributes={
'score': 0.5,
}),
Bbox(5,
6,
2,
2,
attributes={
'x': 2,
'y': '3',
'occluded': False,
}),
]),
DatasetItem(id=3),
],
categories=['label_%s' % i for i in range(4)])
expected = {
'images count': 3,
'annotations count': 10,
'unannotated images count': 1,
'unannotated images': ['3'],
'annotations by type': {
'label': {
'count': 2,
},
'polygon': {
'count': 0,
},
'polyline': {
'count': 0,
},
'bbox': {
'count': 4,
},
'mask': {
'count': 1,
},
'points': {
'count': 1,
},
'caption': {
'count': 2,
},
},
'annotations': {
'labels': {
'count': 6,
'distribution': {
'label_0': [1, 1 / 6],
'label_1': [0, 0.0],
'label_2': [3, 3 / 6],
'label_3': [2, 2 / 6],
},
'attributes': {
'x': {
'count':
2, # unnotations with no label are skipped
'values count': 2,
'values present': ['1', '2'],
'distribution': {
'1': [1, 1 / 2],
'2': [1, 1 / 2],
},
},
'y': {
'count':
2, # unnotations with no label are skipped
'values count': 1,
'values present': ['2'],
'distribution': {
'2': [2, 2 / 2],
},
},
# must not include "special" attributes like "occluded"
}
},
'segments': {
'avg. area': (4 * 2 + 9 * 1) / 3,
'area distribution': [
{
'min': 4.0,
'max': 4.5,
'count': 2,
'percent': 2 / 3
},
{
'min': 4.5,
'max': 5.0,
'count': 0,
'percent': 0.0
},
{
'min': 5.0,
'max': 5.5,
'count': 0,
'percent': 0.0
},
{
'min': 5.5,
'max': 6.0,
'count': 0,
'percent': 0.0
},
{
'min': 6.0,
'max': 6.5,
'count': 0,
'percent': 0.0
},
{
'min': 6.5,
'max': 7.0,
'count': 0,
'percent': 0.0
},
{
'min': 7.0,
'max': 7.5,
'count': 0,
'percent': 0.0
},
{
'min': 7.5,
'max': 8.0,
'count': 0,
'percent': 0.0
},
{
'min': 8.0,
'max': 8.5,
'count': 0,
'percent': 0.0
},
{
'min': 8.5,
'max': 9.0,
'count': 1,
'percent': 1 / 3
},
],
'pixel distribution': {
'label_0': [0, 0.0],
'label_1': [0, 0.0],
'label_2': [4, 4 / 17],
'label_3': [13, 13 / 17],
},
}
},
}
actual = compute_ann_statistics(dataset)
self.assertEqual(expected, actual)
def test_can_compare_projects(self): # just a smoke test
label_categories1 = LabelCategories.from_iterable(['x', 'a', 'b', 'y'])
mask_categories1 = MaskCategories.make_default(len(label_categories1))
point_categories1 = PointsCategories()
for index, _ in enumerate(label_categories1.items):
point_categories1.add(index, ['cat1', 'cat2'], joints=[[0, 1]])
dataset1 = Dataset.from_iterable(
[
DatasetItem(
id=100,
subset='train',
image=np.ones((10, 6, 3)),
annotations=[
Caption('hello', id=1),
Caption('world', id=2, group=5),
Label(2, id=3, attributes={
'x': 1,
'y': '2',
}),
Bbox(1,
2,
3,
4,
label=0,
id=4,
z_order=1,
attributes={
'score': 1.0,
}),
Bbox(5, 6, 7, 8, id=5, group=5),
Points([1, 2, 2, 0, 1, 1], label=0, id=5, z_order=4),
Mask(label=3, id=5, z_order=2, image=np.ones((2, 3))),
]),
DatasetItem(id=21,
subset='train',
annotations=[
Caption('test'),
Label(2),
Bbox(1, 2, 3, 4, label=2, id=42, group=42)
]),
DatasetItem(
id=2,
subset='val',
annotations=[
PolyLine([1, 2, 3, 4, 5, 6, 7, 8], id=11, z_order=1),
Polygon([1, 2, 3, 4, 5, 6, 7, 8], id=12, z_order=4),
]),
DatasetItem(
id=42, subset='test', attributes={
'a1': 5,
'a2': '42'
}),
DatasetItem(id=42),
DatasetItem(id=43, image=Image(path='1/b/c.qq', size=(2, 4))),
],
categories={
AnnotationType.label: label_categories1,
AnnotationType.mask: mask_categories1,
AnnotationType.points: point_categories1,
})
label_categories2 = LabelCategories.from_iterable(['a', 'b', 'x', 'y'])
mask_categories2 = MaskCategories.make_default(len(label_categories2))
point_categories2 = PointsCategories()
for index, _ in enumerate(label_categories2.items):
point_categories2.add(index, ['cat1', 'cat2'], joints=[[0, 1]])
dataset2 = Dataset.from_iterable(
[
DatasetItem(
id=100,
subset='train',
image=np.ones((10, 6, 3)),
annotations=[
Caption('hello', id=1),
Caption('world', id=2, group=5),
Label(2, id=3, attributes={
'x': 1,
'y': '2',
}),
Bbox(1,
2,
3,
4,
label=1,
id=4,
z_order=1,
attributes={
'score': 1.0,
}),
Bbox(5, 6, 7, 8, id=5, group=5),
Points([1, 2, 2, 0, 1, 1], label=0, id=5, z_order=4),
Mask(label=3, id=5, z_order=2, image=np.ones((2, 3))),
]),
DatasetItem(id=21,
subset='train',
annotations=[
Caption('test'),
Label(2),
Bbox(1, 2, 3, 4, label=3, id=42, group=42)
]),
DatasetItem(
id=2,
subset='val',
annotations=[
PolyLine([1, 2, 3, 4, 5, 6, 7, 8], id=11, z_order=1),
Polygon([1, 2, 3, 4, 5, 6, 7, 8], id=12, z_order=4),
]),
DatasetItem(
id=42, subset='test', attributes={
'a1': 5,
'a2': '42'
}),
DatasetItem(id=42),
DatasetItem(id=43, image=Image(path='1/b/c.qq', size=(2, 4))),
],
categories={
AnnotationType.label: label_categories2,
AnnotationType.mask: mask_categories2,
AnnotationType.points: point_categories2,
})
with TestDir() as test_dir:
with DatasetDiffVisualizer(
save_dir=test_dir,
comparator=DistanceComparator(iou_threshold=0.8),
) as visualizer:
visualizer.save(dataset1, dataset2)
self.assertNotEqual(0, os.listdir(osp.join(test_dir)))