Creating TFRecords

Previous Article – https://wp.me/p6xoZs-3G

TFRecords are special data format which is used to read image data from tensorflow framework. To create TFRecords there are two steps as below,

Step 01 – Convert XML to CSV

From the source code you download direct to research/object_detection; then you can find a python script named xml_to_csv.py.

Change the script as below,
Replace

def main():
    image_path = os.path.join(os.getcwd(), 'annotations')
    xml_df = xml_to_csv(image_path)
    xml_df.to_csv('raccoon_labels.csv', index=None)
    print('Successfully converted xml to csv.')

With

def main():
    for directory in ['train','test']:
        image_path = os.path.join(os.getcwd(), 'images/{}'.format(directory))
        xml_df = xml_to_csv(image_path)
        xml_df.to_csv('data/{}_labels.csv'.format(directory), index=None)
        print('Successfully converted xml to csv.')

mindful on the folder paths and your running directory.

then run as below,

python xml_to_csv.py

Step 02 – Convert CSV to TFRecord

Now from the same location grab the generate_tfrecord.py and modify as below,

# TO-DO replace this with label map
def class_text_to_int(row_label):
    if row_label == 'macncheese':
        return 1
    else:
        None

you can expand the if condition according to no of labels you used to train.

Before run the above script direct to research folder and run below command

python setup.py install

then run as below,

python generate_tfrecord.py --csv_input=../Data_Set/train_labels.csv  --output_path=../Data_Set/train.record --image_dir=../Data_Set/train
python generate_tfrecord.py --csv_input=../Data_Set/test_labels.csv  --output_path=../Data_Set/test.record --image_dir=../Data_Set/test

Next Article – https://wp.me/p6xoZs-3O

Leave a comment