tensorflow2装甲板id识别 1数据集制作

数据集获取

通过opencv借助灯条信息提取装甲板区域，数据集中包含1到8的装甲板数字图像，原始数据的格式为64*64的RGB图像，数据样例如下。

armor1.png

数据集制作

TFRecords
tfrecords是tendorflow支持的一种标准的数据记录格式，可用于将图像数据和标签组成数据集，对于数据集比较大的情况可以更好的利用内存，所以这里使用该格式制作数据集。
.tfrecords文件制作
tfrecords文件使用protocolbuf协议中的消息体实现标签和数据的组合，在tensorflow2中使用tf.train.Example结构进行实现
装甲板数据集中Example实现如下：

example = tf.train.Example(features = tf.train.Features(
    feature = {
         'shape': tf.train.Feature(int64_list=tf.train.Int64List (value=[int(class_name)])),
           'img_raw': tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_raw]))
    }
 ))

其中value 是 Feature 类型的消息体有3种取值，在本次实现中，使用了tf.train.Int64List记录整形，用tf.train.BytesList记录图片的二进制信息
关于Example的更多信息可以参考
https://blog.csdn.net/qq_42686721/article/details/98205816
.tfrecords文件生成实现如下:

import tensorflow as tf
import os

dataset_path = './armor_dataset'
#要生成的记录文件
writer = tf.io.TFRecordWriter('armor_train.tfrecords')
for class_name in os.listdir(dataset_path):
    class_dataset_path = dataset_path + '/' + class_name
    for img_path in os.listdir(class_dataset_path):
        print(class_dataset_path+'/' +img_path)
        img = tf.io.read_file (class_dataset_path+'/' +img_path)
        img_raw = tf.image.decode_bmp (img)
        example = tf.train.Example(features = tf.train.Features(
            feature = {
                'label': tf.train.Feature(int64_list=tf.train.Int64List (value=[int(class_name)])),
                'img_raw': tf.train.Feature(bytes_list=tf.train.BytesList(value=[bytes(img_raw.numpy())]))
            }
        ))
        writer.write(example.SerializeToString())  #序列化
writer.close()

读取tfrecords文件

#定义待解析数据集特征
feature_description = {
    'label': tf.io.FixedLenFeature([] , tf.int64, default_value=-1), # 默认值自己定义
    'img_raw' : tf.io.FixedLenFeature([], tf.string)
}
# 映射函数，用于解析一条example
def _parse_function (exam_proto): 
    return tf.io.parse_single_example (exam_proto, feature_description)
#读取返回数据集
def read_dataset(record_path):
    reader = tf.data.TFRecordDataset(record_path) # 打开一个TFrecord
    #reader = reader.shuffle (buffer_size = 1000) # 在缓冲区中随机打乱数据
    reader = reader.map (_parse_function) # 解析数据
    labels = []
    for row in reader.take(1): #获取指定数量的数据集
    #for row in reader:                #遍历数据集
        print(row['label'])
        print(row['img_raw'])

打印结果符合预期，此处row['label']，row['img_raw']均为tensor

最后编辑于：2020.04.15 07:18:38