Tensorflow生成自己的图片数据集TFrecords

Tensorflow生成自己的图片数据集TFrecords

train.txt保存图片的路径和标签信息

test_image/dog/1.jpg 0
test_image/dog/2.jpg 0
test_image/dog/3.jpg 0
test_image/dog/4.jpg 0
test_image/cat/1.jpg 1
test_image/cat/2.jpg 1
test_image/cat/3.jpg 1
test_image/cat/4.jpg 1

使用下面完整代码，可以生成自己的图片数据集TFrecords，并解析：

# -*- coding: utf-8 -*-
# !/usr/bin/python3.5
# ref url : https://blog.csdn.net/guyuealian/article/details/80857228
# Author  : pan_jinquan
# Date    : 2018.6.29
# Function: image convert to tfrecords
#############################################################################################

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

# 参数设置
###############################################################################################
train_file = 'train.txt'  # 图片路径
output_record_dir= './tfrecords/my_record.tfrecords'
resize_height = 100   # 指定存储图片高度
resize_width = 100    # 指定存储图片宽度


###############################################################################################
# 生成整数型的属性
def _int64_feature(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 float_list_feature(value):
  return tf.train.Feature(float_list=tf.train.FloatList(value=value))

# 显示图片
def show_image(image_name,image):
    # plt.figure("show_image")  # 图像窗口名称
    plt.imshow(image)
    plt.axis('on')  # 关掉坐标轴为 off
    plt.title(image_name)  # 图像题目
    plt.show()

#载图txt文件，文件中每行为一个图片信息，且以空格隔开：图像路径 标签，如：test_image/1.jpg 0
def load_txt_file(examples_list_file):
    lines = np.genfromtxt(examples_list_file, delimiter=" ", dtype=[('col1', 'S120'), ('col2', 'i8')])
    examples = []
    labels = []
    for example, label in lines:
        examples.append(example.decode('ascii'))
        labels.append(label)
    return np.asarray(examples), np.asarray(labels), len(lines)

# 读取原始图像数据
def read_image(filename, resize_height, resize_width):
    # image = cv2.imread(filename)
    # image = cv2.resize(image, (resize_height, resize_width))
    # b, g, r = cv2.split(image)
    # rgb_image = cv2.merge([r, g, b])
    rgb_image=Image.open(filename)
    rgb_image=rgb_image.resize((resize_width,resize_height))
    image=np.asanyarray(rgb_image)
    # show_image("src resize image",image)
    return image


#保存record文件
def create_record(train_file, output_record_dir, resize_height, resize_width):
    _examples, _labels, examples_num = load_txt_file(train_file)
    writer = tf.python_io.TFRecordWriter(output_record_dir)
    for i, [example, label] in enumerate(zip(_examples, _labels)):
        print('No.%d' % (i))
        image = read_image(example, resize_height, resize_width)
        print('shape: %d, %d, %d, label: %d' % (image.shape[0], image.shape[1], image.shape[2], label))
        image_raw = image.tostring()
        example = tf.train.Example(features=tf.train.Features(feature={
            'image_raw': _bytes_feature(image_raw),
            'height': _int64_feature(image.shape[0]),
            'width': _int64_feature(image.shape[1]),
            'depth': _int64_feature(image.shape[2]),
            'label': _int64_feature(label)
        }))
        writer.write(example.SerializeToString())
    writer.close()

#解析record文件，并显示，主要用于验证
def disp_records(tfrecord_list_file):
    filename_queue = tf.train.string_input_producer([tfrecord_list_file])
    reader = tf.TFRecordReader()
    _, serialized_example = reader.read(filename_queue)
    features = tf.parse_single_example(
        serialized_example,
        features={
            'image_raw': tf.FixedLenFeature([], tf.string),
            'height': tf.FixedLenFeature([], tf.int64),
            'width': tf.FixedLenFeature([], tf.int64),
            'depth': tf.FixedLenFeature([], tf.int64),
            'label': tf.FixedLenFeature([], tf.int64)
        }
    )
    image = tf.decode_raw(features['image_raw'], tf.uint8)
    # print(repr(image))
    height = features['height']
    width = features['width']
    depth = features['depth']
    label = tf.cast(features['label'], tf.int32)
    init_op = tf.initialize_all_variables()
    resultImg = []
    resultLabel = []
    with tf.Session() as sess:
        sess.run(init_op)
        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(sess=sess, coord=coord)
        for i in range(4):
            image_eval = image.eval()
            resultLabel.append(label.eval())
            image_eval_reshape = image_eval.reshape([height.eval(), width.eval(), depth.eval()])
            print("image.shape=",height.eval(), width.eval(), depth.eval())
            resultImg.append(image_eval_reshape)
            # pilimg = Image.fromarray(np.asarray(image_eval_reshape))
            # pilimg.show()
            show_image("decode_from_tfrecords",image_eval_reshape)
        coord.request_stop()
        coord.join(threads)
        sess.close()
    return resultImg, resultLabel

#解析record文件
def read_record(filename_queuetemp):
    filename_queue = tf.train.string_input_producer([filename_queuetemp])
    reader = tf.TFRecordReader()
    _, serialized_example = reader.read(filename_queue)
    features = tf.parse_single_example(
        serialized_example,
        features={
            'image_raw': tf.FixedLenFeature([], tf.string),
            'width': tf.FixedLenFeature([], tf.int64),
            'depth': tf.FixedLenFeature([], tf.int64),
            'label': tf.FixedLenFeature([], tf.int64)
        }
    )
    image = tf.decode_raw(features['image_raw'], tf.uint8)#获得图像原始的数据
    image=tf.reshape(image, [100, 100, 3]) # 设置图像的维度
    # image = tf.cast(image, tf.float32) * (1. / 255) - 0.5 # 归一化
    label = tf.cast(features['label'], tf.int32) # label
    return image, label



def test():
    create_record(train_file, output_record_dir, resize_height, resize_width)  # 产生record文件

    # img, label = disp_records(output_record_dir)  # 显示函数

    image_data, label_data= read_record(output_record_dir)  # 读取函数


    #使用shuffle_batch可以随机打乱输入
    image_batch, label_batch = tf.train.shuffle_batch([image_data, label_data], batch_size=30, capacity=2000,min_after_dequeue=1000)

    init = tf.global_variables_initializer()
    with tf.Session() as sess:#开始一个会话
        sess.run(init)
        coord = tf.train.Coordinator()
        threads = tf.train.start_queue_runners(coord=coord)
        for i in range(4):
            images, labels= sess.run([image_batch, label_batch])#在会话中取出image和label
            #我们也可以根据需要对val， l进行处理
            #l = to_categorical(l, 12)
            show_image("image",images[i,:,:,:])
            print(images.shape, labels)
        #停止所有线程
        coord.request_stop()
        coord.join(threads)
        sess.close()#关闭会话


if __name__ == '__main__':
    test()

假设我们已经生成了output.tfrecords，其中保存有：

image_raw：图像的特征向量，有784维
pixels：图像分辨率大小28
label：图像的标签

 可以使用下面的方法，进行训练网络：

#coding=utf-8
import tensorflow as tf

# 模型相关的参数
INPUT_NODE = 784
OUTPUT_NODE = 10
LAYER1_NODE = 500
REGULARAZTION_RATE = 0.0001   
TRAINING_STEPS = 5000


files = tf.train.match_filenames_once("./output.tfrecords")
filename_queue = tf.train.string_input_producer(files, shuffle=False)


# 读取文件。

reader = tf.TFRecordReader()
_,serialized_example = reader.read(filename_queue)

# 解析读取的样例。
features = tf.parse_single_example(
    serialized_example,
    features={
        'image_raw':tf.FixedLenFeature([],tf.string),
        'pixels':tf.FixedLenFeature([],tf.int64),
        'label':tf.FixedLenFeature([],tf.int64)
    })

decoded_images = tf.decode_raw(features['image_raw'],tf.uint8)
retyped_images = tf.cast(decoded_images, tf.float32)
labels = tf.cast(features['label'],tf.int32)
#pixels = tf.cast(features['pixels'],tf.int32)
images = tf.reshape(retyped_images, [784])
min_after_dequeue = 10000
batch_size = 100
capacity = min_after_dequeue + 3 * batch_size

image_batch, label_batch = tf.train.shuffle_batch([images, labels], 
                                                    batch_size=batch_size, 
                                                    capacity=capacity, 
                                                    min_after_dequeue=min_after_dequeue)
def inference(input_tensor, weights1, biases1, weights2, biases2):
        layer1 = tf.nn.relu(tf.matmul(input_tensor, weights1) + biases1)
        return tf.matmul(layer1, weights2) + biases2



weights1 = tf.Variable(tf.truncated_normal([INPUT_NODE, LAYER1_NODE], stddev=0.1))
biases1 = tf.Variable(tf.constant(0.1, shape=[LAYER1_NODE]))

weights2 = tf.Variable(tf.truncated_normal([LAYER1_NODE, OUTPUT_NODE], stddev=0.1))
biases2 = tf.Variable(tf.constant(0.1, shape=[OUTPUT_NODE]))

y = inference(image_batch, weights1, biases1, weights2, biases2)
    
# 计算交叉熵及其平均值
cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=y, labels=label_batch)
cross_entropy_mean = tf.reduce_mean(cross_entropy)
    
# 损失函数的计算
regularizer = tf.contrib.layers.l2_regularizer(REGULARAZTION_RATE)
regularaztion = regularizer(weights1) + regularizer(weights2)
loss = cross_entropy_mean + regularaztion

# 优化损失函数
train_step = tf.train.GradientDescentOptimizer(0.01).minimize(loss)
    
# 初始化回话并开始训练过程。
with tf.Session() as sess:
    tf.global_variables_initializer().run()
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess=sess, coord=coord)
    # 循环的训练神经网络。
    for i in range(TRAINING_STEPS):
        if i % 1000 == 0:
            print("After %d training step(s), loss is %g " % (i, sess.run(loss)))
                  
        sess.run(train_step) 
    coord.request_stop()
    coord.join(threads)

参考资料：

[1]https://blog.csdn.net/happyhorizion/article/details/77894055  (五星推荐)

[2]https://blog.csdn.net/ywx1832990/article/details/78462582

[3]https://blog.csdn.net/csuzhaoqinghui/article/details/51377941