怎么在tensorflow中讀取tfrecord文件

今天就跟大家聊聊有關(guān)怎么在tensorflow中讀取tfrecord文件，可能很多人都不太了解，為了讓大家更加了解，小編給大家總結(jié)了以下內(nèi)容，希望大家根據(jù)這篇文章可以有所收獲。

1、生成tfrecord文件

import os
import numpy as np
import tensorflow as tf
from PIL import Image

filenames = [
'images/cat/1.jpg',
'images/cat/2.jpg',
'images/dog/1.jpg',
'images/dog/2.jpg',
'images/pig/1.jpg',
'images/pig/2.jpg',]

labels = {'cat':0, 'dog':1, 'pig':2}

def int64_feature(values):
	if not isinstance(values, (tuple, list)):
		values = [values]
	return tf.train.Feature(int64_list=tf.train.Int64List(value=values))

def bytes_feature(values):
	return tf.train.Feature(bytes_list=tf.train.BytesList(value=[values]))

with tf.Session() as sess:
	output_filename = os.path.join('images/train.tfrecords')
	with tf.python_io.TFRecordWriter(output_filename) as tfrecord_writer:
		for filename in filenames:
			#讀取圖像
			image_data = Image.open(filename)
			#圖像灰度化
			image_data = np.array(image_data.convert('L'))
			#將圖像轉(zhuǎn)化為bytes
			image_data = image_data.tobytes()
			#讀取label
			label = labels[filename.split('/')[-2]]
			#生成protocol數(shù)據(jù)類型
			example = tf.train.Example(features=tf.train.Features(feature={'image': bytes_feature(image_data),
																			'label': int64_feature(label)}))
			tfrecord_writer.write(example.SerializeToString())

2、讀取tfrecord文件

import tensorflow as tf
import matplotlib.pyplot as plt
from PIL import Image

# 根據(jù)文件名生成一個(gè)隊(duì)列
filename_queue = tf.train.string_input_producer(['images/train.tfrecords'])
reader = tf.TFRecordReader()
# 返回文件名和文件
_, serialized_example = reader.read(filename_queue)
features = tf.parse_single_example(serialized_example, 
									features={'image': tf.FixedLenFeature([], tf.string), 
												'label': tf.FixedLenFeature([], tf.int64)})
# 獲取圖像數(shù)據(jù)
image = tf.decode_raw(features['image'], tf.uint8)
# 恢復(fù)圖像原始尺寸[高，寬]
image = tf.reshape(image, [60, 160])
# 獲取label
label = tf.cast(features['label'], tf.int32)

with tf.Session() as sess:
	# 創(chuàng)建一個(gè)協(xié)調(diào)器，管理線程
	coord = tf.train.Coordinator()
	# 啟動(dòng)QueueRunner, 此時(shí)文件名隊(duì)列已經(jīng)進(jìn)隊(duì)
	threads = tf.train.start_queue_runners(sess=sess, coord=coord)

	for i in range(6):
		image_b, label_b = sess.run([image, label])
		img = Image.fromarray(image_b, 'L')
		plt.imshow(img)
		plt.axis('off')
		plt.show()
		print(label_b)

	# 通知其他線程關(guān)閉
	coord.request_stop()
	# 其他所有線程關(guān)閉之后，這一函數(shù)才能返回
	coord.join(threads)

看完上述內(nèi)容，你們對(duì)怎么在tensorflow中讀取tfrecord文件有進(jìn)一步的了解嗎？如果還想了解更多知識(shí)或者相關(guān)內(nèi)容，請(qǐng)關(guān)注億速云行業(yè)資訊頻道，感謝大家的支持。