如何基于TFRecord与.pbtxt文件构建带标签的训练数据集？

1. 读取PBTXT标签映射 #

PBTXT格式的标签通常是类似这样的结构（以目标检测或分类场景为例）：

item {
  id: 1
  name: "cat"
}
item {
  id: 2
  name: "dog"
}

你需要先解析这个文件，生成标签ID与名称的映射字典，有两种实现方式：

方式1：基于TensorFlow Object Detection API（官方标准格式） #

如果你的PBTXT是官方目标检测API的格式，可以用protobuf工具解析：

import tensorflow as tf
from google.protobuf import text_format
from object_detection.protos import string_int_label_map_pb2

def load_label_map(label_map_path):
    label_map = string_int_label_map_pb2.StringIntLabelMap()
    with open(label_map_path, 'r') as f:
        text_format.Merge(f.read(), label_map)
    
    id_to_name = {}
    for item in label_map.item:
        id_to_name[item.id] = item.name
    
    # 可选：生成名称到ID的反向映射
    name_to_id = {v: k for k, v in id_to_name.items()}
    return id_to_name, name_to_id

方式2：手动解析（自定义简单PBTXT） #

如果是自定义的简化版PBTXT，直接文本解析更灵活：

def load_simple_label_map(label_map_path):
    id_to_name = {}
    with open(label_map_path, 'r') as f:
        lines = f.readlines()
        current_id, current_name = None, None
        for line in lines:
            line = line.strip()
            if line.startswith('id:'):
                current_id = int(line.split(':')[1].strip())
            elif line.startswith('name:'):
                current_name = line.split(':')[1].strip().strip('"')
                if current_id is not None:
                    id_to_name[current_id] = current_name
                    current_id, current_name = None, None
    return id_to_name

2. 解析TFRecord中的样本 #

TFRecord里的每个样本是序列化的tf.train.Example，需要定义解析函数提取图像和标签ID（假设你的TFRecord存储了image/encoded（图像二进制）和image/class/label（标签ID）字段，可根据实际存储调整）：

def parse_tfrecord_example(example_proto, id_to_name=None):
    # 定义TFRecord的特征描述
    feature_description = {
        'image/encoded': tf.io.FixedLenFeature([], tf.string),
        'image/class/label': tf.io.FixedLenFeature([], tf.int64),
        # 可添加其他字段：如image/height、image/width等
    }
    
    # 解析单个样本
    parsed_features = tf.io.parse_single_example(example_proto, feature_description)
    
    # 解码图像（JPEG/PNG格式通用）
    image = tf.io.decode_jpeg(parsed_features['image/encoded'], channels=3)
    # 可选：统一图像尺寸
    image = tf.image.resize(image, (224, 224))
    
    # 获取标签ID
    label_id = parsed_features['image/class/label']
    
    # 若需要标签名称，用映射表转换（TensorFlow图模式下要用StaticHashTable）
    if id_to_name is not None:
        keys = tf.constant(list(id_to_name.keys()), dtype=tf.int64)
        values = tf.constant(list(id_to_name.values()), dtype=tf.string)
        label_table = tf.lookup.StaticHashTable(
            tf.lookup.KeyValueTensorInitializer(keys, values),
            default_value='unknown'
        )
        label_name = label_table.lookup(label_id)
        return image, label_id, label_name
    else:
        # 仅返回图像和标签ID，直接用于训练
        return image, label_id

3. 整合数据集并训练 #

把读取TFRecord、解析样本、标签映射整合到数据管道，即可直接用于模型训练：

# 1. 加载标签映射
label_map_path = 'path/to/your/labels.pbtxt'
id_to_name, _ = load_label_map(label_map_path)  # 或用load_simple_label_map

# 2. 读取TFRecord文件
filenames = ['path/to/your/dataset.tfrecord']
raw_dataset = tf.data.TFRecordDataset(filenames)

# 3. 解析样本并关联标签
# 若不需要标签名称，去掉id_to_name参数即可
parsed_dataset = raw_dataset.map(lambda x: parse_tfrecord_example(x, id_to_name))

# 4. 数据集预处理（洗牌、批量、预取，提升训练效率）
batch_size = 32
train_dataset = parsed_dataset.shuffle(1000).batch(batch_size).prefetch(tf.data.AUTOTUNE)

# 5. 示例训练流程
model = tf.keras.applications.ResNet50(weights=None, input_shape=(224,224,3), classes=len(id_to_name))
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model.fit(train_dataset, epochs=10)

关键注意事项 #

• 确保TFRecord中存储的标签ID与PBTXT中的ID完全对应，否则会出现标签不匹配问题。
• 如果TFRecord中没有存储标签ID，需要额外建立样本（如文件名）与标签的映射关系，再关联PBTXT中的标签。
• 大规模数据集下，必须用tf.lookup.StaticHashTable做标签转换，避免Python字典在TensorFlow图模式下的兼容性问题。

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