Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Python script to create tfrecords from pascal VOC data set format (one class detection) for Object Detection API Tensorflow, where it divides dataset into (90% train.record and 10% test.record)
import os
import io
import glob
import hashlib
import pandas as pd
import xml.etree.ElementTree as ET
import tensorflow as tf
import random
from PIL import Image
from object_detection.utils import dataset_util
'''
this script automatically divides dataset into training and evaluation (10% for evaluation)
this scripts also shuffles the dataset before converting it into tfrecords
if u have different structure of dataset (rather than pascal VOC ) u need to change
the paths and names input directories(images and annotation) and output tfrecords names.
(note: this script can be enhanced to use flags instead of changing parameters on code).
default expected directories tree:
dataset-
-JPEGImages
-Annotations
dataset_to_tfrecord.py
to run this script:
$ python dataset_to_tfrecord.py
'''
def create_example(xml_file):
#process the xml file
tree = ET.parse(xml_file)
root = tree.getroot()
image_name = root.find('filename').text
file_name = image_name.encode('utf8')
size=root.find('size')
width = int(size[0].text)
height = int(size[1].text)
xmin = []
ymin = []
xmax = []
ymax = []
classes = []
classes_text = []
truncated = []
poses = []
difficult_obj = []
for member in root.findall('object'):
classes_text.append('Person'.encode('utf8'))
xmin.append(float(member[4][0].text) / width)
ymin.append(float(member[4][1].text) / height)
xmax.append(float(member[4][2].text) / width)
ymax.append(float(member[4][3].text) / height)
difficult_obj.append(0)
#if you have more than one classes in dataset you can change the next line
#to read the class from the xml file and change the class label into its
#corresponding integer number, u can use next function structure
'''
def class_text_to_int(row_label):
if row_label == 'Person':
return 1
if row_label == 'car':
return 2
and so on.....
'''
classes.append(1) # i wrote 1 because i have only one class(person)
truncated.append(0)
poses.append('Unspecified'.encode('utf8'))
#read corresponding image
full_path = os.path.join('./JPEGImages', '{}'.format(image_name)) #provide the path of images directory
with tf.gfile.GFile(full_path, 'rb') as fid:
encoded_jpg = fid.read()
encoded_jpg_io = io.BytesIO(encoded_jpg)
image = Image.open(encoded_jpg_io)
if image.format != 'JPEG':
raise ValueError('Image format not JPEG')
key = hashlib.sha256(encoded_jpg).hexdigest()
#create TFRecord Example
example = tf.train.Example(features=tf.train.Features(feature={
'image/height': dataset_util.int64_feature(height),
'image/width': dataset_util.int64_feature(width),
'image/filename': dataset_util.bytes_feature(file_name),
'image/source_id': dataset_util.bytes_feature(file_name),
'image/key/sha256': dataset_util.bytes_feature(key.encode('utf8')),
'image/encoded': dataset_util.bytes_feature(encoded_jpg),
'image/format': dataset_util.bytes_feature('jpeg'.encode('utf8')),
'image/object/bbox/xmin': dataset_util.float_list_feature(xmin),
'image/object/bbox/xmax': dataset_util.float_list_feature(xmax),
'image/object/bbox/ymin': dataset_util.float_list_feature(ymin),
'image/object/bbox/ymax': dataset_util.float_list_feature(ymax),
'image/object/class/text': dataset_util.bytes_list_feature(classes_text),
'image/object/class/label': dataset_util.int64_list_feature(classes),
'image/object/difficult': dataset_util.int64_list_feature(difficult_obj),
'image/object/truncated': dataset_util.int64_list_feature(truncated),
'image/object/view': dataset_util.bytes_list_feature(poses),
}))
return example
def main(_):
writer_train = tf.python_io.TFRecordWriter('train.record')
writer_test = tf.python_io.TFRecordWriter('test.record')
#provide the path to annotation xml files directory
filename_list=tf.train.match_filenames_once("./Annotations/*.xml")
init = (tf.global_variables_initializer(), tf.local_variables_initializer())
sess=tf.Session()
sess.run(init)
list=sess.run(filename_list)
random.shuffle(list) #shuffle files list
i=1
tst=0 #to count number of images for evaluation
trn=0 #to count number of images for training
for xml_file in list:
example = create_example(xml_file)
if (i%10)==0: #each 10th file (xml and image) write it for evaluation
writer_test.write(example.SerializeToString())
tst=tst+1
else: #the rest for training
writer_train.write(example.SerializeToString())
trn=trn+1
i=i+1
print(xml_file)
writer_test.close()
writer_train.close()
print('Successfully converted dataset to TFRecord.')
print('training dataset: # ')
print(trn)
print('test dataset: # ')
print(tst)
if __name__ == '__main__':
tf.app.run()
@AdarshMJ

