def _load_categories(parsed):
categories = {}
parsed_label_cat = parsed['categories'].get(AnnotationType.label.name)
if parsed_label_cat:
label_categories = LabelCategories()
for item in parsed_label_cat['labels']:
label_categories.add(item['name'], parent=item['parent'])
categories[AnnotationType.label] = label_categories
parsed_mask_cat = parsed['categories'].get(AnnotationType.mask.name)
if parsed_mask_cat:
colormap = {}
for item in parsed_mask_cat['colormap']:
colormap[int(item['label_id'])] = \
(item['r'], item['g'], item['b'])
mask_categories = MaskCategories(colormap=colormap)
categories[AnnotationType.mask] = mask_categories
parsed_points_cat = parsed['categories'].get(
AnnotationType.points.name)
if parsed_points_cat:
point_categories = PointsCategories()
for item in parsed_points_cat['items']:
point_categories.add(int(item['label_id']),
item['labels'],
joints=item['joints'])
categories[AnnotationType.points] = point_categories
return categories
def categories(self):
label_cat = LabelCategories()
point_cat = PointsCategories()
for label in range(10):
label_cat.add('label_' + str(label))
point_cat.add(label)
return {
AnnotationType.label: label_cat,
AnnotationType.points: point_cat,
}
def categories(self):
label_categories = LabelCategories()
points_categories = PointsCategories()
for i in range(10):
label_categories.add(str(i))
points_categories.add(i, [])
return {
AnnotationType.label: label_categories,
AnnotationType.points: points_categories,
}
def _load_person_kp_categories(self, loader):
catIds = loader.getCatIds()
cats = loader.loadCats(catIds)
categories = PointsCategories()
for cat in cats:
label_id = self._label_map[cat['id']]
categories.add(label_id=label_id,
labels=cat['keypoints'],
joints=cat['skeleton'])
return categories
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 categories(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'], adjacent=[0, 1])
return {
AnnotationType.label: label_categories,
AnnotationType.mask: mask_categories,
AnnotationType.points: points_categories,
}
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_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)))