From 9c57a79642ecbc878bca43877ab7a2296317a7e1 Mon Sep 17 00:00:00 2001 From: Dmitry Balabka Date: Sun, 21 Feb 2021 00:36:12 +0200 Subject: [PATCH] Fix problem for XML files which additional tags --- docs/source/scripts/generate_tfrecord.py | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/docs/source/scripts/generate_tfrecord.py b/docs/source/scripts/generate_tfrecord.py index caad456..944b9ec 100644 --- a/docs/source/scripts/generate_tfrecord.py +++ b/docs/source/scripts/generate_tfrecord.py @@ -80,15 +80,19 @@ def xml_to_csv(path): for xml_file in glob.glob(path + '/*.xml'): tree = ET.parse(xml_file) root = tree.getroot() + filename = root.find('filename').text + width = int(root.find('size').find('width').text) + height = int(root.find('size').find('height').text) for member in root.findall('object'): - 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) + bndbox = member.find('bndbox') + value = (filename, + width, + height, + member.find('name').text, + int(bndbox.find('xmin').text), + int(bndbox.find('ymin').text), + int(bndbox.find('xmax').text), + int(bndbox.find('ymax').text), ) xml_list.append(value) column_name = ['filename', 'width', 'height',