基于Tensorflow框架的tfrecord文件的生成与读取（附有详细的注释，易理解）

原创

于 2019-04-22 16:05:43 发布 · 置顶 · 1.1k 阅读

3 ·

CC 4.0 BY-SA版权

文章标签：

#tfrecord #tensorflow #python

本文详细介绍了如何使用Tensorflow框架生成和读取tfrecord文件，提供了带注释的示例代码，适用于猫狗分类数据集。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

基于Tensorflow框架的tfrecord文件的生成与读取

一. 根据自己已有数据集，生成tfrecord文件，附有详细的注释：

我所用的数据集只有猫狗两类，下面是部分示例图：
在这里插入图片描述

#coding=utf-8
import cv2 
import os
import tensorflow as tf
import tensorflow.contrib.slim as slim
import numpy as np

def writeTFRecord(imgDir, recordPath, imgH, imgW, encoderLabel, recordName):
    '''
         创建tfrecord文件
    param:
        imgDir: 预处理图片路径
        recordPath: 预保存tfrecord文件路径
        imgH: 为了方便训练，将原图resize为指定的高
        imgW: 指定的宽
        encoderLabel: 编码标签，将标签以数字的形式呈现
        recordName: 预保存tfrecord文件名，eg:test.tfrecord
    '''
    imgPathList = [os.path.join(imgDir, img) for img in os.listdir(imgDir)]  #将所有图片路径绝对路径放入list中
    np.random.shuffle(imgPathList)  #打乱imgList中的图片顺序
    
    if not os.path.exists(recordPath):
        os.makedirs(recordPath)
    writer = tf.python_io.TF