-
Notifications
You must be signed in to change notification settings - Fork 56
/
anno_coco2voc.py
executable file
·154 lines (138 loc) · 6.36 KB
/
anno_coco2voc.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
import argparse, json
import cytoolz
from lxml import etree, objectify
import os, re
def instance2xml_base(anno):
E = objectify.ElementMaker(annotate=False)
anno_tree = E.annotation(
E.folder('VOC2014_instance/{}'.format(anno['category_id'])),
E.filename(anno['file_name']),
E.source(
E.database('MS COCO 2014'),
E.annotation('MS COCO 2014'),
E.image('Flickr'),
E.url(anno['coco_url'])
),
E.size(
E.width(anno['width']),
E.height(anno['height']),
E.depth(3)
),
E.segmented(0),
)
return anno_tree
def instance2xml_bbox(anno, bbox_type='xyxy'):
"""bbox_type: xyxy (xmin, ymin, xmax, ymax); xywh (xmin, ymin, width, height)"""
assert bbox_type in ['xyxy', 'xywh']
if bbox_type == 'xyxy':
xmin, ymin, w, h = anno['bbox']
xmax = xmin+w
ymax = ymin+h
else:
xmin, ymin, xmax, ymax = anno['bbox']
E = objectify.ElementMaker(annotate=False)
anno_tree = E.object(
E.name(anno['category_id']),
E.bndbox(
E.xmin(xmin),
E.ymin(ymin),
E.xmax(xmax),
E.ymax(ymax)
),
E.difficult(anno['iscrowd'])
)
return anno_tree
def parse_instance(content, outdir):
categories = {d['id']: d['name'] for d in content['categories']}
# merge images and annotations: id in images vs image_id in annotations
merged_info_list = list(map(cytoolz.merge, cytoolz.join('id', content['images'], 'image_id', content['annotations'])))
# convert category id to name
for instance in merged_info_list:
instance['category_id'] = categories[instance['category_id']]
# group by filename to pool all bbox in same file
for name, groups in cytoolz.groupby('file_name', merged_info_list).items():
anno_tree = instance2xml_base(groups[0])
# if one file have multiple different objects, save it in each category sub-directory
filenames = []
for group in groups:
filenames.append(os.path.join(outdir, re.sub(" ", "_", group['category_id']),
os.path.splitext(name)[0] + ".xml"))
anno_tree.append(instance2xml_bbox(group, bbox_type='xyxy'))
for filename in filenames:
etree.ElementTree(anno_tree).write(filename, pretty_print=True)
print("Formating instance xml file {} done!".format(name))
def keypoints2xml_base(anno):
annotation = etree.Element("annotation")
etree.SubElement(annotation, "folder").text = "VOC2014_keypoints"
etree.SubElement(annotation, "filename").text = anno['file_name']
source = etree.SubElement(annotation, "source")
etree.SubElement(source, "database").text = "MS COCO 2014"
etree.SubElement(source, "annotation").text = "MS COCO 2014"
etree.SubElement(source, "image").text = "Flickr"
etree.SubElement(source, "url").text = anno['coco_url']
size = etree.SubElement(annotation, "size")
etree.SubElement(size, "width").text = str(anno["width"])
etree.SubElement(size, "height").text = str(anno["height"])
etree.SubElement(size, "depth").text = '3'
etree.SubElement(annotation, "segmented").text = '0'
return annotation
def keypoints2xml_object(anno, xmltree, keypoints_dict, bbox_type='xyxy'):
assert bbox_type in ['xyxy', 'xywh']
if bbox_type == 'xyxy':
xmin, ymin, w, h = anno['bbox']
xmax = xmin+w
ymax = ymin+h
else:
xmin, ymin, xmax, ymax = anno['bbox']
key_object = etree.SubElement(xmltree, "object")
etree.SubElement(key_object, "name").text = anno['category_id']
bndbox = etree.SubElement(key_object, "bndbox")
etree.SubElement(bndbox, "xmin").text = str(xmin)
etree.SubElement(bndbox, "ymin").text = str(ymin)
etree.SubElement(bndbox, "xmax").text = str(xmax)
etree.SubElement(bndbox, "ymax").text = str(ymax)
etree.SubElement(key_object, "difficult").text = '0'
keypoints = etree.SubElement(key_object, "keypoints")
for i in range(0, len(keypoints_dict)):
keypoint = etree.SubElement(keypoints, keypoints_dict[i+1])
etree.SubElement(keypoint, "x").text = str(anno['keypoints'][i*3])
etree.SubElement(keypoint, "y").text = str(anno['keypoints'][i*3+1])
etree.SubElement(keypoint, "v").text = str(anno['keypoints'][i*3+2])
return xmltree
def parse_keypoints(content, outdir):
keypoints = dict(zip(range(1, len(content['categories'][0]['keypoints'])+1), content['categories'][0]['keypoints']))
# merge images and annotations: id in images vs image_id in annotations
merged_info_list = map(cytoolz.merge, cytoolz.join('id', content['images'], 'image_id', content['annotations']))
# convert category name to person
for keypoint in merged_info_list:
keypoint['category_id'] = "person"
# group by filename to pool all bbox and keypoint in same file
for name, groups in cytoolz.groupby('file_name', merged_info_list).items():
filename = os.path.join(outdir, os.path.splitext(name)[0]+".xml")
anno_tree = keypoints2xml_base(groups[0])
for group in groups:
anno_tree = keypoints2xml_object(group, anno_tree, keypoints, bbox_type="xyxy")
doc = etree.ElementTree(anno_tree)
doc.write(open(filename, "w"), pretty_print=True)
print("Formating keypoints xml file {} done!".format(name))
def main(args):
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
content = json.load(open(args.anno_file, 'r'))
if args.type == 'instance':
# make subdirectories
sub_dirs = [re.sub(" ", "_", cate['name']) for cate in content['categories']]
for sub_dir in sub_dirs:
sub_dir = os.path.join(args.output_dir, str(sub_dir))
if not os.path.exists(sub_dir):
os.makedirs(sub_dir)
parse_instance(content, args.output_dir)
elif args.type == 'keypoint':
parse_keypoints(content, args.output_dir)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--anno_file", help="annotation file for object instance/keypoint")
parser.add_argument("--type", type=str, help="object instance or keypoint", choices=['instance', 'keypoint'])
parser.add_argument("--output_dir", help="output directory for voc annotation xml file")
args = parser.parse_args()
main(args)