awacke1's picture
Upload 11 files
0e371d1
import xml.etree.ElementTree as ET
import os
from glob import glob
XML_PATH = './dataset/xml'
CLASSES_PATH = './class_names/classes.txt'
TXT_PATH = './dataset/txt/anno.txt'
'''loads the classes'''
def get_classes(classes_path):
with open(classes_path) as f:
class_names = f.readlines()
class_names = [c.strip() for c in class_names]
return class_names
classes = get_classes(CLASSES_PATH)
assert len(classes) > 0, 'no class names detected!'
print(f'num classes: {len(classes)}')
# output file
list_file = open(TXT_PATH, 'w')
for path in glob(os.path.join(XML_PATH, '*.xml')):
in_file = open(path)
# Parse .xml file
tree = ET.parse(in_file)
root = tree.getroot()
# Write object information to .txt file
file_name = root.find('filename').text
print(file_name)
list_file.write(file_name)
for obj in root.iter('object'):
cls = obj.find('name').text
cls_id = classes.index(cls)
xmlbox = obj.find('bndbox')
b = (int(xmlbox.find('xmin').text), int(xmlbox.find('ymin').text), int(xmlbox.find('xmax').text), int(xmlbox.find('ymax').text))
list_file.write(" " + ",".join([str(a) for a in b]) + ',' + str(cls_id))
list_file.write('\n')
list_file.close()