Skip to content

Commit 292865b

Browse files
dbalabkasglvladi
authored andcommitted
Fix problem for XML files which additional tags
1 parent d789f03 commit 292865b

File tree

1 file changed

+12
-8
lines changed

1 file changed

+12
-8
lines changed

docs/source/scripts/generate_tfrecord.py

+12-8
Original file line numberDiff line numberDiff line change
@@ -80,15 +80,19 @@ def xml_to_csv(path):
8080
for xml_file in glob.glob(path + '/*.xml'):
8181
tree = ET.parse(xml_file)
8282
root = tree.getroot()
83+
filename = root.find('filename').text
84+
width = int(root.find('size').find('width').text)
85+
height = int(root.find('size').find('height').text)
8386
for member in root.findall('object'):
84-
value = (root.find('filename').text,
85-
int(root.find('size')[0].text),
86-
int(root.find('size')[1].text),
87-
member[0].text,
88-
int(member[4][0].text),
89-
int(member[4][1].text),
90-
int(member[4][2].text),
91-
int(member[4][3].text)
87+
bndbox = member.find('bndbox')
88+
value = (filename,
89+
width,
90+
height,
91+
member.find('name').text,
92+
int(bndbox.find('xmin').text),
93+
int(bndbox.find('ymin').text),
94+
int(bndbox.find('xmax').text),
95+
int(bndbox.find('ymax').text),
9296
)
9397
xml_list.append(value)
9498
column_name = ['filename', 'width', 'height',

0 commit comments

Comments
 (0)