使用tensorflow中的Dataset来读取制作好的tfrecords文件

 1 import tensorflow as tf
 2 import matplotlib.pyplot as plt
 3 
 4 #定义可以一次获得多张图像的函数
 5 def show_image(image_dir):
 6     plt.imshow(image_dir)
 7     plt.axis(‘on‘)
 8     plt.show()
 9 
10 #单个record的解析函数
11 def decode_example(example):#,resize_height,resize_width,labels_nums):
12     features=tf.io.parse_single_example(example,features={
13         ‘image_raw‘:tf.io.FixedLenFeature([],tf.string),
14         ‘label‘:tf.io.FixedLenFeature([],tf.int64)
15     })
16     tf_image=tf.decode_raw(features[‘image_raw‘],tf.uint8)#这个其实就是图像的像素模式，之前我们使用矩阵来表示图像
17     tf_image=tf.reshape(tf_image,shape=[224,224,3])#对图像的尺寸进行调整，调整成三通道图像
18     tf_image=tf.cast(tf_image,tf.float32)*(1./255)#对图像进行归一化以便保持和原图像有相同的精度
19     tf_label=tf.cast(features[‘label‘],tf.int32)
20     tf_label=tf.one_hot(tf_label,5,on_value=1,off_value=0)#将label转化成用one_hot编码的格式
21     return tf_image,tf_label
22 
23 def batch_test(tfrecords_file):
24     dataset=tf.data.TFRecordDataset(tfrecords_file)
25     dataset=dataset.map(decode_example)
26     dataset=dataset.shuffle(100).batch(4)
27     iterator=tf.compat.v1.data.make_one_shot_iterator(dataset)
28     batch_images,batch_labels=iterator.get_next()
29 
30     init_op=tf.compat.v1.global_variables_initializer()
31     with tf.compat.v1.Session() as sess:
32         sess.run(init_op)
33         coord=tf.train.Coordinator()
34         threads=tf.train.start_queue_runners(coord=coord)
35         for i in range(4):
36             images,labels=sess.run([batch_images,batch_labels])
37             show_image(images[1,:,:,:])
38             print(‘shape:{},tpye:{},labels:{}‘.format(images.shape, images.dtype, labels))
39 
40         coord.request_stop()
41         coord.join(threads)
42 
43 if __name__==‘__main__‘:
44     tfrecords_file=‘D:/软件/pycharmProject/wenyuPy/Dataset/VGG16/record/train.tfrecords‘
45     resize_height=224
46     resize_width=224
47     batch_test(tfrecords_file)