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 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 deleted Polygon([1, 1, 2, 2, 3, 4], label=4), # Should be kept PolyLine([1, 3, 4, 2, 5, 6]), Bbox(4, 3, 2, 1, label=5), ]) ], categories={ AnnotationType.label: LabelCategories.from_iterable( f'label{i}' for i in range(6)), AnnotationType.mask: MaskCategories( colormap=mask_tools.generate_colormap(6)), AnnotationType.points: PointsCategories.from_iterable( [(i, [str(i)]) for i in range(6)]) }) 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), PolyLine([1, 3, 4, 2, 5, 6], label=None), Bbox(4, 3, 2, 1, label=2), ]), ], categories={ AnnotationType.label: LabelCategories.from_iterable( ['label0', 'label9', 'label5']), AnnotationType.mask: MaskCategories(colormap={ i: v for i, v in enumerate({ k: v for k, v in mask_tools.generate_colormap(6).items() if k in { 0, 1, 5 } }.values()) }), AnnotationType.points: PointsCategories.from_iterable( [(0, ['0']), (1, ['1']), (2, ['5'])]) }) actual = transforms.RemapLabels(src_dataset, mapping={ 'label1': 'label9', # rename & join with new label9 (from label3) 'label2': 'label0', # rename & join with existing label0 'label3': 'label9', # rename & join with new label9 (from label1) 'label4': '', # delete the label and associated annotations # 'label5' - unchanged }, default='keep') compare_datasets(self, dst_dataset, actual)
def test_project_labels_maps_secondary_categories(self): source = Dataset.from_iterable([], categories={ AnnotationType.label: LabelCategories.from_iterable([ 'a', 'b', # no parents ('c', 'a'), ('d', 'b') # have parents ]), AnnotationType.points: PointsCategories.from_iterable([ (0, ['a']), (1, ['b']), (2, ['c']) ]), AnnotationType.mask: MaskCategories.generate(4) }) expected = Dataset.from_iterable([], categories={ AnnotationType.label: LabelCategories.from_iterable([ ('c', 'a'), # must keep parent 'a', 'd' # must drop parent because it was removed ]), AnnotationType.points: PointsCategories.from_iterable([ (0, ['c']), (1, ['a']) ]), AnnotationType.mask: MaskCategories(colormap={ i: v for i, v in { { 2: 0, 0: 1, 3: 2 }.get(k): v for k, v in mask_tools.generate_colormap(4).items() }.items() if i is not None }), }) actual = transforms.ProjectLabels(source, dst_labels=['c', 'a', 'd']) compare_datasets(self, expected, actual)
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 generate(cls, size: int = 255, include_background: bool = True) \ -> MaskCategories: """ Generates MaskCategories with the specified size. If include_background is True, the result will include the item "0: (0, 0, 0)", which is typically used as a background color. """ from datumaro.util.mask_tools import generate_colormap return cls( generate_colormap(size, include_background=include_background))
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('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_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 test_can_paint_mask(self): mask = np.zeros((1, 3), dtype=np.uint8) mask[:, 0] = 0 mask[:, 1] = 1 mask[:, 2] = 2 colormap = mask_tools.generate_colormap(3) expected = np.zeros((*mask.shape, 3), dtype=np.uint8) expected[:, 0] = colormap[0][::-1] expected[:, 1] = colormap[1][::-1] expected[:, 2] = colormap[2][::-1] actual = mask_tools.paint_mask(mask, colormap) self.assertTrue(np.array_equal(expected, actual), '%s\nvs.\n%s' % (expected, actual))
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_mapillary_instance_categories(label_map): label_cat = LabelCategories() for label, _ in label_map.items(): label_cat.add(label) has_colors = any([len(color) == 3 for color in label_map.values()]) if not has_colors: colormap = generate_colormap(len(label_map)) else: colormap = { label_cat.find(label)[0]: color for label, color in label_map.items() } mask_cat = MaskCategories(colormap) mask_cat.inverse_colormap # pylint: disable=pointless-statement return {AnnotationType.label: label_cat, AnnotationType.mask: mask_cat}
def make_colormap(task_data): labels = sorted( [label['name'] for _, label in task_data.meta['task']['labels']]) if 'background' not in labels: labels.insert(0, 'background') predefined = parse_default_colors() # NOTE: using pop() to avoid collisions colormap = {k: predefined.pop(normalize_label(k), None) for k in labels} random_labels = [k for k in labels if not colormap[k]] if random_labels: colors = generate_colormap(DEFAULT_COLORMAP_CAPACITY + len(random_labels)) for i, label in enumerate(random_labels): colormap[label] = colors[DEFAULT_COLORMAP_CAPACITY + i] return {l: [c, [], []] for l, c in colormap.items()}
def test_can_unpaint_mask(self): colormap = mask_tools.generate_colormap(3) inverse_colormap = mask_tools.invert_colormap(colormap) mask = np.zeros((1, 3, 3), dtype=np.uint8) mask[:, 0] = colormap[0][::-1] mask[:, 1] = colormap[1][::-1] mask[:, 2] = colormap[2][::-1] expected = np.zeros((1, 3), dtype=np.uint8) expected[:, 0] = 0 expected[:, 1] = 1 expected[:, 2] = 2 actual = mask_tools.unpaint_mask(mask, inverse_colormap) self.assertTrue(np.array_equal(expected, actual), '%s\nvs.\n%s' % (expected, actual))
def make_categories(label_map=None): if label_map is None: label_map = SYNTHIA_LABEL_MAP categories = {} label_categories = LabelCategories() for label in label_map: 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 colormap = { label_id: (desc[0], desc[1], desc[2]) for label_id, desc in enumerate(label_map.values()) } mask_categories = MaskCategories(colormap) mask_categories.inverse_colormap # pylint: disable=pointless-statement categories[AnnotationType.mask] = mask_categories return categories
def _save_item(self, subset_name, subset, item): if self._save_images and item.has_image: self._save_image(item, subdir=osp.join(subset_name, IcdarPath.IMAGES_DIR)) annotation = '' anns = [a for a in item.annotations if a.type == AnnotationType.mask] color_bank = iter( generate_colormap(len(anns), include_background=False).values()) colormap = [(255, 255, 255)] used_colors = set(colormap) if anns: anns = sorted(anns, key=lambda a: int(a.attributes.get('index', 0))) group = anns[0].group for i, ann in enumerate(anns): # Assign new color if it is not defined color = ann.attributes.get('color', '') if color: color = color.split() if len(color) != 3: raise DatumaroError( "Item %s: mask #%s has invalid color" % (item.id, i)) color = tuple(map(int, color)) else: color = next(color_bank) while color in used_colors: color = next(color_bank) colormap.append(color) used_colors.add(color) text = ann.attributes.get('text', '') bbox = ann.get_bbox() if ann.group != group or (not ann.group and anns[0].group != 0): annotation += '\n' if text == ' ': annotation += '#' annotation += ' '.join(str(p) for p in color) annotation += ' %s' % ann.attributes.get('center', '- -') annotation += ' %s %s %s %s' % (bbox[0], bbox[1], bbox[0] + bbox[2], bbox[1] + bbox[3]) annotation += ' \"%s\"' % text annotation += '\n' group = ann.group mask = CompiledMask.from_instance_masks( anns, instance_labels=[m.attributes['index'] + 1 for m in anns]) mask = paint_mask(mask.class_mask, {i: colormap[i] for i in range(len(colormap))}) save_image(osp.join(self._save_dir, subset_name, item.id + '_GT' + IcdarPath.GT_EXT), mask, create_dir=True) anno_file = osp.join(self._save_dir, subset_name, item.id + '_GT' + '.txt') os.makedirs(osp.dirname(anno_file), exist_ok=True) with open(anno_file, 'w', encoding='utf-8') as f: f.write(annotation)
def __init__(self, extractor: IExtractor, dst_labels: Union[Iterable[str], LabelCategories]): super().__init__(extractor) self._categories = {} src_categories = self._extractor.categories() src_label_cat = src_categories.get(AnnotationType.label) if isinstance(dst_labels, LabelCategories): dst_label_cat = deepcopy(dst_labels) else: dst_labels = list(dst_labels) if src_label_cat: dst_label_cat = LabelCategories( attributes=deepcopy(src_label_cat.attributes)) for dst_label in dst_labels: assert isinstance(dst_label, str) src_label = src_label_cat.find(dst_label)[1] if src_label is not None: dst_label_cat.add(dst_label, src_label.parent, deepcopy(src_label.attributes)) else: dst_label_cat.add(dst_label) else: dst_label_cat = LabelCategories.from_iterable(dst_labels) for label in dst_label_cat: if label.parent not in dst_label_cat: label.parent = '' self._categories[AnnotationType.label] = dst_label_cat self._make_label_id_map(src_label_cat, dst_label_cat) src_mask_cat = src_categories.get(AnnotationType.mask) if src_mask_cat is not None: assert src_label_cat is not None dst_mask_cat = MaskCategories( attributes=deepcopy(src_mask_cat.attributes)) for old_id, old_color in src_mask_cat.colormap.items(): new_id = self._map_id(old_id) if new_id is not None and new_id not in dst_mask_cat: dst_mask_cat.colormap[new_id] = deepcopy(old_color) # Generate new colors for new labels, keep old untouched existing_colors = set(dst_mask_cat.colormap.values()) color_bank = iter( mask_tools.generate_colormap( len(dst_label_cat), include_background=False).values()) for new_id, new_label in enumerate(dst_label_cat): if new_label.name in src_label_cat: continue if new_id in dst_mask_cat: continue color = next(color_bank) while color in existing_colors: color = next(color_bank) dst_mask_cat.colormap[new_id] = color self._categories[AnnotationType.mask] = dst_mask_cat src_point_cat = src_categories.get(AnnotationType.points) if src_point_cat is not None: assert src_label_cat is not None dst_point_cat = PointsCategories( attributes=deepcopy(src_point_cat.attributes)) for old_id, old_cat in src_point_cat.items.items(): new_id = self._map_id(old_id) if new_id is not None and new_id not in dst_point_cat: dst_point_cat.items[new_id] = deepcopy(old_cat) self._categories[AnnotationType.points] = dst_point_cat
def make_default(cls, size=256): from datumaro.util.mask_tools import generate_colormap return cls(generate_colormap(size))