tensorflow分类任务——TFRecord制作自己的数据集

制作数据集的前菜——打乱数组

	x, y = imageNumpyData()
    state = np.random.get_state()
    np.random.shuffle(x)
    np.random.set_state(state)
    np.random.shuffle(y)
    # 打乱之后的x，y作为训练数据
    x = np.array(x)
    y = np.array(y)

1. 代码运行之前,x是一个list，y是一个array
2. 使用随机数的方式打乱数据
3. 如果x和y的数据样本数不一样，容易出现错误
4. 打乱之后全部变成array

制作数据集的正餐

数据集的流程：

1. 先读取输入，并把数据保存成固定通道，固定维数。
2. 保存成numpy的形式，之后打乱数据集
3. 数据集一张一张写入TFRecord中

注意事项：

1. Data里面存放数据
2. record里面存放制作成功的数据集
3. 需要提前准备好了record这个文件夹，不然容易报错，如果要自动创建文件夹请看我另一篇博客，这个不做赘述
4. 我的图片格式是JPEG，我在代码中写了判断，同时也可以制作BMP和PNG类型的图片
5. 我的图片是灰度图，所以是单通道，读者若是彩色图片，相应更改代码，我在下面的解释中有说明
   

所有代码

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


def imageNumpyData():
    # the classification file root path
    rootFilePath = './Data/'
    rootAbsPath = os.path.abspath(rootFilePath)
    node1FloderPath = os.listdir(rootAbsPath)
    node1Path = [os.path.join(rootAbsPath, nodePath) for nodePath in node1FloderPath]
    label = 0
    imageLabel = []
    imageData = []
    for imagePath in node1Path:
        display = 0
        image = os.listdir(imagePath)
        image = [os.path.join(imagePath, file) for file in image]
        # 使用tensorflow
        with tf.Session() as sess:
            for image_raw in image:
                try:
                    img = Image.open(image_raw)
                    if img.format == "BMP":
                        image_raw_data = tf.gfile.FastGFile(image_raw, 'rb').read()
                        image_data = tf.image.decode_bmp(image_raw_data)
                    if img.format == "JPEG":
                        image_raw_data = tf.gfile.FastGFile(image_raw, 'rb').read()
                        image_data = tf.image.decode_jpeg(image_raw_data)
                    if img.format == "PNG":
                        image_raw_data = tf.gfile.FastGFile(image_raw, 'rb').read()
                        image_data = tf.image.decode_png(image_raw_data)
                    if image_data.dtype != tf.float32:
                        image_data = tf.image.convert_image_dtype(image_data, dtype=tf.float32)
                    image_data = tf.image.resize_images(image_data, [224, 224])
                    image_temp = np.array(sess.run(image_data))
                    if image_temp.shape[2] == 1:
                        imageData.append(image_temp)
                        imageLabel.append(label)
                        display = display + 1
                        print(str(label) + ':' + str(display))
                except:
                    print(image_raw + '读取失败')
                    # os.remove(image_raw)
        label = label + 1
    one_hot = tf.one_hot(imageLabel, label)
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        imageLabel = sess.run(one_hot)
    return imageData, imageLabel


def tfRecordWrite(filename):
    # 生成整数的属性
    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]))

    # filename = 'record/Imageoutput.tfrecords'
    writer = tf.python_io.TFRecordWriter(filename)
    x, y = imageNumpyData()
    state = np.random.get_state()
    np.random.shuffle(x)
    np.random.set_state(state)
    np.random.shuffle(y)
    # 打乱之后的x，y作为训练数据
    x = np.array(x)
    y = np.array(y)

    # 将每张图片都转为一个Example
    height = x.shape[1]
    width = x.shape[2]
    channels = x.shape[3]
    n_class = y.shape[1]

    x = x.reshape([-1, height, width, channels])

    division = int(x.shape[0] * 0.9)

    x_train = x[:division]
    y_train = y[:division]
    x_test = x[division:]
    y_test = y[division:]
    np.savez('testData.npz', test_X=x_test, test_Y=y_test, height=height, width=width, channels=channels,
             n_class=n_class)
    del x_test, y_test
    for i in range(x_train.shape[0]):
        image = x_train[i].tostring()  # 将图像转为字符串
        example = tf.train.Example(features=tf.train.Features(
            feature={
                'image': _bytes_feature(image),
                'label': _int64_feature(np.argmax(y_train[i]))
            }))
        writer.write(example.SerializeToString())  # 将Example写入TFRecord文件
    writer.close()

    return height, width, channels, n_class


