Skip to content

Instantly share code, notes, and snippets.

@nitroxplunge
Last active March 3, 2021 10:44
Show Gist options
  • Select an option

  • Save nitroxplunge/fdd750465f7ae55afc0d43d37ef65ef5 to your computer and use it in GitHub Desktop.

Select an option

Save nitroxplunge/fdd750465f7ae55afc0d43d37ef65ef5 to your computer and use it in GitHub Desktop.
Convert a .csv file in VOC format to .tfrecord
import tensorflow as tf
import numpy as np
import base64
import csv
import os
from PIL import Image
import io
def int64_feature(value):
return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
def int64_list_feature(value):
return tf.train.Feature(int64_list=tf.train.Int64List(value=value))
def bytes_feature(value):
return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
def bytes_list_feature(value):
return tf.train.Feature(bytes_list=tf.train.BytesList(value=value))
def float_list_feature(value):
return tf.train.Feature(float_list=tf.train.FloatList(value=value))
class_dict = {
1: b'cup', # List of class map Text with byte
2: b'rutensil'
}
def create_tf_example(img_name, img_path, labels):
with tf.gfile.GFile(img_path, 'rb') as fid:
encoded_jpg = fid.read()
image_format = str.encode(img_name.split('.')[1])
xmins = []
xmaxs = []
ymins = []
ymaxs = []
classes_text = []
classes = []
for label in labels:
xmins.append(float(int(label[0][4]) / int(label[0][1])))
xmaxs.append(float(int(label[0][6]) / int(label[0][1])))
ymins.append(float(int(label[0][5]) / int(label[0][2])))
ymaxs.append(float(int(label[0][7]) / int(label[0][2])))
classes_text.append(label[0][3])
class_ = 0
if (label[0][3] == "cup"):
class_ = 1
elif (label[0][3] == "rutensil"):
class_ = 2
classes.append(class_)
tf_example = tf.train.Example(features=tf.train.Features(feature={
'image/height': int64_feature(int(labels[0][0][2])),
'image/width': int64_feature(int(labels[0][0][1])),
'image/filename': bytes_feature(str.encode(labels[0][0][0])),
'image/source_id': bytes_feature(labels[0][0][0]),
'image/encoded': bytes_feature(encoded_jpg),
'image/format': bytes_feature(image_format),
'image/object/bbox/xmin': float_list_feature(xmins),
'image/object/bbox/xmax': float_list_feature(xmaxs),
'image/object/bbox/ymin': float_list_feature(ymins),
'image/object/bbox/ymax': float_list_feature(ymaxs),
'image/object/class/text': bytes_list_feature(classes_text),
'image/object/class/label': int64_list_feature(classes),
}))
return tf_example
def checkHowManyAhead(csv_reader, i, relative):
try:
if csv_reader[i + relative][0][0] == csv_reader[i + relative + 1][0][0]:
return checkHowManyAhead(csv_reader, i, relative + 1)
else:
return relative
except IndexError:
print "End of list"
return 0
def main():
writer = tf.python_io.TFRecordWriter('test.tfrecords')
data_path = 'testandtrainimgs/test'
images = os.listdir(data_path)
label_csv = 'test_labels.csv'
csvee_reader = csv.reader(open(label_csv, 'r'))
csv_reader = zip(csvee_reader)
for i in range(0, len(images)):
labels = []
ahead = checkHowManyAhead(csv_reader, i, 0)
for j in range(0, ahead + 1):
labels.append(csv_reader[i + j])
tf_example = create_tf_example(images[i], data_path + "/" + images[i], labels)
writer.write(tf_example.SerializeToString())
writer.close()
main()
#if __name__ == '__main__':
# tf.app.run()
@phuonganhniie1808
Copy link
Copy Markdown

Excuse me but how to use this script?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment