def test_shapes_to_boxes(self):
source_dataset = Dataset.from_iterable([
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]],
), id=1),
Polygon([1, 1, 4, 1, 4, 4, 1, 4], id=2),
PolyLine([1, 1, 2, 1, 2, 2, 1, 2], id=3),
Points([2, 2, 4, 2, 4, 4, 2, 4], id=4),
]
),
])
target_dataset = Dataset.from_iterable([
DatasetItem(id=1, image=np.zeros((5, 5, 3)),
annotations=[
Bbox(0, 0, 4, 4, id=1),
Bbox(1, 1, 3, 3, id=2),
Bbox(1, 1, 1, 1, id=3),
Bbox(2, 2, 2, 2, id=4),
]
),
])
actual = transforms.ShapesToBoxes(source_dataset)
compare_datasets(self, target_dataset, actual)
def test_can_resize(self):
small_dataset = Dataset.from_iterable([
DatasetItem(id=i, image=np.ones((4, 4)) * i, annotations=[
Label(1),
Bbox(1, 1, 2, 2, label=2),
Polygon([1, 1, 1, 2, 2, 2, 2, 1], label=1),
PolyLine([1, 1, 1, 2, 2, 2, 2, 1], label=2),
Points([1, 1, 1, 2, 2, 2, 2, 1], label=2),
Mask(np.array([
[0, 0, 1, 1],
[1, 0, 0, 1],
[0, 1, 1, 0],
[1, 1, 0, 0],
]))
]) for i in range(3)
], categories=['a', 'b', 'c'])
big_dataset = Dataset.from_iterable([
DatasetItem(id=i, image=np.ones((8, 8)) * i, annotations=[
Label(1),
Bbox(2, 2, 4, 4, label=2),
Polygon([2, 2, 2, 4, 4, 4, 4, 2], label=1),
PolyLine([2, 2, 2, 4, 4, 4, 4, 2], label=2),
Points([2, 2, 2, 4, 4, 4, 4, 2], label=2),
Mask(np.array([
[0, 0, 0, 0, 1, 1, 1, 1],
[0, 0, 0, 0, 1, 1, 1, 1],
[1, 1, 0, 0, 0, 0, 1, 1],
[1, 1, 0, 0, 0, 0, 1, 1],
[0, 0, 1, 1, 1, 1, 0, 0],
[0, 0, 1, 1, 1, 1, 0, 0],
[1, 1, 1, 1, 0, 0, 0, 0],
[1, 1, 1, 1, 0, 0, 0, 0],
]))
]) for i in range(3)
], categories=['a', 'b', 'c'])
with self.subTest('upscale'):
actual = transforms.ResizeTransform(small_dataset, width=8, height=8)
compare_datasets(self, big_dataset, actual)
with self.subTest('downscale'):
actual = transforms.ResizeTransform(big_dataset, width=4, height=4)
compare_datasets(self, small_dataset, actual)
def test_can_match_lines_when_line_not_approximated(self):
source0 = Dataset.from_iterable([
DatasetItem(1, annotations=[
PolyLine([1, 1, 2, 1, 3, 5, 5, 5, 8, 3]),
]),
])
source1 = Dataset.from_iterable([
DatasetItem(1, annotations=[
PolyLine([1, 1, 8, 3]),
]),
])
expected = Dataset.from_iterable([
DatasetItem(1, annotations=[
PolyLine([1, 1, 2, 1, 3, 5, 5, 5, 8, 3]),
]),
], categories=[])
merger = IntersectMerge(conf={'quorum': 1, 'pairwise_dist': 0.1})
merged = merger([source0, source1])
compare_datasets(self, expected, merged, ignored_attrs={'score'})
self.assertEqual(0, len(merger.errors))
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))
elif ann_type == AnnotationType.cuboid_3d:
loaded.append(
Cuboid3d(ann.get('position'),
ann.get('rotation'),
ann.get('scale'),
label=label_id,
id=ann_id,
attributes=attributes,
group=group))
else:
raise NotImplementedError()
return loaded
def test_can_load_video(self):
expected_dataset = Dataset.from_iterable([
DatasetItem(id='frame_000010', subset='annotations',
image=255 * 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',
}),
], attributes={'frame': 10}),
DatasetItem(id='frame_000013', subset='annotations',
image=255 * 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,
}),
], attributes={'frame': 13}),
DatasetItem(id='frame_000016', subset='annotations',
image=Image(path='frame_0000016.png', size=(20, 25)),
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,
}),
], attributes={'frame': 16}),
], categories={
AnnotationType.label: LabelCategories.from_iterable([
['klhg', '', {'hgl'}],
['z U k'],
['II']
]),
})
parsed_dataset = Dataset.import_from(DUMMY_VIDEO_DATASET_DIR, 'cvat')
compare_datasets(self, expected_dataset, parsed_dataset)
def test_can_save_and_load(self):
src_label_cat = LabelCategories(attributes={'occluded', 'common'})
for i in range(10):
src_label_cat.add(str(i))
src_label_cat.items[2].attributes.update(['a1', 'a2', 'empty'])
source_dataset = Dataset.from_iterable([
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, 'common': 't' }),
Points([1, 1, 3, 2, 2, 3],
label=2,
attributes={ 'a1': 'x', 'a2': 42, 'empty': '',
'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))),
], categories={ AnnotationType.label: src_label_cat })
target_label_cat = LabelCategories(
attributes={'occluded'}) # unable to represent a common attribute
for i in range(10):
target_label_cat.add(str(i), attributes={'common'})
target_label_cat.items[2].attributes.update(['a1', 'a2', 'empty', 'common'])
target_dataset = Dataset.from_iterable([
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, 'common': 't' }),
Points([1, 1, 3, 2, 2, 3],
label=2,
attributes={ 'occluded': False, 'empty': '',
'a1': 'x', 'a2': '42' }),
Label(1),
Label(2, attributes={ 'a1': 'y', 'a2': '44' }),
], attributes={'frame': 0}
),
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 }),
], attributes={'frame': 1}
),
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 }),
], attributes={'frame': 0}
),
DatasetItem(id=3, subset='s3', image=Image(
path='3.jpg', size=(2, 4)),
attributes={'frame': 0}),
], categories={ AnnotationType.label: target_label_cat })
with TestDir() as test_dir:
self._test_save_and_load(source_dataset,
partial(CvatConverter.convert, save_images=True), test_dir,
target_dataset=target_dataset)
def test_can_match_shapes(self):
source0 = Dataset.from_iterable([
DatasetItem(1, annotations=[
# unique
Bbox(1, 2, 3, 4, label=1),
# common
Mask(label=2, z_order=2, image=np.array([
[0, 0, 0, 0],
[0, 0, 0, 0],
[1, 1, 1, 0],
[1, 1, 1, 0],
])),
Polygon([1, 0, 3, 2, 1, 2]),
# an instance with keypoints
Bbox(4, 5, 2, 4, label=2, z_order=1, group=1),
Points([5, 6], label=0, group=1),
Points([6, 8], label=1, group=1),
PolyLine([1, 1, 2, 1, 3, 1]),
]),
], categories=['a', 'b', 'c'])
source1 = Dataset.from_iterable([
DatasetItem(1, annotations=[
# common
Mask(label=2, image=np.array([
[0, 0, 0, 0],
[0, 1, 1, 1],
[0, 1, 1, 1],
[0, 1, 1, 1],
])),
Polygon([0, 2, 2, 0, 2, 1]),
# an instance with keypoints
Bbox(4, 4, 2, 5, label=2, z_order=1, group=2),
Points([5.5, 6.5], label=0, group=2),
Points([6, 8], label=1, group=2),
PolyLine([1, 1.5, 2, 1.5]),
]),
], categories=['a', 'b', 'c'])
source2 = Dataset.from_iterable([
DatasetItem(1, annotations=[
# common
Mask(label=2, z_order=3, image=np.array([
[0, 0, 1, 1],
[0, 1, 1, 1],
[1, 1, 1, 1],
[1, 1, 1, 0],
])),
Polygon([3, 1, 2, 2, 0, 1]),
# an instance with keypoints, one is missing
Bbox(3, 6, 2, 3, label=2, z_order=4, group=3),
Points([4.5, 5.5], label=0, group=3),
PolyLine([1, 1.25, 3, 1, 4, 2]),
]),
], categories=['a', 'b', 'c'])
expected = Dataset.from_iterable([
DatasetItem(1, annotations=[
# unique
Bbox(1, 2, 3, 4, label=1),
# common
# nearest to mean bbox
Mask(label=2, z_order=3, image=np.array([
[0, 0, 0, 0],
[0, 1, 1, 1],
[0, 1, 1, 1],
[0, 1, 1, 1],
])),
Polygon([1, 0, 3, 2, 1, 2]),
# an instance with keypoints
Bbox(4, 5, 2, 4, label=2, z_order=4, group=1),
Points([5, 6], label=0, group=1),
Points([6, 8], label=1, group=1),
PolyLine([1, 1.25, 3, 1, 4, 2]),
]),
], categories=['a', 'b', 'c'])
merger = IntersectMerge(conf={'quorum': 1, 'pairwise_dist': 0.1})
merged = merger([source0, source1, source2])
compare_datasets(self, expected, merged, ignored_attrs={'score'})
self.assertEqual(
[
NoMatchingAnnError(item_id=('1', DEFAULT_SUBSET_NAME),
sources={2}, ann=source0.get('1').annotations[5]),
NoMatchingAnnError(item_id=('1', DEFAULT_SUBSET_NAME),
sources={1, 2}, ann=source0.get('1').annotations[0]),
],
sorted((e for e in merger.errors
if isinstance(e, NoMatchingAnnError)),
key=lambda e: len(e.sources))
)
def _parse_shape_ann(cls, ann, categories):
ann_id = ann.get('id', 0)
ann_type = ann['type']
attributes = ann.get('attributes') or {}
if 'occluded' in categories[AnnotationType.label].attributes:
attributes['occluded'] = ann.get('occluded', False)
if 'outside' in ann:
attributes['outside'] = ann['outside']
if 'keyframe' in ann:
attributes['keyframe'] = ann['keyframe']
if 'track_id' in ann:
attributes['track_id'] = ann['track_id']
group = ann.get('group')
label = ann.get('label')
label_id = categories[AnnotationType.label].find(label)[0]
z_order = ann.get('z_order', 0)
points = ann.get('points', [])
if ann_type == 'polyline':
return PolyLine(points,
label=label_id,
z_order=z_order,
id=ann_id,
attributes=attributes,
group=group)
elif ann_type == 'polygon':
return Polygon(points,
label=label_id,
z_order=z_order,
id=ann_id,
attributes=attributes,
group=group)
elif ann_type == 'points':
return Points(points,
label=label_id,
z_order=z_order,
id=ann_id,
attributes=attributes,
group=group)
elif ann_type == 'box':
x, y = points[0], points[1]
w, h = points[2] - x, points[3] - y
return Bbox(x,
y,
w,
h,
label=label_id,
z_order=z_order,
id=ann_id,
attributes=attributes,
group=group)
else:
raise NotImplementedError("Unknown annotation type '%s'" %
ann_type)
def test_can_compare_projects(self): # just a smoke test
label_categories1 = LabelCategories.from_iterable(['x', 'a', 'b', 'y'])
mask_categories1 = MaskCategories.generate(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.generate(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 DiffVisualizer(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)))
def test_dataset(self):
label_categories = LabelCategories(attributes={'a', 'b', 'score'})
for i in range(5):
label_categories.add('cat' + str(i), attributes={'x', 'y'})
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,
attributes={
'a': 1.5,
'b': 'text',
}),
Points([1, 2, 2, 0, 1, 1],
label=0,
id=5,
z_order=4,
attributes={
'x': 1,
'y': '2',
}),
Mask(label=3,
id=5,
z_order=2,
image=np.ones((2, 3)),
attributes={
'x': 1,
'y': '2',
}),
]),
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=1,
subset='test',
annotations=[
Cuboid3d([1.0, 2.0, 3.0], [2.0, 2.0, 4.0],
[1.0, 3.0, 4.0],
id=6,
label=0,
attributes={'occluded': True},
group=6)
]),
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,
})