def test_can_import(self): with TestDir() as temp_dir: self.COCO_dataset_generate(temp_dir.path) project = Project.import_from(temp_dir.path, 'ms_coco') dataset = project.make_dataset() self.assertListEqual(['val'], sorted(dataset.subsets())) self.assertEqual(1, len(dataset)) item = next(iter(dataset)) self.assertTrue(item.has_image) self.assertEqual(np.sum(item.image), np.prod(item.image.shape)) self.assertEqual(4, len(item.annotations)) ann_1 = find(item.annotations, lambda x: x.id == 1) ann_1_poly = find(item.annotations, lambda x: \ x.group == ann_1.id and x.type == AnnotationType.polygon) self.assertFalse(ann_1 is None) self.assertFalse(ann_1_poly is None) ann_2 = find(item.annotations, lambda x: x.id == 2) ann_2_mask = find(item.annotations, lambda x: \ x.group == ann_2.id and x.type == AnnotationType.mask) self.assertFalse(ann_2 is None) self.assertFalse(ann_2_mask is None)
def _test_save_and_load(self, source_dataset, converter_type, test_dir, importer_params=None): converter = converter_type() converter(source_dataset, test_dir.path) if not importer_params: importer_params = {} project = Project.import_from(test_dir.path, 'ms_coco', **importer_params) parsed_dataset = project.make_dataset() source_subsets = [ s if s else DEFAULT_SUBSET_NAME for s in source_dataset.subsets() ] self.assertListEqual( sorted(source_subsets), sorted(parsed_dataset.subsets()), ) self.assertEqual(len(source_dataset), len(parsed_dataset)) for item_a in source_dataset: item_b = find(parsed_dataset, lambda x: x.id == item_a.id) self.assertFalse(item_b is None) self.assertEqual(len(item_a.annotations), len(item_b.annotations)) for ann_a in item_a.annotations: ann_b = find(item_b.annotations, lambda x: \ x.id == ann_a.id and \ x.type == ann_a.type and x.group == ann_a.group) self.assertEqual(ann_a, ann_b, 'id: ' + str(ann_a.id))
def compare_datasets(test, expected, actual, ignored_attrs=None, require_images=False): compare_categories(test, expected.categories(), actual.categories()) test.assertEqual(sorted(expected.subsets()), sorted(actual.subsets())) test.assertEqual(len(expected), len(actual)) for item_a in expected: item_b = find(actual, lambda x: x.id == item_a.id and \ x.subset == item_a.subset) test.assertFalse(item_b is None, item_a.id) test.assertEqual(item_a.attributes, item_b.attributes) if require_images or \ item_a.has_image and item_a.image.has_data and \ item_b.has_image and item_b.image.has_data: test.assertEqual(item_a.image, item_b.image, item_a.id) test.assertEqual(len(item_a.annotations), len(item_b.annotations)) for ann_a in item_a.annotations: # We might find few corresponding items, so check them all ann_b_matches = [ x for x in item_b.annotations if x.type == ann_a.type ] test.assertFalse(len(ann_b_matches) == 0, 'ann id: %s' % ann_a.id) ann_b = find( ann_b_matches, lambda x: _compare_annotations( x, ann_a, ignored_attrs=ignored_attrs)) if ann_b is None: test.fail('ann %s, candidates %s' % (ann_a, ann_b_matches)) item_b.annotations.remove(ann_b) # avoid repeats
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 _load_categories(self, label_map_source): if label_map_source == LabelmapType.voc.name: # use the default VOC colormap label_map = make_voc_label_map() elif label_map_source == LabelmapType.source.name and \ AnnotationType.mask not in self._extractor.categories(): # generate colormap for input labels labels = self._extractor.categories() \ .get(AnnotationType.label, LabelCategories()) label_map = OrderedDict( (item.name, [None, [], []]) for item in labels.items) elif label_map_source == LabelmapType.source.name and \ AnnotationType.mask in self._extractor.categories(): # use source colormap labels = self._extractor.categories()[AnnotationType.label] colors = self._extractor.categories()[AnnotationType.mask] label_map = OrderedDict() for idx, item in enumerate(labels.items): color = colors.colormap.get(idx) if color is not None: label_map[item.name] = [color, [], []] elif isinstance(label_map_source, dict): label_map = OrderedDict( sorted(label_map_source.items(), key=lambda e: e[0])) elif isinstance(label_map_source, str) and osp.isfile(label_map_source): label_map = parse_label_map(label_map_source) else: raise Exception("Wrong labelmap specified: '%s', " "expected one of %s or a file path" % \ (label_map_source, ', '.join(t.name for t in LabelmapType))) # 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, 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[0] 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) self._categories = make_voc_categories(label_map) # Update colors with assigned values colormap = self._categories[AnnotationType.mask].colormap for label_id, color in colormap.items(): label_desc = label_map[self._categories[ AnnotationType.label].items[label_id].name] label_desc[0] = color self._label_map = label_map self._label_id_mapping = self._make_label_id_map()
def save_meta_file(path, categories): dataset_meta = {} labels = [label.name for label in categories[AnnotationType.label]] dataset_meta['labels'] = labels if categories.get(AnnotationType.mask): label_map = {} segmentation_colors = [] for i, color in categories[AnnotationType.mask].colormap.items(): if color: segmentation_colors.append( [int(color[0]), int(color[1]), int(color[2])]) label_map[str(i)] = labels[i] dataset_meta['label_map'] = label_map dataset_meta['segmentation_colors'] = segmentation_colors bg_label = find(categories[AnnotationType.mask].colormap.items(), lambda x: x[1] == (0, 0, 0)) if bg_label is not None: dataset_meta['background_label'] = str(bg_label[0]) meta_file = path if osp.isdir(path): meta_file = get_meta_file(path) dump_json_file(meta_file, dataset_meta, indent=True)
def test_can_load_voc_det(self): with TestDir() as test_dir: generated_subsets = generate_dummy_voc(test_dir.path) extractor = VocDetectionExtractor(test_dir.path) self.assertEqual(len(generated_subsets), len(extractor.subsets())) subset_name = 'train' generated_subset = generated_subsets[subset_name] for id_ in generated_subset: parsed_subset = extractor.get_subset(subset_name) self.assertEqual(len(generated_subset), len(parsed_subset)) item = find(parsed_subset, lambda x: x.id == id_) self.assertFalse(item is None) obj1 = find(item.annotations, lambda x: x.type == AnnotationType.bbox and \ get_label(extractor, x.label) == VOC.VocLabel(1).name) self.assertFalse(obj1 is None) self.assertListEqual([1, 2, 2, 2], obj1.get_bbox()) self.assertDictEqual( { 'pose': VOC.VocPose(1).name, 'truncated': True, 'difficult': False, }, obj1.attributes) obj2 = find(item.annotations, lambda x: x.type == AnnotationType.bbox and \ get_label(extractor, x.label) == VOC.VocLabel.person.name) self.assertFalse(obj2 is None) self.assertListEqual([4, 5, 2, 2], obj2.get_bbox()) self.assertEqual(2, len(item.annotations)) subset_name = 'test' generated_subset = generated_subsets[subset_name] for id_ in generated_subset: parsed_subset = extractor.get_subset(subset_name) self.assertEqual(len(generated_subset), len(parsed_subset)) item = find(parsed_subset, lambda x: x.id == id_) self.assertFalse(item is None) self.assertEqual(0, len(item.annotations))
def compare_datasets_3d(test, expected: IDataset, actual: IDataset, ignored_attrs: Union[None, Literal['*'], Collection[str]] = None, require_point_cloud: bool = False): compare_categories(test, expected.categories(), actual.categories()) if actual.subsets(): test.assertEqual(sorted(expected.subsets()), sorted(actual.subsets())) test.assertEqual(len(expected), len(actual)) for item_a in expected: item_b = find(actual, lambda x: x.id == item_a.id) test.assertFalse(item_b is None, item_a.id) if ignored_attrs and ignored_attrs != IGNORE_ALL: test.assertEqual( item_a.attributes, filter_dict(item_b.attributes, exclude_keys=ignored_attrs), item_a.id) elif not ignored_attrs: test.assertEqual(item_a.attributes, item_b.attributes, item_a.id) if (require_point_cloud and item_a.has_point_cloud) or \ (item_a.has_point_cloud and item_b.has_point_cloud): test.assertEqual(item_a.point_cloud, item_b.point_cloud, item_a.id) test.assertEqual(set(img.path for img in item_a.related_images), set(img.path for img in item_b.related_images), item_a.id) test.assertEqual(len(item_a.annotations), len(item_b.annotations)) for ann_a in item_a.annotations: # We might find few corresponding items, so check them all ann_b_matches = [ x for x in item_b.annotations if x.type == ann_a.type ] test.assertFalse(len(ann_b_matches) == 0, 'ann id: %s' % ann_a.id) ann_b = find( ann_b_matches, lambda x: _compare_annotations( x, ann_a, ignored_attrs=ignored_attrs)) if ann_b is None: test.fail('ann %s, candidates %s' % (ann_a, ann_b_matches)) item_b.annotations.remove(ann_b) # avoid repeats
def test_can_load_voc_action(self): with TestDir() as test_dir: generated_subsets = generate_dummy_voc(test_dir.path) extractor = VocActionExtractor(test_dir.path) self.assertEqual(len(generated_subsets), len(extractor.subsets())) subset_name = 'train' generated_subset = generated_subsets[subset_name] for id_ in generated_subset: parsed_subset = extractor.get_subset(subset_name) self.assertEqual(len(generated_subset), len(parsed_subset)) item = find(parsed_subset, lambda x: x.id == id_) self.assertFalse(item is None) obj2 = find(item.annotations, lambda x: x.type == AnnotationType.bbox and \ get_label(extractor, x.label) == VOC.VocLabel.person.name) self.assertFalse(obj2 is None) self.assertListEqual([4, 5, 2, 2], obj2.get_bbox()) count = 1 for action in VOC.VocAction: if action.value % 2 == 1: count += 1 ann = find(item.annotations, lambda x: x.type == AnnotationType.label and \ get_label(extractor, x.label) == action.name) self.assertFalse(ann is None) self.assertTrue(obj2.id == ann.group) self.assertEqual(count, len(item.annotations)) subset_name = 'test' generated_subset = generated_subsets[subset_name] for id_ in generated_subset: parsed_subset = extractor.get_subset(subset_name) self.assertEqual(len(generated_subset), len(parsed_subset)) item = find(parsed_subset, lambda x: x.id == id_) self.assertFalse(item is None) self.assertEqual(0, len(item.annotations))
def convert_instance(self, instance, item): points_ann = find(item.annotations, lambda x: \ x.type == AnnotationType.points and x.group == instance[0].group) if not points_ann: return None elem = super().convert_instance(instance, item) elem.update(self.convert_points_object(points_ann)) return elem
def compare_datasets(test, expected, actual): compare_categories(test, expected.categories(), actual.categories()) test.assertEqual(sorted(expected.subsets()), sorted(actual.subsets())) test.assertEqual(len(expected), len(actual)) for item_a in expected: item_b = find(actual, lambda x: x.id == item_a.id and \ x.subset == item_a.subset) test.assertFalse(item_b is None, item_a.id) test.assertEqual(len(item_a.annotations), len(item_b.annotations)) for ann_a in item_a.annotations: # We might find few corresponding items, so check them all ann_b_matches = [x for x in item_b.annotations if x.id == ann_a.id and \ x.type == ann_a.type and x.group == ann_a.group] test.assertFalse(len(ann_b_matches) == 0, 'ann id: %s' % ann_a.id) ann_b = find(ann_b_matches, lambda x: x == ann_a) test.assertEqual(ann_a, ann_b, 'ann: %s' % ann_to_str(ann_a)) item_b.annotations.remove(ann_b) # avoid repeats
def test_can_load_voc_segm(self): with TestDir() as test_dir: generated_subsets = generate_dummy_voc(test_dir.path) extractor = VocSegmentationExtractor(test_dir.path) self.assertEqual(len(generated_subsets), len(extractor.subsets())) subset_name = 'train' generated_subset = generated_subsets[subset_name] for id_ in generated_subset: parsed_subset = extractor.get_subset(subset_name) self.assertEqual(len(generated_subset), len(parsed_subset)) item = find(parsed_subset, lambda x: x.id == id_) self.assertFalse(item is None) cls_mask = find(item.annotations, lambda x: x.type == AnnotationType.mask and \ x.attributes.get('class') == True) self.assertFalse(cls_mask is None) self.assertFalse(cls_mask.image is None) inst_mask = find(item.annotations, lambda x: x.type == AnnotationType.mask and \ x.attributes.get('instances') == True) self.assertFalse(inst_mask is None) self.assertFalse(inst_mask.image is None) self.assertEqual(2, len(item.annotations)) subset_name = 'test' generated_subset = generated_subsets[subset_name] for id_ in generated_subset: parsed_subset = extractor.get_subset(subset_name) self.assertEqual(len(generated_subset), len(parsed_subset)) item = find(parsed_subset, lambda x: x.id == id_) self.assertFalse(item is None) self.assertEqual(0, len(item.annotations))
def test_can_load_voc_layout(self): with TestDir() as test_dir: generated_subsets = generate_dummy_voc(test_dir.path) extractor = VocLayoutExtractor(test_dir.path) self.assertEqual(len(generated_subsets), len(extractor.subsets())) subset_name = 'train' generated_subset = generated_subsets[subset_name] for id_ in generated_subset: parsed_subset = extractor.get_subset(subset_name) self.assertEqual(len(generated_subset), len(parsed_subset)) item = find(parsed_subset, lambda x: x.id == id_) self.assertFalse(item is None) obj2 = find(item.annotations, lambda x: x.type == AnnotationType.bbox and \ get_label(extractor, x.label) == VOC.VocLabel.person.name) self.assertFalse(obj2 is None) self.assertListEqual([4, 5, 2, 2], obj2.get_bbox()) obj2head = find(item.annotations, lambda x: x.type == AnnotationType.bbox and \ get_label(extractor, x.label) == VOC.VocBodyPart(1).name) self.assertTrue(obj2.id == obj2head.group) self.assertListEqual([5, 6, 2, 2], obj2head.get_bbox()) self.assertEqual(2, len(item.annotations)) subset_name = 'test' generated_subset = generated_subsets[subset_name] for id_ in generated_subset: parsed_subset = extractor.get_subset(subset_name) self.assertEqual(len(generated_subset), len(parsed_subset)) item = find(parsed_subset, lambda x: x.id == id_) self.assertFalse(item is None) self.assertEqual(0, len(item.annotations))
def test_can_load_voc_cls(self): with TestDir() as test_dir: generated_subsets = generate_dummy_voc(test_dir.path) extractor = VocClassificationExtractor(test_dir.path) self.assertEqual(len(generated_subsets), len(extractor.subsets())) subset_name = 'train' generated_subset = generated_subsets[subset_name] for id_ in generated_subset: parsed_subset = extractor.get_subset(subset_name) self.assertEqual(len(generated_subset), len(parsed_subset)) item = find(parsed_subset, lambda x: x.id == id_) self.assertFalse(item is None) count = 0 for label in VOC.VocLabel: if label.value % 2 == 1: count += 1 ann = find(item.annotations, lambda x: x.type == AnnotationType.label and \ x.label == label.value) self.assertFalse(ann is None) self.assertEqual(count, len(item.annotations)) subset_name = 'test' generated_subset = generated_subsets[subset_name] for id_ in generated_subset: parsed_subset = extractor.get_subset(subset_name) self.assertEqual(len(generated_subset), len(parsed_subset)) item = find(parsed_subset, lambda x: x.id == id_) self.assertFalse(item is None) self.assertEqual(0, len(item.annotations))
def find_instance_parts(self, group, img_width, img_height): boxes = [a for a in group if a.type == AnnotationType.bbox] polygons = [a for a in group if a.type == AnnotationType.polygon] masks = [a for a in group if a.type == AnnotationType.mask] anns = boxes + polygons + masks leader = anno_tools.find_group_leader(anns) bbox = anno_tools.max_bbox(anns) mask = None polygons = [p.points for p in polygons] if self._context._segmentation_mode == SegmentationMode.guess: use_masks = True == leader.attributes.get( 'is_crowd', find(masks, lambda x: x.label == leader.label) is not None) elif self._context._segmentation_mode == SegmentationMode.polygons: use_masks = False elif self._context._segmentation_mode == SegmentationMode.mask: use_masks = True else: raise NotImplementedError("Unexpected segmentation mode '%s'" % \ self._context._segmentation_mode) if use_masks: if polygons: mask = mask_tools.rles_to_mask(polygons, img_width, img_height) if masks: masks = (m.image for m in masks) if mask is not None: masks += chain(masks, [mask]) mask = mask_tools.merge_masks(masks) if mask is not None: mask = mask_tools.mask_to_rle(mask) polygons = [] else: if masks: mask = mask_tools.merge_masks(m.image for m in masks) polygons += mask_tools.mask_to_polygons(mask) mask = None return [leader, polygons, mask, bbox]
def save_annotations(self, item): for ann in item.annotations: if ann.type != AnnotationType.points: continue elem = { 'id': self._get_ann_id(ann), 'image_id': _cast(item.id, int, 0), 'category_id': _cast(ann.label, int, -1) + 1, } if 'score' in ann.attributes: elem['score'] = float(ann.attributes['score']) keypoints = [] points = ann.get_points() visibility = ann.visibility for index in range(0, len(points), 2): kp = points[index:index + 2] state = visibility[index // 2].value keypoints.extend([*kp, state]) num_visible = len([v for v in visibility \ if v == PointsObject.Visibility.visible]) bbox = find(item.annotations, lambda x: \ x.group == ann.group and \ x.type == AnnotationType.bbox and x.label == ann.label) if bbox is None: bbox = BboxObject(*ann.get_bbox()) elem.update({ 'segmentation': bbox.get_polygon(), 'area': bbox.area(), 'bbox': bbox.get_bbox(), 'iscrowd': 0, 'keypoints': keypoints, 'num_keypoints': num_visible, }) self.annotations.append(elem)
def write_meta_file(path, label_map): # Uses custom format with extra fields dataset_meta = {} labels = [] labels_dict = {} segmentation_colors = [] parts = {} actions = {} for i, (label_name, label_desc) in enumerate(label_map.items()): labels.append(label_name) if label_desc[0]: labels_dict[str(i)] = label_name segmentation_colors.append([ int(label_desc[0][0]), int(label_desc[0][1]), int(label_desc[0][2]) ]) parts[str(i)] = label_desc[1] actions[str(i)] = label_desc[2] dataset_meta['labels'] = labels if any(segmentation_colors): dataset_meta['label_map'] = labels_dict dataset_meta['segmentation_colors'] = segmentation_colors bg_label = find(label_map.items(), lambda x: x[1] == (0, 0, 0)) if bg_label is not None: dataset_meta['background_label'] = str(bg_label[0]) if any(parts): dataset_meta['parts'] = parts if any(actions): dataset_meta['actions'] = actions dump_json_file(get_meta_file(path), dataset_meta)
def merge_annotations(self, sources): self._make_mergers(sources) clusters = self._match_annotations(sources) joined_clusters = sum(clusters.values(), []) group_map = self._find_cluster_groups(joined_clusters) annotations = [] for t, clusters in clusters.items(): for cluster in clusters: self._check_cluster_sources(cluster) merged_clusters = self._merge_clusters(t, clusters) for merged_ann, cluster in zip(merged_clusters, clusters): attributes = self._find_cluster_attrs(cluster, merged_ann) attributes = { k: v for k, v in attributes.items() if k not in self.conf.ignored_attributes } attributes.update(merged_ann.attributes) merged_ann.attributes = attributes new_group_id = find(enumerate(group_map), lambda e: id(cluster) in e[1][0]) if new_group_id is None: new_group_id = 0 else: new_group_id = new_group_id[0] + 1 merged_ann.group = new_group_id if self.conf.close_distance: self._check_annotation_distance(t, merged_clusters) annotations += merged_clusters if self.conf.groups: self._check_groups(annotations) return annotations
def __call__(self, path): from datumaro.components.project import Project # cyclic import project = Project() for task_name, extractor_type, task_dir in self._TASKS: task_dir = osp.join(path, task_dir) if not osp.isdir(task_dir): continue dir_items = os.listdir(task_dir) if not find(dir_items, lambda x: x == task_name): continue project.add_source(task_name, { 'url': task_dir, 'format': extractor_type, }) if len(project.config.sources) == 0: raise Exception("Failed to find 'voc_results' dataset at '%s'" % \ path) return project
def _load_categories(self, label_map_source=None): if label_map_source == LabelmapType.voc.name: # use the default VOC colormap label_map = make_voc_label_map() elif label_map_source == LabelmapType.source.name and \ AnnotationType.mask not in self._extractor.categories(): # generate colormap for input labels labels = self._extractor.categories() \ .get(AnnotationType.label, LabelCategories()) label_map = OrderedDict() label_map['background'] = [None, [], []] for item in labels.items: label_map[item.name] = [None, [], []] elif label_map_source == LabelmapType.source.name and \ AnnotationType.mask in self._extractor.categories(): # use source colormap labels = self._extractor.categories()[AnnotationType.label] colors = self._extractor.categories()[AnnotationType.mask] label_map = OrderedDict() has_black = False for idx, item in enumerate(labels.items): color = colors.colormap.get(idx) if idx is not None: if color == (0, 0, 0): has_black = True label_map[item.name] = [color, [], []] if not has_black and 'background' not in label_map: label_map['background'] = [(0, 0, 0), [], []] label_map.move_to_end('background', last=False) elif label_map_source in [LabelmapType.guess.name, None]: # generate colormap for union of VOC and input dataset labels label_map = make_voc_label_map() rebuild_colormap = False source_labels = self._extractor.categories() \ .get(AnnotationType.label, LabelCategories()) for label in source_labels.items: if label.name not in label_map: rebuild_colormap = True if label.attributes or label.name not in label_map: label_map[label.name] = [None, [], label.attributes] if rebuild_colormap: for item in label_map.values(): item[0] = None elif isinstance(label_map_source, dict): label_map = label_map_source elif isinstance(label_map_source, str) and osp.isfile(label_map_source): label_map = parse_label_map(label_map_source) has_black = find( label_map.items(), lambda e: e[0] == 'background' or e[1][0] == (0, 0, 0)) if not has_black and 'background' not in label_map: label_map['background'] = [(0, 0, 0), [], []] label_map.move_to_end('background', last=False) else: raise Exception("Wrong labelmap specified, " "expected one of %s or a file path" % \ ', '.join(t.name for t in LabelmapType)) self._categories = make_voc_categories(label_map) self._label_map = label_map colormap = self._categories[AnnotationType.mask].colormap for label_id, color in colormap.items(): label_desc = label_map[self._categories[ AnnotationType.label].items[label_id].name] label_desc[0] = color self._label_id_mapping = self._make_label_id_map()
def save_annotations(self, item): annotations = item.annotations.copy() while len(annotations) != 0: ann = annotations.pop() if ann.type == AnnotationType.bbox and ann.label is not None: pass elif ann.type == AnnotationType.polygon and ann.label is not None: pass elif ann.type == AnnotationType.mask and ann.label is not None: pass else: continue bbox = None segmentation = None if ann.type == AnnotationType.bbox: is_crowd = ann.attributes.get('is_crowd', False) bbox = ann.get_bbox() elif ann.type == AnnotationType.polygon: is_crowd = ann.attributes.get('is_crowd', False) elif ann.type == AnnotationType.mask: is_crowd = ann.attributes.get('is_crowd', True) if is_crowd: segmentation = ann area = None # If ann in a group, try to find corresponding annotations in # this group, otherwise try to infer them. if bbox is None and ann.group is not None: bbox = find(annotations, lambda x: \ x.group == ann.group and \ x.type == AnnotationType.bbox and \ x.label == ann.label) if bbox is not None: bbox = bbox.get_bbox() if is_crowd: # is_crowd=True means there should be a mask if segmentation is None and ann.group is not None: segmentation = find(annotations, lambda x: \ x.group == ann.group and \ x.type == AnnotationType.mask and \ x.label == ann.label) if segmentation is not None: binary_mask = np.array(segmentation.image, dtype=np.bool) binary_mask = np.asfortranarray(binary_mask, dtype=np.uint8) segmentation = mask_utils.encode(binary_mask) area = mask_utils.area(segmentation) segmentation = mask_tools.convert_mask_to_rle(binary_mask) else: # is_crowd=False means there are some polygons polygons = [] if ann.type == AnnotationType.polygon: polygons = [ann] if ann.group is not None: # A single object can consist of several polygons polygons += [p for p in annotations if p.group == ann.group and \ p.type == AnnotationType.polygon and \ p.label == ann.label] if polygons: segmentation = [p.get_points() for p in polygons] h, w, _ = item.image.shape rles = mask_utils.frPyObjects(segmentation, h, w) rle = mask_utils.merge(rles) area = mask_utils.area(rle) if ann.group is not None: # Mark the group as visited to prevent repeats for a in annotations[:]: if a.group == ann.group: annotations.remove(a) if segmentation is None: is_crowd = False segmentation = [ann.get_polygon()] area = ann.area() if bbox is None: bbox = ann.get_bbox() elem = { 'id': self._get_ann_id(ann), 'image_id': _cast(item.id, int, 0), 'category_id': _cast(ann.label, int, -1) + 1, 'segmentation': segmentation, 'area': float(area), 'bbox': bbox, 'iscrowd': int(is_crowd), } if 'score' in ann.attributes: elem['score'] = float(ann.attributes['score']) self.annotations.append(elem)
def find_stage(self, stage): if stage == 'root': return self.root elif stage == 'head': return self.head return find(self.stages, lambda x: x.name == stage or x == stage)
def save_subsets(self): subsets = self._extractor.subsets() if len(subsets) == 0: subsets = [ None ] for subset_name in subsets: if subset_name: subset = self._extractor.get_subset(subset_name) else: subset_name = DEFAULT_SUBSET_NAME subset = self._extractor class_lists = OrderedDict() clsdet_list = OrderedDict() action_list = OrderedDict() layout_list = OrderedDict() segm_list = OrderedDict() for item in subset: item_id = str(item.id) if self._save_images: data = item.image if data is not None: save_image(osp.join(self._images_dir, str(item_id) + VocPath.IMAGE_EXT), data) labels = [] bboxes = [] masks = [] for a in item.annotations: if a.type == AnnotationType.label: labels.append(a) elif a.type == AnnotationType.bbox: bboxes.append(a) elif a.type == AnnotationType.mask: masks.append(a) if len(bboxes) != 0: root_elem = ET.Element('annotation') if '_' in item_id: folder = item_id[ : item_id.find('_')] else: folder = '' ET.SubElement(root_elem, 'folder').text = folder ET.SubElement(root_elem, 'filename').text = \ item_id + VocPath.IMAGE_EXT if item.has_image: h, w, c = item.image.shape size_elem = ET.SubElement(root_elem, 'size') ET.SubElement(size_elem, 'width').text = str(w) ET.SubElement(size_elem, 'height').text = str(h) ET.SubElement(size_elem, 'depth').text = str(c) item_segmented = 0 < len(masks) if item_segmented: ET.SubElement(root_elem, 'segmented').text = '1' objects_with_parts = [] objects_with_actions = defaultdict(dict) main_bboxes = [] layout_bboxes = [] for bbox in bboxes: label = self.get_label(bbox.label) if label in self._LABELS: main_bboxes.append(bbox) elif label in self._BODY_PARTS: layout_bboxes.append(bbox) for new_obj_id, obj in enumerate(main_bboxes): attr = obj.attributes obj_elem = ET.SubElement(root_elem, 'object') ET.SubElement(obj_elem, 'name').text = self.get_label(obj.label) pose = attr.get('pose') if pose is not None: ET.SubElement(obj_elem, 'pose').text = VocPose[pose].name truncated = attr.get('truncated') if truncated is not None: ET.SubElement(obj_elem, 'truncated').text = '%d' % truncated difficult = attr.get('difficult') if difficult is not None: ET.SubElement(obj_elem, 'difficult').text = '%d' % difficult bbox = obj.get_bbox() if bbox is not None: _write_xml_bbox(bbox, obj_elem) for part in VocBodyPart: part_bbox = find(layout_bboxes, lambda x: \ obj.id == x.group and \ self.get_label(x.label) == part.name) if part_bbox is not None: part_elem = ET.SubElement(obj_elem, 'part') ET.SubElement(part_elem, 'name').text = part.name _write_xml_bbox(part_bbox.get_bbox(), part_elem) objects_with_parts.append(new_obj_id) actions = [x for x in labels if obj.id == x.group and \ self.get_label(x.label) in self._ACTIONS] if len(actions) != 0: actions_elem = ET.SubElement(obj_elem, 'actions') for action in VocAction: presented = find(actions, lambda x: \ self.get_label(x.label) == action.name) is not None ET.SubElement(actions_elem, action.name).text = \ '%d' % presented objects_with_actions[new_obj_id][action] = presented if self._task in [None, VocTask.detection, VocTask.person_layout, VocTask.action_classification]: with open(osp.join(self._ann_dir, item_id + '.xml'), 'w') as f: f.write(ET.tostring(root_elem, encoding='unicode', pretty_print=True)) clsdet_list[item_id] = True layout_list[item_id] = objects_with_parts action_list[item_id] = objects_with_actions for label_obj in labels: label = self.get_label(label_obj.label) if label not in self._LABELS: continue class_list = class_lists.get(item_id, set()) class_list.add(label_obj.label) class_lists[item_id] = class_list clsdet_list[item_id] = True for mask_obj in masks: if mask_obj.attributes.get('class') == True: self.save_segm(osp.join(self._segm_dir, item_id + VocPath.SEGM_EXT), mask_obj, self._mask_categories.colormap) if mask_obj.attributes.get('instances') == True: self.save_segm(osp.join(self._inst_dir, item_id + VocPath.SEGM_EXT), mask_obj, VocInstColormap) segm_list[item_id] = True if len(item.annotations) == 0: clsdet_list[item_id] = None layout_list[item_id] = None action_list[item_id] = None segm_list[item_id] = None if self._task in [None, VocTask.classification, VocTask.detection, VocTask.action_classification, VocTask.person_layout]: self.save_clsdet_lists(subset_name, clsdet_list) if self._task in [None, VocTask.classification]: self.save_class_lists(subset_name, class_lists) if self._task in [None, VocTask.action_classification]: self.save_action_lists(subset_name, action_list) if self._task in [None, VocTask.person_layout]: self.save_layout_lists(subset_name, layout_list) if self._task in [None, VocTask.segmentation]: self.save_segm_lists(subset_name, segm_list)