tf.Example msg
tensorflow提供了一种统一的格式.tfrecord来存储图像数据.用的是自家的google protobuf.就是把图像数据序列化成自定义格式的二进制数据.
To read data efficiently it can be helpful to serialize your data and store it in a set of files (100-200MB each) that can each be read linearly. This is especially true if the data is being streamed over a network. This can also be useful for caching any data-preprocessing.
The TFRecord format is a simple format for storing a sequence of binary records.
protobuf消息的格式如下:
https://github.com/tensorflow/tensorflow/blob/r2.0/tensorflow/core/example/feature.proto
message BytesList {
repeated bytes value = 1;
}
message FloatList {
repeated float value = 1 [packed = true];
}
message Int64List {
repeated int64 value = 1 [packed = true];
}
// Containers for non-sequential data.
message Feature {
// Each feature can be exactly one kind.
oneof kind {
BytesList bytes_list = 1;
FloatList float_list = 2;
Int64List int64_list = 3;
}
};
message Features {
map<string, Feature> feature = 1;
};
message FeatureList {
repeated Feature feature = 1;
};
message FeatureLists {
map<string, FeatureList> feature_list = 1;
};
tf.Example是一个map. map的格式为{"string": tf.train.Feature}
tf.train.Feature基本的格式有3种:
- tf.train.BytesList
- string
- byte
- tf.train.FloatList
- float(float32)
- double(float64)
- tf.train.Int64List
- bool
- enum
- int32
- unit32
- int64
- uint64
将自己的数据制作为tfrecord格式
完整代码
from __future__ import absolute_import, division, print_function, unicode_literals
import tensorflow as tf
import numpy as np
import IPython.display as display
import os
import cv2 as cv
import argparse
def _bytes_feature(value):
"""Returns a bytes_list from a string / byte."""
if isinstance(value, type(tf.constant(0))):
value = value.numpy() # BytesList won't unpack a string from an EagerTensor.
return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
def _float_feature(value):
"""Returns a float_list from a float / double."""
return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))
def _int64_feature(value):
"""Returns an int64_list from a bool / enum / int / uint."""
return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
def convert_to_tfexample(img_data,label,height=320,width=320):
"""convert one img matrix into tf.Example"""
image_format = 'png'
example = tf.train.Example(features=tf.train.Features(feature={
'image/encoded': _bytes_feature(img_data),
'image/format': _bytes_feature(tf.compat.as_bytes(image_format)),
'image/class/label': _int64_feature(label),
'image/height': _int64_feature(height),
'image/width': _int64_feature(width),
}))
return example
#path="/home/sc/disk/data/lishui/1"
def read_dataset(path):
imgs=[]
labels=[]
for root, dirs, files in os.walk(path):
for one_file in files:
#print(os.path.join(path,one_file))
one_file = os.path.join(path,one_file)
if one_file.endswith("png"):
label_file = one_file.replace('png','txt')
if not os.path.isfile(label_file):
continue
f = open(label_file)
class_index = int(f.readline().split(' ')[0])
labels.append(class_index)
img = tf.gfile.GFile(one_file, 'rb').read()
print(type(img))
imgs.append(img)
return imgs,labels
def arg_parse():
parser = argparse.ArgumentParser()
#parser.add_argument('--help',help='ex:python create_tfrecord.py -d /home/sc/disk/data/lishui/1 -o train.tfrecord')
parser.add_argument('-d','--dir',type=str,default='./data',required='True',help='dir store images/label file')
parser.add_argument('-o','--output',type=str,default='./outdata.tfrecord',required='True',help='output tfrecord file name')
args = parser.parse_args()
return args
def main():
args = arg_parse()
writer = tf.io.TFRecordWriter(args.output)
imgs,labels = read_dataset(args.dir)
examples = map(convert_to_tfexample,imgs,labels)
for example in examples:
writer.write(example.SerializeToString())
writer.close()
print("write done")
if __name__ == '__main__':
"""
usage:python create_tfrecord.py [data_path] [outrecordfile_path]
ex:python create_tfrecord.py -d /home/sc/disk/data/lishui/1 -o train.tfrecord
"""
main()
首先就是需要有工具函数把byte/string/float/int..等等类型的数据转换为tf.train.Feature
def _bytes_feature(value):
"""Returns a bytes_list from a string / byte."""
if isinstance(value, type(tf.constant(0))):
value = value.numpy() # BytesList won't unpack a string from an EagerTensor.
return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
def _float_feature(value):
"""Returns a float_list from a float / double."""
return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))
def _int64_feature(value):
"""Returns an int64_list from a bool / enum / int / uint."""
return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))
接下来,对于图片矩阵和标签数据,我们调用上述工具函数,将单幅图片及其标签信息转换为tf.ttrain.Example消息.
def convert_to_tfexample(img,label):
"""convert one img matrix into tf.Example"""
img_raw = img.tostring()
example = tf.train.Example(features=tf.train.Features(feature={
'label': _int64_feature(label),
'img': _bytes_feature(img_raw)}))
return example
对于我的数据,图片以及label文件位于同一目录.比如dir下有图片a.png及相应的标签信息a.txt.
def read_dataset(path):
imgs=[]
labels=[]
for root, dirs, files in os.walk(path):
for one_file in files:
#print(os.path.join(path,one_file))
one_file = os.path.join(path,one_file)
if one_file.endswith("png"):
label_file = one_file.replace('png','txt')
if not os.path.isfile(label_file):
continue
f = open(label_file)
class_index = int(f.readline().split(' ')[0])
labels.append(class_index)
img = tf.gfile.GFile(one_file, 'rb').read()
print(type(img))
imgs.append(img)
return imgs,labels
遍历data目录,完成图片读取,及label读取. 如果你的数据不是这么存放的,就修改这个函数好了,返回值仍然是imgs,labels
最后就是调用 tf.io.TFRecordWriter将每一个tf.train.Example消息写入文件保存.
def main():
args = arg_parse()
writer = tf.io.TFRecordWriter(args.output)
#path="/home/sc/disk/data/lishui/1"
imgs,labels = read_dataset(args.dir)
examples = map(convert_to_tfexample,imgs,labels)
for example in examples:
writer.write(example.SerializeToString())
writer.close()
print("write done")