如何修改代码将JPEG文件夹转为单图对应独立TFRecord文件

单图对应独立TFRecord的代码修改 #

原代码的核心逻辑是创建1个TFRecordWriter，循环将所有图像写入同一个文件，只需调整写入逻辑：每处理1张图像就新建独立的写入器，写完单张图像后立刻关闭写入流即可实现单图单文件。
另外原代码存在一个适配bug：将图像通道数depth硬编码为1，而你当前的数据集是3通道RGB图像（输出shape为(300,256,256,3)），需要同步修正这个参数避免后续读取错误。
修改后的可直接使用的转换函数如下：

import tensorflow as tf
import os
import numpy as np
from tqdm import tqdm
import skimage.io
import random

def _int64_feature(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))

def _bytes_feature(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

def convert_to_single_tfrecord(images, labels, output_directory, name_prefix, file_name_list=None):
    num_examples = labels.shape[0]
    if images.shape[0] != num_examples:
        raise ValueError("Images size %d does not match label size %d." %
                         (images.shape[0], num_examples))
    rows = images.shape[1]
    cols = images.shape[2]
    # 自动适配通道数，替换原硬编码depth=1的逻辑
    depth = images.shape[3] if len(images.shape) == 4 else 1
    
    # 提前创建输出文件夹
    os.makedirs(output_directory, exist_ok=True)

    for index in tqdm(range(num_examples)):
        # 生成独立TFRecord文件名，传入原文件名列表时可和原JPEG文件名对齐
        if file_name_list is not None:
            base_name = os.path.splitext(file_name_list[index])[0]
            filename = os.path.join(output_directory, f"{base_name}.tfrecords")
        else:
            filename = os.path.join(output_directory, f"{name_prefix}_{index}.tfrecords")
        
        # 每个文件单独创建写入器
        writer = tf.io.TFRecordWriter(filename)
        # TF1版本可将上行替换为 writer = tf.python_io.TFRecordWriter(filename)
        image_raw = images[index].tobytes()
        example = tf.train.Example(features=tf.train.Features(feature={
            'height': _int64_feature(rows),
            'width': _int64_feature(cols),
            'depth': _int64_feature(depth),
            'label': _int64_feature(int(labels[index])),
            'image_raw': _bytes_feature(image_raw)}))
        writer.write(example.SerializeToString())
        # 单张图写完立刻关闭写入流
        writer.close()

调用时替换原有的convert_to函数即可，如果需要输出的TFRecord文件名和原JPEG文件名一一对应，把存了所有原文件名的image_list传入即可：

# 原有图像读取逻辑不变
def read_image(file_name, images_path):
    image = skimage.io.imread(images_path + file_name)
    return image

def extract_image_index_make_label(img_name):
    label = random.randint(1,300)
    return label

images_path = "/content/monet_jpg/"
image_list = os.listdir(images_path)
images = []
labels = []
for img_name in tqdm(image_list):
    images.append(read_image(img_name, images_path))
    labels.append(extract_image_index_make_label(img_name))
images_array = np.array(images)
labels = np.array(labels)

# 替换原有调用，输出文件存到./monet_tfrecord文件夹，文件名和原jpg对齐
convert_to_single_tfrecord(images_array, labels, "./monet_tfrecord", "monet", file_name_list=image_list)

如果后续处理的图像数量更大，建议去掉提前把所有图像读入images_array的逻辑，改成遍历图片时边读边写，可以大幅降低内存占用。

TFRecord文件夹对比JPEG文件夹的效率说明 #

• 适配成本最低：你的启动代码已经基于TFRecord格式编写，直接输出符合要求的单文件TFRecord不需要再修改输入流水线逻辑，不会出现格式不兼容问题。
• 读写效率更高：TFRecord是为TensorFlow输入流水线优化的二进制序列化格式，配合tf.data的并行读取、预取、缓存机制时，比零散存储的JPEG文件IO开销更低。零散JPEG文件在读取时需要频繁做文件系统元数据查询、逐文件打开关闭，在机械硬盘、分布式存储场景下额外开销尤其明显。
• 数据一致性更好：转换时将图像、标签统一序列化存储，不会出现训练过程中文件损坏、标签和图像错配、文件名乱序导致的标注错误问题。
• 注意：单图单TFRecord的模式相比将多张图打包为少量TFRecord文件，还是存在小文件IO开销，但如果你的启动代码强制要求单图单文件，这个方案完全满足需求，整体效率仍然优于直接读取JPEG文件夹。

内容的提问来源于stack exchange，提问作者Gazoo