TensorFlow Example

发表于 | 阅读次数

测试环境：Windows 10 ; TensorFlow 1.1

测试样例：输入是30张128*128 jpg格式图片，大小96.4KB；标签是[0:30]

第一种：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
def generate_TFrecord(files_list,labels_list,to_file="train.tfrecords"):
    """
    生成一个TFrecord文件
    :param files_list: 文件列表，例如 ['1.jpg','/dataset/2.jpg']
    :param labels_list: 和文件列表对应的标签列表,例如[0,4]
    :param to_file: 写入目标TFrecord文件
    :return: NULL
    """
    writer = tf.python_io.TFRecordWriter(to_file)
    for (file,label) in zip(files_list,labels_list):
        img = Image.open(file)
        img_raw = img.tobytes()  #转成byte
        example = tf.train.Example(features=tf.train.Features(
            feature={
                "label": tf.train.Feature(int64_list=tf.train.Int64List(value=[label])),
                'img_raw': tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_raw]))
            }))
        writer.write(example.SerializeToString())#写入TFrecord文件
    writer.close()

生成TFrecord大小：1.40MB

第二种：

直接用数组形式存下每个像素的值。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
def generate_TFrecord2(files_list,labels_list,to_file="train.tfrecords"):
    writer = tf.python_io.TFRecordWriter(to_file)
    for (file,label) in zip(files_list,labels_list):
        img = Image.open(file)
        #img_raw = img.tobytes()  #转成byte
        img_raw = (numpy.array(img).reshape(-1).tolist())
        example = tf.train.Example(features=tf.train.Features(
            feature={
                "label": tf.train.Feature(int64_list=tf.train.Int64List(value=[label])),
                'img_raw': tf.train.Feature(int64_list=tf.train.Int64List(value=img_raw))
            }))
        writer.write(example.SerializeToString())#写入TFrecord文件
    writer.close()

测试结果：1.79MB

第三种

使用tf.train.SequenceExample代替tf.train.Example

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
def generate_TFrecord_SequenceExample(files_list,labels_list,to_file="train.tfrecords"):
    writer = tf.python_io.TFRecordWriter(to_file)
    for (file,label) in zip(files_list,labels_list):
        img = Image.open(file)
        img_raw = img.tobytes()  #转成byte
        feature_input=tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_raw]))
        feature_label= tf.train.Feature(int64_list=tf.train.Int64List(value=[label]))
        feature_list_input= tf.train.FeatureList(feature=[feature_input])
        feature_list_label=tf.train.FeatureList(feature=[feature_label])
        feature_list = {
            'imgs': feature_list_input,
            'labels': feature_list_label
        }
        feature_lists = tf.train.FeatureLists(feature_list=feature_list)
        example= tf.train.SequenceExample(feature_lists=feature_lists)
        #以上语句等同于下面3行。
        # example=tf.train.SequenceExample()
        # example.feature_lists.feature_list["imgs"].feature.add().bytes_list.value.append(img_raw)
        # example.feature_lists.feature_list["labels"].feature.add().int64_list.value.append(label)
        writer.write(example.SerializeToString())#写入TFrecord文件
    writer.close()

测试结果：1.40MB