This comment has been minimized.

Copy link

@AdarshMJ AdarshMJ commented Dec 26, 2017

I keep getting this error. Any idea why?

Traceback (most recent call last):
File "dataset_to_tfrecord.py", line 134, in
tf.app.run()
File "/Library/Python/2.7/site-packages/tensorflow/python/platform/app.py", line 48, in run
_sys.exit(main(_sys.argv[:1] + flags_passthrough))
File "dataset_to_tfrecord.py", line 116, in main
example = create_example(xml_file)
File "dataset_to_tfrecord.py", line 51, in create_example
xmin.append(float(member[4][0].text) / width)
File "/System/Library/Frameworks/Python.framework/Versions/2.7/lib/python2.7/xml/etree/ElementTree.py", line 266, in getitem
return self._children[index]
IndexError: list index out of range

@arvinxian

This comment has been minimized.

Copy link

@arvinxian arvinxian commented Apr 23, 2018

I got the same issue, someone pls tell me how to solve it. thx

@mepotts

This comment has been minimized.

Copy link

@mepotts mepotts commented Feb 21, 2019

Not sure if you solved but I ran into the same issue. Seems like I had a slightly different xml than the one this is meant for. It is still pascal VOC but has an added tag in the xml file called occluded. If you change the member[4][0].text to member[5][0].text it should work. Or just look at the xml and make sure it is referencing the right tag to get the bndbox data. @arvinxian @AdarshMJ @saghiralfasly

@MAkC8

This comment has been minimized.

Copy link

@MAkC8 MAkC8 commented Jul 1, 2019

use this

from absl import app, flags, logging
...


if __name__ == '__main__':
    app.run(main)
@stiv-yakovenko

This comment has been minimized.

Copy link

@stiv-yakovenko stiv-yakovenko commented Jul 12, 2019

This crashes for me like this:

tensorflow.python.framework.errors_impl.NotFoundError: ./JPEGImages/0711; No such file or directory

the file is place, but its named 0711.jpg

@heizie

This comment has been minimized.

Copy link

@heizie heizie commented Sep 11, 2019

This crashes for me like this:

tensorflow.python.framework.errors_impl.NotFoundError: ./JPEGImages/0711; No such file or directory

the file is place, but its named 0711.jpg

your .xml file point to ./JPEGImages/0711 instead of ./JPEGImages/0711.jpg

@harisarapakis

This comment has been minimized.

Copy link

@harisarapakis harisarapakis commented Mar 25, 2020

An exception has occurred, use %tb to see the full traceback.

SystemExit

Any solutions??

@monicaglez

This comment has been minimized.

Copy link

@monicaglez monicaglez commented May 15, 2020

I have this error:

TypeError: run() missing 1 required positional argument: 'fetches'

Can anyone help? thanks!

@monicaglez

This comment has been minimized.

Copy link

@monicaglez monicaglez commented May 16, 2020

I have this error:

TypeError: run() missing 1 required positional argument: 'fetches'

Can anyone help? thanks!

It was solved by converting the tuple init to a list:
init = [(tf.global_variables_initializer(), tf.local_variables_initializer()]

But then I get a similar error as @harisarapakis, although the tf_records were successfuly converted:

'An exception has occurred, use %tb to see the full traceback.

SystemExit

%tb

File "C:\Users\monic\AppData\Local\conda\conda\envs\tensorflow\lib\site-packages\absl\app.py", line 251, in _run_main
sys.exit(main(argv))

SystemExit

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.