annotation_xml_parser.py
1.19 KB
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
import xml.etree.ElementTree as elemTree
import os
CLASSES = ['dog']
def parseFile(path):
result = ""
annotation = elemTree.parse(path)
filename = annotation.find('filename')
size = annotation.find('size')
width = size.find('width')
height = size.find('height')
result += filename.text + ' ' + width.text + ' ' + height.text
for obj in annotation.findall('./object'):
name = obj.find('name')
box = obj.find('bndbox')
label_index = CLASSES.index(name.text)
xmin = box.find('xmin')
ymin = box.find('ymin')
xmax = box.find('xmax')
ymax = box.find('ymax')
result += ' ' + str(label_index) + ' ' + xmin.text + ' ' + ymin.text + ' ' + xmax.text + ' ' + ymax.text
return result
def parseDir(path):
result = ""
for file in os.listdir(path):
if file.split('.')[-1] != 'xml':
continue
result += parseFile(path+'/'+file) + '\n'
return result
directory = input('Input directory, or 0 to exit: ')
if directory != '0':
result = parseDir(directory)
file = open(directory+'/'+'annotation.txt','w')
file.write(result)
file.close()
print("Complete!")