TFRecords数据的生成与读取

布满荆棘的人生 2022-05-29 08:06 173阅读 0赞

TFRecords 作为TensorFlow标准支持格式，将所有信息（包括图片信息）写入到一个tfrecords文件中，便于管理数据。TFRecords 是二进制文件，有特定的写入和读取方式。

需要将下图中的image图片文件和anno.txt文件生成为tfrecords格式。

![Image 1][]

anno.txt中标注数据如下：

E:/anno/image/Drivingdecordr_01.jpg 1 648 512 793 549
    E:/anno/image/Drivingdecordr_02.jpg 1 1330 446 1388 477
    E:/anno/image/Drivingdecordr_03.jpg 1 26 509 95 543
    E:/anno/image/Drivingdecordr_04.jpg 1 1430 437 1503 466
    E:/anno/image/Drivingdecordr_05.jpg 1 1582 521 1714 567
    E:/anno/image/Drivingdecordr_06.jpg 1 888 502 1053 547
    E:/anno/image/Drivingdecordr_07.jpg 1 1419 453 1493 478
    E:/anno/image/Drivingdecordr_08.jpg 1 31 495 102 530
    E:/anno/image/Drivingdecordr_09.jpg 1 862 471 989 508

以下代码中trans2tfrecords()将标注文件转换为tfrecords格式数据，read\_tfrecords()为训练模型读取tfrecords格式数据。

# -*- coding: utf-8 -*-
    
    import tensorflow as tf
    import cv2
    import sys
    import random
    
    def get_data(data_dir):
        imagelist = open(data_dir, 'r')
        dataset = []
        for line in imagelist.readlines():
            info = line.strip().split(' ')
            data_example = dict()
            box = dict()
            data_example['filename'] = info[0]
            data_example['label'] = int(info[1])
            box['xmin'] = float(info[2])
            box['ymin'] = float(info[3])
            box['xmax'] = float(info[4])
            box['ymax'] = float(info[5])
            data_example['box'] = box
            dataset.append(data_example)
        return dataset
    
    def _int64_feature(value):
        """Wrapper for insert int64 feature into Example proto."""
        if not isinstance(value, list):
            value = [value]
        return tf.train.Feature(int64_list=tf.train.Int64List(value=value))
    
    
    def _float_feature(value):
        """Wrapper for insert float features into Example proto."""
        if not isinstance(value, list):
            value = [value]
        return tf.train.Feature(float_list=tf.train.FloatList(value=value))
    
    
    def _bytes_feature(value):
        """Wrapper for insert bytes features into Example proto."""
        if not isinstance(value, list):
            value = [value]
        return tf.train.Feature(bytes_list=tf.train.BytesList(value=value))
    
    def trans2tfrecords(tf_filename, dataset):
    	with tf.python_io.TFRecordWriter(tf_filename) as writer:
    		for i, image_example in enumerate(dataset):
    			sys.stdout.write('\r>> Converting image %d/%d\n' % (i + 1, len(dataset)))
                #sys.stdout.flush()
    			filename = image_example['filename']
    			image_data = extract(filename)
    			class_label = image_example['label']
    			box = image_example['box']
    			roi = [box['xmin'], box['ymin'], box['xmax'], box['ymax']]
    			#tf.train.Example协议内存块，包含字段Features,Features中包含Feature的字典
    			example = tf.train.Example(features = tf.train.Features(feature={     
    				'image': _bytes_feature(image_data),
    				'label': _int64_feature(class_label),
    				'roi': _float_feature(roi)
    				}))
    				
    			writer.write(example.SerializeToString())
    	
    def extract(filename):
    	image = cv2.imread(filename)
    	image_data = image.tostring()   #将numpy类转化为string类
    	return image_data
    	
    def read_tfrecords(tfrecord_file):
        filename_queue = tf.train.string_input_producer([tfrecord_file],shuffle=True) #生成队列，并随机打乱顺序
        reader = tf.TFRecordReader()    #文件读取器
        _, serialized_example = reader.read(filename_queue)
        image_features = tf.parse_single_example(              #解析器
                serialized_example,
                features = {
                        'image': tf.FixedLenFeature([], tf.string),
                        'label':tf.FixedLenFeature([], tf.int64),
                        'roi': tf.FixedLenFeature([4], tf.float32)
                        })
        image = tf.decode_raw(image_features['image'], tf.uint8)      #解码器
        image = tf.reshape(image, [128, 128, 3])
        image = (tf.cast(image, tf.float32)-127.5) / 128
        
        label = tf.cast(image_features['label'], tf.float32)
        roi = tf.cast(image_features['roi'],tf.float32)
        return image, label, roi
    	
    if __name__ == '__main__':
        data_dir = 'E:/anno/anno.txt'
        imagelist = open(data_dir, 'r')
    	
        output_dir = 'E:/anno'
        tf_filename = '%s/train_data.tfrecord' %output_dir
    	
        shuffling = 'True'
    	
        dataset = get_data(data_dir)
        if shuffling:
            random.shuffle(dataset)
    	
        trans2tfrecords(tf_filename, dataset)
        print('Finished converting dataset!')
    	
        image, label, roi = read_tfrecords(tf_filename)

[Image 1]: