将数组写成tfrecord

Posted 2021-01-23 rongye

#!/usr/bin/env python


import tensorflow as tf
import numpy as np
# write = tf.python_io.TFRecordWriter(‘train.tfrecords‘)
# #img_raw = np.random.random_integers(0, 255, size=(7, 30))  # 创建7*30，取值在0-255之间随机数组
# list_a=[6,5,4]
# list_b=[7]
# img_raw = np.array(list_a,dtype=np.float32)
# img_label =np.array(list_b,dtype=np.int64)
# #img_label = img_label.tostring()
# img_raw = img_raw.tostring()
# example = tf.train.Example(features=tf.train.Features(
#     feature={
#         ‘label‘: tf.train.Feature(int64_list=tf.train.Int64List(value=[img_label])),
#         ‘img_raw‘: tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_raw]))
#     }))
# write.write(example.SerializeToString())
# write.close()
def create_file(path):
    write = tf.python_io.TFRecordWriter(‘train02.tfrecords‘)
    with open(path,‘r‘) as file:
        lines = file.readlines()
    # print lines.__len__()
        count = 0
        data = []
        featuresList = []
        labelList = []
        label = []
        for line in lines:
            word = line.split(" ")
            features = []
            for i in range(1, len(word)):
                if i < (len(word) - 1):
                    features.append(word[i].split(":")[1])
                else:
                    features.append(word[len(word) - 1].split(":")[1].split("
")[0])
            label.append(int(word[0]))
            count = count + 1
            #print(count)
            featuresList.append(features)
            labelList.append(label)
            img_raw = np.array(featuresList[-1], dtype=np.float32)
            img_label = np.array(label[-1],dtype=np.int64)
            print(img_label)
            print(img_raw)
            img_raw = img_raw.tostring()
            example = tf.train.Example(features=tf.train.Features(
                feature={
                    ‘label‘: tf.train.Feature(int64_list=tf.train.Int64List(value=[img_label])),
                    ‘img_raw‘: tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_raw]))
                }))
            write.write(example.SerializeToString())
        write.close()
create_file("train02.txt")