def test_remap_labels(self):
src_dataset = Dataset.from_iterable(
[
DatasetItem(
id=1,
annotations=[
# Should be remapped
Label(1),
Bbox(1, 2, 3, 4, label=2),
Mask(image=np.array([1]), label=3),
# Should be kept
Polygon([1, 1, 2, 2, 3, 4], label=4),
PolyLine([1, 3, 4, 2, 5, 6])
])
],
categories={
AnnotationType.label:
LabelCategories.from_iterable('label%s' % i for i in range(5)),
AnnotationType.mask:
MaskCategories(colormap=mask_tools.generate_colormap(5)),
})
dst_dataset = Dataset.from_iterable(
[
DatasetItem(id=1,
annotations=[
Label(1),
Bbox(1, 2, 3, 4, label=0),
Mask(image=np.array([1]), label=1),
Polygon([1, 1, 2, 2, 3, 4], label=2),
PolyLine([1, 3, 4, 2, 5, 6], label=None)
]),
],
categories={
AnnotationType.label:
LabelCategories.from_iterable(['label0', 'label9', 'label4']),
AnnotationType.mask:
MaskCategories(
colormap={
k: v
for k, v in mask_tools.generate_colormap(5).items()
if k in {0, 1, 3, 4}
})
})
actual = transforms.RemapLabels(src_dataset,
mapping={
'label1': 'label9',
'label2': 'label0',
'label3': 'label9',
},
default='keep')
compare_datasets(self, dst_dataset, actual)
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 make_voc_categories(label_map=None):
if label_map is None:
label_map = make_voc_label_map()
categories = {}
label_categories = LabelCategories()
label_categories.attributes.update(['difficult', 'truncated', 'occluded'])
for label, desc in label_map.items():
label_categories.add(label, attributes=desc[2])
for part in OrderedDict((k, None) for k in chain(
*(desc[1] for desc in label_map.values()))):
label_categories.add(part)
categories[AnnotationType.label] = label_categories
has_colors = any(v[0] is not None for v in label_map.values())
if not has_colors: # generate new colors
colormap = generate_colormap(len(label_map))
else: # only copy defined colors
label_id = lambda label: label_categories.find(label)[0]
colormap = { label_id(name): desc[0]
for name, desc in label_map.items() if desc[0] is not None }
mask_categories = MaskCategories(colormap)
mask_categories.inverse_colormap # pylint: disable=pointless-statement
categories[AnnotationType.mask] = mask_categories
return categories
def make_camvid_categories(label_map=None):
if label_map is None:
label_map = CamvidLabelMap
# There must always be a label with color (0, 0, 0) at index 0
bg_label = find(label_map.items(), lambda x: x[1] == (0, 0, 0))
if bg_label is not None:
bg_label = bg_label[0]
else:
bg_label = 'background'
if bg_label not in label_map:
has_colors = any(v is not None for v in label_map.values())
color = (0, 0, 0) if has_colors else None
label_map[bg_label] = color
label_map.move_to_end(bg_label, last=False)
categories = {}
label_categories = LabelCategories()
for label, desc in label_map.items():
label_categories.add(label)
categories[AnnotationType.label] = label_categories
has_colors = any(v is not None for v in label_map.values())
if not has_colors: # generate new colors
colormap = generate_colormap(len(label_map))
else: # only copy defined colors
label_id = lambda label: label_categories.find(label)[0]
colormap = {
label_id(name): (desc[0], desc[1], desc[2])
for name, desc in label_map.items()
}
mask_categories = MaskCategories(colormap)
mask_categories.inverse_colormap # pylint: disable=pointless-statement
categories[AnnotationType.mask] = mask_categories
return categories
def __init__(self, extractor, mapping, default=None):
super().__init__(extractor)
assert isinstance(default, (str, self.DefaultAction))
if isinstance(default, str):
default = self.DefaultAction[default]
assert isinstance(mapping, (dict, list))
if isinstance(mapping, list):
mapping = dict(mapping)
self._categories = {}
src_label_cat = self._extractor.categories().get(AnnotationType.label)
if src_label_cat is not None:
self._make_label_id_map(src_label_cat, mapping, default)
src_mask_cat = self._extractor.categories().get(AnnotationType.mask)
if src_mask_cat is not None:
assert src_label_cat is not None
dst_mask_cat = MaskCategories(attributes=src_mask_cat.attributes)
dst_mask_cat.colormap = {
id: src_mask_cat.colormap[id]
for id, _ in enumerate(src_label_cat.items)
if self._map_id(id) or id == 0
}
self._categories[AnnotationType.mask] = dst_mask_cat
src_points_cat = self._extractor.categories().get(
AnnotationType.points)
if src_points_cat is not None:
assert src_label_cat is not None
dst_points_cat = PointsCategories(
attributes=src_points_cat.attributes)
dst_points_cat.items = {
id: src_points_cat.items[id]
for id, item in enumerate(src_label_cat.items)
if self._map_id(id) or id == 0
}
self._categories[AnnotationType.points] = dst_points_cat
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_cat = LabelCategories()
label_cat.add('label0')
label_cat.add('label1')
label_cat.add('label2')
label_cat.add('label3')
label_cat.add('label4')
mask_cat = MaskCategories(
colormap=mask_tools.generate_colormap(5))
return {
AnnotationType.label: label_cat,
AnnotationType.mask: mask_cat,
}
def categories(self):
label_cat = LabelCategories()
label_cat.add('label0')
label_cat.add('label9')
label_cat.add('label4')
mask_cat = MaskCategories(colormap={
k: v for k, v in mask_tools.generate_colormap(5).items()
if k in { 0, 1, 3, 4 }
})
return {
AnnotationType.label: label_cat,
AnnotationType.mask: mask_cat,
}
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 _make_voc_categories():
categories = {}
label_categories = LabelCategories()
for label in chain(VocLabel, VocAction, VocBodyPart):
label_categories.add(label.name)
categories[AnnotationType.label] = label_categories
def label_id(class_index):
class_label = VocLabel(class_index).name
label_id, _ = label_categories.find(class_label)
return label_id
colormap = { label_id(idx): tuple(color) \
for idx, color in VocColormap.items() }
mask_categories = MaskCategories(colormap)
mask_categories.inverse_colormap # init inverse colormap
categories[AnnotationType.mask] = mask_categories
return categories
def _load_segmentation_items(self):
items = {}
image_dir = osp.join(self._path, IcdarPath.IMAGES_DIR)
if osp.isdir(image_dir):
images = {
osp.splitext(osp.relpath(p, image_dir))[0]: p
for p in find_images(image_dir, recursive=True)
}
else:
images = {}
for path in iglob(osp.join(self._path, '**', '*.txt'), recursive=True):
item_id = osp.splitext(osp.relpath(path, self._path))[0]
item_id = item_id.replace('\\', '/')
if item_id.endswith('_GT'):
item_id = item_id[:-3]
if item_id not in items:
items[item_id] = DatasetItem(item_id,
subset=self._subset,
image=images.get(item_id))
annotations = items[item_id].annotations
colors = [(255, 255, 255)]
chars = ['']
centers = [0]
groups = [0]
group = 1
number_in_group = 0
with open(path, encoding='utf-8') as f:
for line in f:
line = line.strip()
if line == '':
if number_in_group == 1:
groups[len(groups) - 1] = 0
else:
group += 1
number_in_group = 0
continue
objects = line.split()
if objects[0][0] == '#':
objects[0] = objects[0][1:]
objects[9] = '\" \"'
objects.pop()
if len(objects) != 10:
raise Exception(
"Line %s contains the wrong number "
"of arguments, e.g. '241 73 144 1 4 0 3 1 4 \"h\""
% line)
centers.append(objects[3] + ' ' + objects[4])
groups.append(group)
colors.append(tuple(int(o) for o in objects[:3]))
char = objects[9]
if char[0] == '\"' and char[-1] == '\"':
char = char[1:-1]
chars.append(char)
number_in_group += 1
if number_in_group == 1:
groups[len(groups) - 1] = 0
mask_categories = MaskCategories(
{i: colors[i]
for i in range(len(colors))})
inverse_cls_colormap = mask_categories.inverse_colormap
gt_path = osp.join(self._path, item_id + '_GT' + IcdarPath.GT_EXT)
if osp.isfile(gt_path):
# load mask through cache
mask = lazy_mask(gt_path, inverse_cls_colormap)
mask = mask()
classes = np.unique(mask)
for label_id in classes:
if label_id == 0:
continue
i = int(label_id)
annotations.append(
Mask(group=groups[i],
image=self._lazy_extract_mask(mask, label_id),
attributes={
'index': i - 1,
'color': ' '.join(str(p) for p in colors[i]),
'text': chars[i],
'center': centers[i]
}))
return items
def test_can_merge_categories(self):
source0 = Dataset.from_iterable(
[
DatasetItem(1, annotations=[
Label(0),
]),
],
categories={
AnnotationType.label:
LabelCategories.from_iterable(['a', 'b']),
AnnotationType.points:
PointsCategories.from_iterable([
(0, ['l0', 'l1']),
(1, ['l2', 'l3']),
]),
AnnotationType.mask:
MaskCategories({
0: (0, 1, 2),
1: (1, 2, 3),
}),
})
source1 = Dataset.from_iterable(
[
DatasetItem(1, annotations=[
Label(0),
]),
],
categories={
AnnotationType.label:
LabelCategories.from_iterable(['c', 'b']),
AnnotationType.points:
PointsCategories.from_iterable([
(0, []),
(1, ['l2', 'l3']),
]),
AnnotationType.mask:
MaskCategories({
0: (0, 2, 4),
1: (1, 2, 3),
}),
})
expected = Dataset.from_iterable(
[
DatasetItem(1, annotations=[
Label(0),
Label(2),
]),
],
categories={
AnnotationType.label:
LabelCategories.from_iterable(['a', 'b', 'c']),
AnnotationType.points:
PointsCategories.from_iterable([
(0, ['l0', 'l1']),
(1, ['l2', 'l3']),
(2, []),
]),
AnnotationType.mask:
MaskCategories({
0: (0, 1, 2),
1: (1, 2, 3),
2: (0, 2, 4),
}),
})
merger = IntersectMerge()
merged = merger([source0, source1])
compare_datasets(self, expected, merged, ignored_attrs={'score'})
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)))
def _load_segmentation_items(self):
items = {}
for path in glob(osp.join(self._path, '*.txt')):
item_id = osp.splitext(osp.basename(path))[0]
if item_id.endswith('_GT'):
item_id = item_id[:-3]
image_path = osp.join(self._path, IcdarPath.IMAGES_DIR,
item_id + IcdarPath.IMAGE_EXT)
if item_id not in items:
items[item_id] = DatasetItem(item_id,
subset=self._subset,
image=image_path)
annotations = items[item_id].annotations
colors = [(255, 255, 255)]
chars = ['']
centers = [0]
groups = [0]
group = 1
number_in_group = 0
with open(path, encoding='utf-8') as f:
for line in f:
line = line.strip()
if line == '':
if number_in_group == 1:
groups[len(groups) - 1] = 0
else:
group += 1
number_in_group = 0
continue
objects = line.split()
if objects[0][0] == '#':
objects[0] = objects[0][1:]
objects[9] = '\" \"'
objects.pop()
if len(objects) != 10:
continue
centers.append(objects[3] + ' ' + objects[4])
groups.append(group)
colors.append(tuple(int(o) for o in objects[:3]))
char = objects[9]
if char[0] == '\"' and char[-1] == '\"':
char = char[1:-1]
chars.append(char)
number_in_group += 1
if number_in_group == 1:
groups[len(groups) - 1] = 0
mask_categories = MaskCategories(
{i: colors[i]
for i in range(len(colors))})
inverse_cls_colormap = mask_categories.inverse_colormap
gt_path = osp.join(self._path, item_id + '_GT' + IcdarPath.GT_EXT)
if osp.isfile(gt_path):
# load mask through cache
mask = lazy_mask(gt_path, inverse_cls_colormap)
mask = mask()
classes = np.unique(mask)
for label_id in classes:
if label_id == 0:
continue
i = int(label_id)
annotations.append(
Mask(group=groups[i],
image=self._lazy_extract_mask(mask, label_id),
attributes={
'index': i - 1,
'color': ' '.join(str(p) for p in colors[i]),
'text': chars[i],
'center': centers[i]
}))
return items