filename = r'./record\Imageoutput.tfrecords'
height, width, channels, n_class = tfRecordWrite(filename)
print(height, width, channels, n_class)
print('data processing success')

关键代码解释

                 try:
                    img = Image.open(image_raw)
                    if img.format == "BMP":
                        image_raw_data = tf.gfile.FastGFile(image_raw, 'rb').read()
                        image_data = tf.image.decode_bmp(image_raw_data)
                    if img.format == "JPEG":
                        image_raw_data = tf.gfile.FastGFile(image_raw, 'rb').read()
                        image_data = tf.image.decode_jpeg(image_raw_data)
                    if img.format == "PNG":
                        image_raw_data = tf.gfile.FastGFile(image_raw, 'rb').read()
                        image_data = tf.image.decode_png(image_raw_data)

解释：尝试打开图片，并且图片格式是JPEG，BMP或PNG， tf.gfile.FastGFile的这个地方必须要‘rb’，我看一些书上仅仅是写了一个‘r’，这样并不对。

 if image_data.dtype != tf.float32:
                            image_data = tf.image.convert_image_dtype(image_data, dtype=tf.float32)

解释：我做CNN的中，需要float32，这里必须转格式

image_data = tf.image.resize_images(image_data, [224, 224])

解释：把图片更改大小

 if image_temp.shape[2] == 1:
                            imageData.append(image_temp)
                            imageLabel.append(label)

解释:我的是灰度图，所以这里**‘’==1‘’，若是三通道，则‘==3’**

one_hot = tf.one_hot(imageLabel, label)
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        imageLabel = sess.run(one_hot)

解释：把得到的label，转换成one-hot形式

 x = x.reshape([-1, height, width, channels])

    division = int(x.shape[0] * 0.9)

    x_train = x[:division]
    y_train = y[:division]
    x_test = x[division:]
    y_test = y[division:]
    np.savez('testData.npz', test_X=x_test, test_Y=y_test, height=height, width=width, channels=channels,
             n_class=n_class)
    del x_test, y_test

解释：这里是我拿出一部分数据保存成‘npz’的形式作为我的测试集，因为我用tfrecord格式的数据作为训练集

 # 生成整数的属性
    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]))

解释：这是我们必须遵从的规则，用整数属性写入label，用字符串形式写入图像数据

# filename = 'record/Imageoutput.tfrecords'
    writer = tf.python_io.TFRecordWriter(filename)

解释：开始写入，并生成文件

    for i in range(x_train.shape[0]):
        image = x_train[i].tostring()  # 将图像转为字符串
        example = tf.train.Example(features=tf.train.Features(
            feature={
                'image': _bytes_feature(image),
                'label': _int64_feature(np.argmax(y_train[i]))
            }))
        writer.write(example.SerializeToString())  # 将Example写入TFRecord文件
    writer.close()

解释：根据数据集大小，一个一个的写入

数据集地址：
链接：https://pan.baidu.com/s/1aIHzKsxUb67sJZAFrGH1ZQ
提取码：lvjp

工程地址：
链接：https://pan.baidu.com/s/1XGAA6UQ0JByhvDYQ__my4g
提取码：dxpn

结束语

希望大家积极学习，有不懂的地方可以加我的微信:ChaofeiLi
如果读者觉得写得好的可以稍微打赏下作者