Spaces:
Runtime error
Runtime error
| 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() | |