使用tensorflow将图像数据集制作成tfrecords格式(代码)

最新推荐文章于 2024-06-03 15:05:10 发布

CROD鳄鱼

最新推荐文章于 2024-06-03 15:05:10 发布

阅读量1.2k

点赞数

文章标签： python

本文链接：https://blog.csdn.net/qiuuu/article/details/123872353

版权

本文提供了一段使用tensorflow-gpu 2以上版本的代码，快速将图像数据集转换为tfrecords格式。虽然代码简洁，但未深入解释原理。代码主要功能包括读取文件夹内的图像数据，按类别组织，并将数据转化为tf.train.Example格式。

摘要由CSDN通过智能技术生成

只有几串代码，快速实现，但是原理没有详细说明。

使用tensorflow-gpu 2以上版本：

BATCH_SIZE = 2
train_dir = "C:\\Users\Desktop\泸州老窖精品头曲组合装\\"
train_tfrecord = "C:\\Users\Desktop\\train.tfrecords"
dataset_to_tfrecord(dataset_dir=train_dir, tfrecord_name=train_tfrecord)

自定义函数如下：

def dataset_to_tfrecord(dataset_dir, tfrecord_name):
    image_paths, image_labels = get_images_and_labels(dataset_dir)
    image_paths_and_labels_dict = {}
    for i in range(len(image_paths)):
        image_paths_and_labels_dict[image_paths[i]] = image_labels[i]
    # shuffle the dict
    image_paths_and_labels_dict = shuffle_dict(image_paths_and_labels_dict) # 打乱数据
    with tf.io.TFRecordWriter(path=tfrecord_name) as writer:
        for image_path, label in image_paths_and_labels_dict.items():
            print("Writing to tfrecord: {}".format(image_path))
            image_string = open(