从LabelImg建立TFRecord格式的训练样本

完整内容见 作者的另一篇博客

图像样本转化为TFRecord

图像数据处理

用OpenCV读取图像数据

用OpenCV读取图像数据的格式为numpy.ndarray

import cv2

# 读取分辩率为1920*1080的3通道图像文件
img = cv2.imread("d:/temp/test_folder/image2.jpg", cv2.IMREAD_COLOR)  # cv2.IMREAD_COLOR值为1
print(img.shape)  # 输出: (1080, 1920, 3)
print(type(img))  # 输出: <class 'numpy.ndarray'>
print(img.dtype)  # 输出: uint8

# 修改宽高
img = cv2.resize(img, (480, 270))  # 注意dsize参数的元组按目标宽高排列

# 将BGR格式转为TensorFlow常用的RGB格式
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

# 显示图像
cv2.imshow("bgr", img)
cv2.imshow("rgb", img_rgb)
cv2.waitKey(0)

上述代码显示图像如下：

用TensorFlow读取图像数据

用TensorFlow读入图像文件的类型为bytes，即string类型

解析图像文件必须用Session运行tensor才能执行，而且必须注意float类型的归一化及其与uint8的转换

import cv2
import tensorflow as tf

# 读取分辩率为1920*1080的3通道图像文件
with tf.gfile.FastGFile("d:/temp/test_folder/image2.jpg", "rb") as fid:
    img = fid.read()
    print(type(img))  # 输出: <class 'bytes'>

    # 建立解析jpg图像为原始数据的tensor
    img_tensor = tf.image.decode_jpeg(img)
    print(type(img_tensor))  # 输出: <class 'tensorflow.python.framework.ops.Tensor'>
    print(img_tensor.shape)  # 输出: (?, ?, ?)
    print(img_tensor.dtype)  # 输出: <dtype: 'uint8'>

    # 指定tensor的shape
    # 无论是否设置shape，decode_jpeg都能正确解析图像
    img_tensor.set_shape([1080, 1920, 3])
    print(img_tensor.shape)  # 输出: (1080, 1920, 3)

    # 修改宽高
    # 特别注意resize_images输出类型为float32，即将uint8转化为float32类型
    # resize_images即可以操作图像batch的4D数据，也可以转化单幅图像的3D数据
    img_resize_float = tf.image.resize_images(img_tensor, size=(270, 480))
    print(type(img_resize_float))  # 输出: <class 'tensorflow.python.framework.ops.Tensor'>
    print(img_resize_float.shape)  # 输出: (270, 480, 3)，如果img_tensor未指定shape，则输出为(270, 480, ?)
    print(img_resize_float.dtype)  # 输出: <dtype: 'float32'>

    # 转换类型为uint8
    # 特别注意convert_image_dtype中float32类型的图像都必须归一化至0~1范围
    # img_resize_float/255用于归一化至0~1的范围
    img_resize_uint = tf.image.convert_image_dtype(img_resize_float/255, dtype=tf.uint8, saturate=True)

    # 用Session完成tensor的计算
    with tf.Session() as sess:
        img_rgb_float = sess.run(img_resize_float)
        print(type(img_rgb_float))  # 输出: <class 'numpy.ndarray'>
        print(img_rgb_float.shape)  # 输出: (270, 480, 3)
        print(img_rgb_float.dtype)  # 输出: float32

        img_rgb = sess.run(img_resize_uint)
        print(type(img_rgb))  # 输出: <class 'numpy.ndarray'>
        print(img_rgb.shape)  # 输出: (270, 480, 3)
        print(img_rgb.dtype)  # 输出: uint8

        # 用OpenCV显示生成的图像
        img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)

        cv2.imshow("float rgb", img_rgb_float/255)  # imshow显示float类型图像，必须归一化至0~1范围
        cv2.imshow("bgr", img_bgr)
        cv2.imshow("rgb", img_rgb)
        cv2.waitKey(0)

上述代码是显示图像如下：

TFRecord文件的数据读写方法

TFRecord的写入

import tensorflow as tf
"""
TFRecord的写入主要在于以下各类
tf.python_io.TFRecordWriter: tf.python_io.TFRecordWriter.write()执行写操作
tf.train.Example: tf.train.Example.SerializeToString()生成写操作对象
tf.train.Features: 包含tf.train.Feature
tf.train.Feature: 建立tf.train.Example使用的feature
tf.train.BytesList: 建立字符串列表的feature
tf.train.FloatList: 建立浮点数列表的feature
tf.train.Int64List: 建立整型数列表的feature
"""

# tf.train.Feature支持的3种list类型
# 字符串、浮点数、int64整型数
data_bytes = b"\x11\x22\x33\x44"  # b'\x11"3D'
data_float = 3.2
data_int64 = 0

# 建立5个TFRecord文件保存数据
for i in range(5):
    filename = "d:/temp/test_folder/rec{0}.tfrecords".format(i)
    with tf.python_io.TFRecordWriter(filename) as writer:  # 用with而不是调用writer.close()
        example = tf.train.Example(features=tf.train.Features(
            feature={
   
                # feature参数输入为字典，key为字符串，值为tf.train.Feature对象
                # tf.train.Feature三种list分别传入不同类型的list
                # bytes_list=tf.train.BytesList(value=字符串列表)
                # float_list=tf.train.FloatList(value=浮点数列表)
                # int64_list=tf.train.Int64List(value=int64整型数列表)
                "bytes": tf.train.Feature(bytes_list=tf.train.BytesList(value=[data_bytes])),
                "float": tf.train.Feature(float_list=tf.train