diff --git a/xml_to_csv.py b/xml_to_csv.py index 03692c4cfa0..a4c24e34d83 100644 --- a/xml_to_csv.py +++ b/xml_to_csv.py @@ -14,6 +14,11 @@ import xml.etree.ElementTree as ET +def confine(value, lower, upper): + return max(min(upper, value), lower) + + + def xml_to_csv(path): """Iterates through all .xml files (generated by labelImg) in a given directory and combines them in a single Pandas datagrame. @@ -32,16 +37,28 @@ def xml_to_csv(path): tree = ET.parse(xml_file) root = tree.getroot() for member in root.findall("object"): - classes_names.append(member[0].text) + classname = member.find("name").text + classes_names.append(classname) + + size = root.find("size") + width = int(size.find("width").text) + height = int(size.find("height").text) + + box = member.find("bndbox") + xmin = int(box.find("xmin").text) + xmax = int(box.find("xmax").text) + ymin = int(box.find("ymin").text) + ymax = int(box.find("ymax").text) + value = ( root.find("filename").text, - int(root.find("size")[0].text), - int(root.find("size")[1].text), - member[0].text, - int(member[4][0].text), - int(member[4][1].text), - int(member[4][2].text), - int(member[4][3].text), + width, + height, + classname, + confine(xmin, 0, width), + confine(ymin, 0, width), + confine(xmax, 0, height), + confine(ymax, 0, height) ) xml_list.append(value) column_name = [