-
Notifications
You must be signed in to change notification settings - Fork 4
/
coco_to_xml.py
120 lines (92 loc) · 3.43 KB
/
coco_to_xml.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
import argparse
parser = argparse.ArgumentParser(description='Coco Json to Pascal VOC XML Converter.')
parser.add_argument('--coco_json', required=True, help='Target json file to convert to csv.')
parser.add_argument('--coco_folder', required=False, default='', help='Target folder to find images.')
parser.add_argument('--save_xml', required=True, help='The folder to save annotations xmls.')
parser.add_argument('--database_name', required=False, default='', help='The name of database.')
parser.add_argument('--no_skip_background', dest='skip_background', action='store_false', help='Do not skip \'background\' category.')
parser.set_defaults(skip_background=True)
args = parser.parse_args()
from pycocotools.coco import COCO
from PIL import Image
from pathlib import Path
import os
def write_to_xml(image_name, bboxes, image_folder_name, data_folder, save_folder, database_name):
with Image.open(os.path.join(data_folder, image_name)) as img:
width, height = img.size
if img.mode == 'YCbCr':
depth = 3
else:
depth = len(img.mode)
objects = ''
for bbox in bboxes:
objects = objects + '''
<object>
<name>{category_name}</name>
<pose>Unspecified</pose>
<truncated>0</truncated>
<difficult>0</difficult>
<bndbox>
<xmin>{xmin}</xmin>
<ymin>{ymin}</ymin>
<xmax>{xmax}</xmax>
<ymax>{ymax}</ymax>
</bndbox>
</object>'''.format(
category_name = bbox[0],
xmin = bbox[1],
ymin = bbox[2],
xmax = bbox[3],
ymax = bbox[4]
)
xml = '''<annotation>
<folder>{image_folder_name}</folder>
<filename>{image_name}</filename>
<source>
<database>{database_name}</database>
</source>
<size>
<width>{width}</width>
<height>{height}</height>
<depth>{depth}</depth>
</size>
<segmented>0</segmented>{objects}
</annotation>'''.format(
image_folder_name = image_folder_name,
image_name = image_name,
database_name = database_name,
width = width,
height = height,
depth = depth,
objects = objects
)
anno_path = os.path.join(save_folder, os.path.splitext(image_name)[0] + '.xml')
with open(anno_path, 'w') as file:
file.write(xml)
Path(args.save_xml).mkdir(parents=True, exist_ok=True)
image_folder_name = os.path.basename(os.path.abspath(args.coco_folder))
coco = COCO(args.coco_json)
imgIds = coco.getImgIds()
if not args.database_name:
args.database_name = image_folder_name
print('Write annotations file...')
total = len(imgIds)
now = 1
for imgId in imgIds:
img = coco.loadImgs(imgId)[0]
if not os.path.isfile(os.path.join(args.coco_folder, img['file_name'])):
print('The image {} is not exist.'.format(img['file_name']))
exit()
anno_list = []
annIds = coco.getAnnIds(imgIds=imgId)
for annId in annIds:
ann = coco.loadAnns(annId)[0]
cat = coco.loadCats(ann['category_id'])[0]
if cat['name'] == 'background' and args.skip_background == True:
continue
box = ann['bbox']
anno_list.append([cat['name'], int(box[0]), int(box[1]), int(box[0]+box[2]), int(box[1]+box[3])])
write_to_xml(img['file_name'], anno_list, image_folder_name, args.coco_folder, args.save_xml, args.database_name)
print('Write xml files ({} / {})'.format(now, total))
now = now + 1
print('Annotations file was written!')