Tensorflow --- creation and reading of tfrecord file

One 、 Environmental requirements

Python==3.7
Tensorflow==2.4.0

Two 、 Data preparation

The data set demonstrated in this blog can be downloaded through the following link
Baidu online disk extraction code ：9q0f

3、 ... and 、TFRecord File creation

import os
import cv2
import tensorflow as tf
import tqdm as tqdm

if __name__ == '__main__':
    #  Root directory of data set to be processed 
    Dataset_path = 'Split_Dataset'
    #  Used to generate the file path to be written 
    file_names = [i.split('.')[0] for i in os.listdir(os.path.join(Dataset_path, 'DSM'))]
    dsm_filename = [os.path.join(Dataset_path, 'DSM', i + '.tif') for i in file_names]
    label_filename = [os.path.join(Dataset_path, 'Label', i + '.png') for i in file_names]
    rgb_filename = [os.path.join(Dataset_path, 'RGB', i + '.png') for i in file_names]
    #  To begin TFRecords File construction 
    with tf.io.TFRecordWriter('Potsdam.tfrecords') as writer:
        tqdm_file = tqdm.tqdm(iterable=zip(dsm_filename, rgb_filename, label_filename), total=len(dsm_filename))
        for dsm, rgb, label in tqdm_file:
            #  Read the elevation map 
            dsm_data = cv2.imread(dsm, -1)
            dsm_data = dsm_data.ravel()

            #  Read the optical map 
            rgb_data = open(rgb, 'rb').read()

            #  Read the label 
            label_data = open(label, 'rb').read()

            #  establish tf.train.Feature Dictionaries 
            feature = {
    
                'dsm': tf.train.Feature(float_list=tf.train.FloatList(value=dsm_data)),
                'rgb': tf.train.Feature(bytes_list=tf.train.BytesList(value=[rgb_data])),
                'label': tf.train.Feature(bytes_list=tf.train.BytesList(value=[label_data]))
            }

            #  Create... From a dictionary Example
            example = tf.train.Example(features=tf.train.Features(feature=feature))
            #  take Example Serialize and write TFRecords file 
            writer.write(example.SerializeToString())
        tqdm_file.close()

Four 、TFRecord File reading

import tensorflow as tf
import matplotlib.pyplot as plt

#  structure Feature structure , Tell the decoder everyone Feature What is it? 
feature_description = {
    
    'dsm': tf.io.FixedLenFeature([512, 512, 1], tf.float32),
    'rgb': tf.io.FixedLenFeature([], tf.string),
    'label': tf.io.FixedLenFeature([], tf.string)
}


# Example Analytic function of 
def parse_example(example_string):
    feature = tf.io.parse_single_example(serialized=example_string, features=feature_description)

    feature['rgb'] = tf.image.decode_png(feature['rgb'], channels=3)
    feature['rgb'] = tf.image.resize(feature['rgb'], [512, 512])
    feature['rgb'] = tf.cast(feature['rgb'], tf.float32)
    feature['rgb'] = feature['rgb'] / 255

    feature['label'] = tf.image.decode_png(feature['label'], channels=1)
    feature['label'] = tf.image.resize(feature['label'], [512, 512])
    feature['label'] = tf.cast(feature['label'], tf.int64)
    feature['label'] = tf.squeeze(feature['label'])

    return feature['dsm'], feature['rgb'], feature['label']


if __name__ == '__main__':
    dataset = tf.data.TFRecordDataset(r'Potsdam.tfrecords')
    dataset = dataset.map(parse_example)
    dataset = dataset.shuffle(buffer_size=500)

    for dsm, rgb, label in dataset.take(1):
        temp = (dsm - tf.reduce_min(input_tensor=dsm, axis=(0, 1, 2))) / (
                tf.reduce_max(input_tensor=dsm, axis=(0, 1, 2)) - tf.reduce_min(input_tensor=dsm, axis=(0, 1, 2)))

        plt.figure(figsize=(15, 5), dpi=150)
        plt.subplot(1, 3, 1)
        plt.imshow(temp)
        plt.axis('off')

        plt.subplot(1, 3, 2)
        plt.imshow(rgb)
        plt.axis('off')

        plt.subplot(1, 3, 3)
        plt.imshow(label)
        plt.axis('off')

        plt.show()

5、 ... and 、TFRecord File reading results

6、 ... and 、 Downloading project files

The project source file of this article can be obtained by clicking the following link
Project source file