xml_to_csv.py
3.34 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
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
"""
Usage:
# Create train data:
python xml_to_csv.py -i [PATH_TO_IMAGES_FOLDER]/train -o [PATH_TO_ANNOTATIONS_FOLDER]/train_labels.csv
# Create test data:
python xml_to_csv.py -i [PATH_TO_IMAGES_FOLDER]/test -o [PATH_TO_ANNOTATIONS_FOLDER]/test_labels.csv
"""
import os
import glob
import pandas as pd
import argparse
import xml.etree.ElementTree as ET
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.
Parameters:
----------
path : {str}
The path containing the .xml files
Returns
-------
Pandas DataFrame
The produced dataframe
"""
classes_names = []
xml_list = []
for xml_file in glob.glob(path + "/*.xml"):
tree = ET.parse(xml_file)
root = tree.getroot()
for member in root.findall("object"):
classes_names.append(member[0].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),
)
xml_list.append(value)
column_name = [
"filename",
"width",
"height",
"class",
"xmin",
"ymin",
"xmax",
"ymax",
]
xml_df = pd.DataFrame(xml_list, columns=column_name)
classes_names = list(set(classes_names))
classes_names.sort()
return xml_df, classes_names
def main():
# Initiate argument parser
parser = argparse.ArgumentParser(
description="Sample TensorFlow XML-to-CSV converter"
)
parser.add_argument(
"-i",
"--inputDir",
help="Path to the folder where the input .xml files are stored",
type=str,
)
parser.add_argument(
"-o", "--outputFile", help="Name of output .csv file (including path)", type=str
)
parser.add_argument(
"-l",
"--labelMapDir",
help="Directory path to save label_map.pbtxt file is specified.",
type=str,
default="",
)
args = parser.parse_args()
if args.inputDir is None:
args.inputDir = os.getcwd()
if args.outputFile is None:
args.outputFile = args.inputDir + "/labels.csv"
assert os.path.isdir(args.inputDir)
os.makedirs(os.path.dirname(args.outputFile), exist_ok=True)
xml_df, classes_names = xml_to_csv(args.inputDir)
xml_df.to_csv(args.outputFile, index=None)
print("Successfully converted xml to csv.")
if args.labelMapDir:
os.makedirs(args.labelMapDir, exist_ok=True)
label_map_path = os.path.join(args.labelMapDir, "label_map.pbtxt")
print("Generate `{}`".format(label_map_path))
# Create the `label_map.pbtxt` file
pbtxt_content = ""
for i, class_name in enumerate(classes_names):
pbtxt_content = (
pbtxt_content
+ "item {{\n id: {0}\n name: '{1}'\n}}\n\n".format(
i + 1, class_name
)
)
pbtxt_content = pbtxt_content.strip()
with open(label_map_path, "w") as f:
f.write(pbtxt_content)
if __name__ == "__main__":
main()