#coding=utf-8
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import datetime
import os
root=os.path.abspath(os.path.dirname(__file__))
os.chdir(root)
image_train_path="train"
#label_train_path="train.txt"
tfRecord_train="data/train.tfrecords"
image_test_path="test"
#label_test_path="test.txt"
tfRecord_test="data/test.tfrecords"
resize_height=32
resize_width=32
data_path="data"
def write_record(tf_record_name,image_path,label_data):
writer = tf.io.TFRecordWriter(tf_record_name)
pic_id=0
for con in label_data:
img_path=image_path+"/"+con
img= Image.open(img_path)
img_str=img.tobytes()
label=label_data[con]
feature={
'img_str':tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_str])),
'label':tf.train.Feature(bytes_list=tf.train.BytesList(value=[label.encode()]))
}
example=tf.train.Example(features=tf.train.Features(feature=feature))
writer.write(example.SerializeToString())
pic_id+=1
print ("the pic is:",pic_id)
writer.close()
print ("The data is writed successful!")
def get_label_data(img_path):
if not os.path.exists(img_path):
print ("input data error!")
return
root = os.listdir(img_path)
list_data={}
for r in root :
list_data[r]=r[0].upper()
return list_data
def generate_tfRecord():
if not os.path.exists(data_path):
os.makedirs(data_path)
print ("root create successfully!")
else:
print ("root already exists!")
train_label_data=get_label_data(image_train_path)
#print (train_label_data)
write_record(tfRecord_train,image_train_path,train_label_data)
test_label_data=get_label_data(image_test_path)
write_record(tfRecord_test,image_test_path,test_label_data)
def main():
generate_tfRecord()
if __name__ == "__main__":
main()
result:
root already exists!
the pic is: 1
the pic is: 2
the pic is: 3
the pic is: 4
the pic is: 5
the pic is: 6
the pic is: 7
the pic is: 8
the pic is: 9
the pic is: 10
the pic is: 11
the pic is: 12
the pic is: 13
the pic is: 14
the pic is: 15
the pic is: 16
the pic is: 17
the pic is: 18
the pic is: 19
the pic is: 20
the pic is: 21
the pic is: 22
the pic is: 23
the pic is: 24
the pic is: 25
the pic is: 26
the pic is: 27
the pic is: 28
the pic is: 29
the pic is: 30
The data is writed successful!
the pic is: 1
the pic is: 2
the pic is: 3
the pic is: 4
the pic is: 5
the pic is: 6
The data is writed successful!
project root
train data
test data
output file