tensorflow 2.0-根据图片制作数据集

	#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