Python:制作Tensorflow需要的tfrecord格式

回顾

之前几篇文章，算是在弯路中崎岖前行。
因为没有成功安装labeliamge工具，所以走了两条路：

直接用Matlab标记好的数据生成CSV文件；
模拟labeliamge工具，标记数据先为每张图生成xml文件，然后用转换代码生成CSV文件。

第二个方法是多此一举的意思，主要是当时在排查错误源头，不过也算是学习了很多读写文件的方式。

回顾一下我之前写的几篇文章。

机器学习之Matlab制作自己的数据集
 借助Mtalab制作数据集之保存到CSV

tensorflow数据集处理：CSV中绝对路径转相对路径（C++写）
Matlab制作Tensorflow数据集：将数据写入XML文件

标准CSV文件

由于Matlab那块没有目标类别标记功能，因此每次只能做一个目标的标记，然后将多个CSV文件合并。这个直接复制粘贴就可以，不是太费事。

这个是我生成的一个CSV文件：

在这里插入图片描述

注意文件头一定是这几个标签，不可以随意修改；否则你需要去寻找代码，在对应地方修改。

CSV生成tfrecord

直接代码是参照别人的：

目标检测Tensorflow object detection API之构建自己的模型

"""
Usage:
  # From tensorflow/models/
  # Create train data:
python csv_tfrecord.py --csv_input=data/result.csv  --output_path=data/result.record
  # Create test data:
python generate_TFR.py --csv_input=data/test.csv  --output_path=data/test.record
  需要修改三处
  os.chdir('D:\\python3\\models-master\\research\\object_detection\\')
  path = os.path.join(os.getcwd(), 'images/train')
  def class_text_to_int(row_label): #对应的标签返回一个整数，后面会有文件用到
    if row_label == 'ZJL':
        return 1
    elif row_label == 'CYX':
        return 2
    else:
        None
"""



import os
import io
import pandas as pd
import tensorflow as tf

from PIL import Image
from utils import dataset_util
# import dataset_util
from collections import namedtuple, OrderedDict

os.chdir('D:\\TF_0913\\models-master\\models-master\\research\\object_detection')

flags = tf.app.flags
flags.DEFINE_string('csv_input', '', 'Path to the CSV input')
flags.DEFINE_string('output_path', '', 'Path to output TFRecord')
FLAGS = flags.FLAGS


# TO-DO replace this with label map
def class_text_to_int(row_label):
    if row_label == 'car':
        return 1
    elif row_label == 'boat':
        return 2
    else:
        None


def split(df, group):
    data = namedtuple('data', ['filename', 'object'])
    gb = df.groupby(group)
    return [data(filename, gb.get_group(x)) for filename, x in zip(gb.groups.keys(), gb.groups)]


def create_tf_example(group, path):
    with tf.gfile.GFile(os.path.join(path, '{}'.format(group.filename)), 'rb') as fid:
        encoded_jpg = fid.read()
    encoded_jpg_io = io.BytesIO(encoded_jpg)
    image = Image.open(encoded_jpg_io)
    width, height = image.size

    filename = group.filename.encode('utf8')
    image_format = b'jpg'
    xmins = []
    xmaxs = []
    ymins = []
    ymaxs = []
    classes_text = []
    classes = []

    for index, row in group.object.iterrows():
        xmins.append(row['xmin'] / width)
        xmaxs.append(row['xmax'] / width)
        ymins.append(row['ymin'] / height)
        ymaxs.append(row['ymax'] / height)
        classes_text.append(row['class'].encode('utf8'))
        classes.append(class_text_to_int(row['class']))

    tf_example = tf.train.Example(features=tf.train.Features(feature={
        'image/height': dataset_util.int64_feature(height),
        'image/width': dataset_util.int64_feature(width),
        'image/filename': dataset_util.bytes_feature(filename),
        'image/source_id': dataset_util.bytes_feature(filename),
        'image/encoded': dataset_util.bytes_feature(encoded_jpg),
        'image/format': dataset_util.bytes_feature(image_format),
        'image/object/bbox/xmin': dataset_util.float_list_feature(xmins),
        'image/object/bbox/xmax': dataset_util.float_list_feature(xmaxs),
        'image/object/bbox/ymin': dataset_util.float_list_feature(ymins),
        'image/object/bbox/ymax': dataset_util.float_list_feature(ymaxs),
        'image/object/class/text': dataset_util.bytes_list_feature(classes_text),
        'image/object/class/label': dataset_util.int64_list_feature(classes),
    }))
    return tf_example


def main(_):
    writer = tf.python_io.TFRecordWriter(FLAGS.output_path)
    path = os.path.join(os.getcwd(), 'images/images_train') #20180418做了修改
    examples = pd.read_csv(FLAGS.csv_input)
    grouped = split(examples, 'filename')
    for group in grouped:
        tf_example = create_tf_example(group, path)
        writer.write(tf_example.SerializeToString())

    writer.close()
    output_path = os.path.join(os.getcwd(), FLAGS.output_path)
    print('Successfully created the TFRecords: {}'.format(output_path))


if __name__ == '__main__':
    tf.app.run()

需要修改的地方，作者已经很明了的介绍了，因为自己跳过的坑，还是分享一下：

这部分代码是在cmd环境运行，并不是在IDE；（我不清楚在IDE可不可以）
路径是当前代码所在的路径，这个工程目录很大，自己下载后，将转换代码放在object_detection目录下。
第一处修改的路径为当前要操作的路径，也就是代码所在路径；
图片路径保存在这个工程目录下的，新建images文件夹保存；
标记好的CSV文件放在路径下的data文件夹下。

在这里插入图片描述

椭圆圈中的地方，为图片所在文件夹路径，相对于系统当前打开路径的位置，是相对位置。

系统打开的路径，为下图第一个矩形框所添加的路径。

在这里插入图片描述

代码运行
按照作者的思路，修改了三处：

1、系统需要打开的路径；
2、图片类别保存为标签；
3、图片所在文件夹相对于第一个路径的路径。

我在运行时候这句代码报错：from object_detection.utils import dataset_util
提示找不到object_detection
解决方法：

1、将utils import dataset_util.py文件拷贝到当前目录，使其可以找到；
2、上句修改为from utils import dataset_util

没有报错的自行忽略。

运行并不是在IDE，而是在文件目录下：
shift+右键，中间位置有个powershell点集进入，即可切换至当前代码所在路径。
然后将这句话复制进去即可运行。

//第一句为转换代码名称
//第二句为CSV文件相对于打开路径的相对路径：尝试过了，用绝对路径也可以
//第三句为需要生成的tfrecord文件所在路径
python csv_tfrecord.py --csv_input=data/result.csv  --output_path=data/result.record

在这里插入图片描述

哈哈哈，终于成功了，撒花花~~~~