Tensorflow之TFRecord读写自己的数据（一）

记录如何使用自己的数据训练神经网络

安装：
pip install scikit-image
pip install numpy
pip install matplotlib
pip imstall tensorflow
pip install Pillow

---

导入的函数模块

import tensorflow as tf
import skimage.io as io
import  numpy as np
import math
import os
import matplotlib.pyplot as plt
from PIL import Image

Step1:

从本地（pc)端获取数据——图像

def get _file(file_dir):
	images =  [] # 用于存放图像名称的list
	temp = [] # 存放图像路径的list
	lable = [] # 存放图像标签的list
	for root,sub_folders,files in os.walk(file_dir): # os.walk()可以直接获取路径下的根目录，子文件夹及子文件
		for name in files:
			if name.endswith == (.jpg):
				images.append(os.path.join(root,name)) # 如果文件名是以.jpg结尾，将其保存进images[]
		for name in sub_folders:
			temp.append(os.path.join(root,name))
		for  one_folder in temp:
			n_img = len(os.listdir(one_folder)) # 计算子文件夹中的图像个数
			class_name = one_folder.split("\")[-1]# 类名称就是图像路径的最后一个“\"后面的名字，也就是子文件夹的文件名。
			if class_name == 'cats': # 把cats文件夹里的图片都打上标签’1‘
				labels = np.append(labels,n_img*[1])
			elif class_name == 'dogs': # 把dogs文件夹里的标签都打上标签’2‘
				labels = np.append(labels,n_img*[2])

	# shuffle
	temp = np.array([images,labels])  # (2,40)
	temp = temp.transpose()  # (40,2)
	np.random.shuffle(temp)
	image_list = list(temp[:,0])  #取矩阵的第一列
	lable_list = list(temp[:,1]) #取矩阵的第二列
	return image_list,label_list	# 返回图片和标签列表

Step2:

将获取到的图像转化为tfrecord格式

def _int64_feature(value):
	if not isinstance(value,list):
		value  = [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 convert_to_tfrecord(images,labels,save_dir,name):
	filename = os.path.join(save_dir,name+'.tfrecords') #tfrecord文件的文件名称
	n_labels = len(labels) #label 的个数 
	# 检测图像数量是否与标签数量相符
	if np.shape(images) != n_labels:
		raise ValueError('Image size %d dose not match labels size %d.' % (images.size(),labels.size()))
		
		writer = tf.python_io.TFRecordWriter(filename)
		print('Transform start...') #开始写入
		m=n=0
		for i in np.arange(0,n_labels):
			try:
				m+=1
				image = Image.open(images[i])
				image = image.resize(IMAGE_SIZE) #统一图像大小
				image_raw = image.tobytes() #将图片转换为二进制格式，便于写入TFRecord文件
				label = int(labels[i])
				example = tf.train.Example(features=tf.train.Features(feature={'label':_int64_feature(label),'image_raw':_bytes_feature(image_raw)}))
				writer.write(example.SerializeToString())
			except IOError as e:
				n+=1
				print('Could not read: ’，images[i])
				print("Error type:%s" %e)
				print(‘Skip it!\n')
		writer.close()
		print("Transform done!")
		print("Success:%d\nfailed:%d" % (m,n)) #写入成功次数和失败次数
		return filename

Step3：

从.tfrecord文件读取并将二进制格式转化为Tensor

• tf.train.string_input_producer()函数需要传入一个文件名list,系统自动将它转换为一个文件名列队（该函数目前已被官方弃用）

def read_and_decord(tfrecords_file,batch_size):
	filename_queue = tf.train.string_input_producer([tfrecords_file])、
	reader = tf.TFRecordReader()
	_,serialized_example = reader.read(filename_queue)
	img_features = tf.parse_single_example(serialized_example,
											features={'label':tf.FixedLenFeature([],tf.int64),
											'image_raw':tf.FixedLenFeature([],tf.string)})
	image = tf.decode(features['image_raw'],tf.uint8)
	image = tf.reshape(image,[IMAGE_SIZE[0],IMAGE_SIZE[1],3])
	label = tf.cast(feature['label'],tf.int32)
	image_batch,label_batch = tf.train.batch([image,label],
											batch_size=batch_size,
											num_threads=64,
											capacity=2000)
	return image_batch,tf.reshape(label_batch,[batch_size])

Test code:

• 在使用tf.train.string_input_producer()创建文件名列队后，整个系统其实还处于’停滞状态’，也就是说，文件名并没真正加入列队。
• 而使用tf.train.start_queue_runners()函数后，才会启动填充列队的线程。此后，计算单元才会拿到数据并进行计算，整个程序也就运行起来了。

FILE_DIR = '/TFRecord'
IMAGE_SIZE = (208,208)
BATCH_SIZE = 9
# 通过文件路径 获取图像及对应的标签
image_list,label_list = get_file(FILE_DIR)
# 将图像转换为.tfrecord文件并返回其文件名
filename = convert_to_tfrecord(image_list,label_list,FILE_DIR+'/data,'training')
# 读取.tfredord文件，转化为Tensor,并将图像分批（batch）
image_batch,label_batch = read_and_decode(filename,BATCH_SIZE)

# Session会话
with tf.Session() as sess:
	i = 0
	coord = tf.train.Coordinator()
	threads = tf.train.strat_queue_runners(coord=coord)
	try:
		while not coord.should_stop() and i<1:
			image,label = sess.run([image_batch,label_batch])
			plt.figure()
			for j in np.arange(1,10):
				plt.subplot(3,3,j)
				plt.imshow(image[j-1])
				if label[j-1] == 1:
					label_name ='cat'
				elif label[j-1] == 2:
					label_name = 'dog'
				plt.title(label_name,fontsize=14)
			plt.show()
			i+=1
		except tf.errors.OutOfRangeError:
			print('done!')
		finally:
			coord.request_stop()
		coord.join(threads)

Complete code:

import tensorflow as tf
import numpy as np
import os
import matplotlib.pyplot as plt
import skimage.io as io
import math
from PIL import Image

def get_file(file_dir):
	images = []
	temp = []
	for root,sub_folders,files in os.walk(file_dir):
		print('root:%s\nsub_folders:%s\nfiles:%s' % (root,sub_folders,files))
		for name in files:
			if name.endswith('.jpg'):
				images.append(os.path.join(root,name))	
		for name in sub_folders:
			temp.append(os.path.join(root,name))
	print('iamges:',images)
	print('\n')
	print('temp:',temp)
	
	labels = []
	for one_folder in temp:
		n_img = len(os.listdir(one_folder))
		print('n_img:',n_img)
		class_name = one_folder.split('\\')[-1]
		if class_name == 'cats':
			labels = np.append(labels,n_img*[1])
		elif class_name == 'dogs':
			labels = np.append(labels,n_img*[2])
		
	print('labels:',labels)

	# shuffle
	temp = np.array([images,labels]) # [2,40]
	print('temp:',temp)
	temp = temp.transpose()  # [40,2]
	print('temp transpose:',temp)
	np.random.shuffle(temp)
	print('temp shuffle:',np.random.shuffle(temp))
	image_list = list(temp[:,0]) # 取矩阵第一列
	label_list = list(temp[:,1]) # 取矩阵第二列
	# for i in label_list:
	# 	label_list[int(float(i))] = int(float(i))
	label_list = [int(float(i)) for i in label_list]
	print("image list:%s \nlabel list:%s"%(image_list,label_list))
	return image_list , label_list

def _int64_feature(value):
	if not isinstance(value,list):
		value=[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 convert_to_tfrecord(images,labels,save_dir,name):
	filename = os.path.join(save_dir , name + '.tfrecords')
	n_samples = len(labels)
	print('len labels:',n_samples)

	if np.shape(images)[0] != n_samples:
		raise ValueError('Images size %d does not match labels size %d.' % (images.size(),labels.size()))

	writer = tf.python_io.TFRecordWriter(filename)
	print('\nTransform start...')
	m=n=0
	for i in np.arange(0,n_samples):
		try:
			m+=1
			print('images[i]:',images[i])
			image = Image.open(images[i])
			image = image.resize(IMAGE_SIZE)

			image_raw = image.tobytes()
			label = int(labels[i])
			example = tf.train.Example(features=tf.train.Features(feature={
				'label':_int64_feature(label),
				'image_raw':_bytes_feature(image_raw)
				}))
			writer.write(example.SerializeToString())
		except IOError as e:
			n+=1
			print('Could not read:',images[i])
			print('Error type:%s'%e)
			print('Skip it!\n')
	writer.close()	
	print('Transform done!')
	print('success:%d\nfailed:%d\n'%(m,n))
	return filename

def read_and_decode(tfrecords_file,batch_size):
	filename_queue = tf.train.string_input_producer([tfrecords_file])
	# filename_queue = tf.data.Dataset.from_tensor_slices([tfrecords_file])
	reader = tf.TFRecordReader()
	_,serialized_example = reader.read(filename_queue)
	img_features = tf.parse_single_example(
											serialized_example,
											features={
											'label':tf.FixedLenFeature([],tf.int64),
											'image_raw':tf.FixedLenFeature([],tf.string)
											})
	image = tf.decode_raw(img_features['image_raw'],tf.uint8)

	#
	#you can put data augmentation here
	#

	image = tf.reshape(image,[IMAGE_SIZE[0],IMAGE_SIZE[1],3])
	label = tf.cast(img_features['label'],tf.int32)
	image_batch,label_batch = tf.train.batch([image,label],
										batch_size=batch_size,
										num_threads=64,
										capacity=2000,
											)
	print('iamge batch:',image_batch)
	print('label batch:',label_batch)
	return image_batch ,tf.reshape(label_batch,[batch_size])


FILE_DIR = 'F:/OpenCV-Python/TFRecord'
IMAGE_SIZE = (208,208)
BATCH_SIZE = 16
image_list,label_list = get_file(FILE_DIR)

filename = convert_to_tfrecord(image_list,label_list,FILE_DIR+'/data',"training")

image_batch,label_batch = read_and_decode(filename,BATCH_SIZE)

print('image shape:',image_batch.shape)

with tf.Session() as sess:
	i =0
	coord = tf.train.Coordinator()
	threads = tf.train.start_queue_runners(coord=coord)
	try:
		while not coord.should_stop() and i<1:
			image,label = sess.run([image_batch,label_batch])
			plt.figure()
			for j in np.arange(1,10):
				plt.subplot(3,3,j)
				plt.imshow(image[j-1])
				if label[j-1] == 1:
					label_name = 'cat'
				elif label[j-1] == 2:
					label_name = 'dog'
				plt.title(label_name,fontsize=14)
			plt.show()
			print('label:',label)
			i+=1
	except tf.errors.OutOfRangeError:
		print('done!')
	finally:
		coord.request_stop()	
	coord.join(threads)