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)