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