tensorflow2下csv转tfrecord格式的python代码

最新推荐文章于 2022-06-01 19:50:33 发布

Ben Dan

最新推荐文章于 2022-06-01 19:50:33 发布

阅读量589

点赞数 1

分类专栏：深度学习文章标签： python tensorflow csv 深度学习机器学习

本文链接：https://blog.csdn.net/weixin_44574507/article/details/118107190

版权

该博客介绍了如何在TensorFlow2环境下将CSV格式的标注数据转换为TFRecord，以便于后续的深度学习和机器学习任务。文章提供了一段可修改的Python代码，包括替换图片后缀、CSV文件路径、输出TFRecord文件名和图片路径等步骤，并展示了TFRecord文件的结构。对于多类别标注，只需增加相应的elif语句。

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

将csv格式标注数据转为tfrecord供tensorflow使用

在tensorflow2版本下运行

from __future__ import division
from __future__ import print_function
from __future__ import absolute_import

import os
import io
import pandas as pd
import tensorflow.compat.v1 as tf
from PIL import Image
from object_detection.utils import dataset_util
from collections import namedtuple, OrderedDict

flags = tf.app.flags
flags.DEFINE_string('csv_input', '', 'Path to the CSV input')
flags.DEFINE_string('output_path', '', 'Path to output TFRecord')
FLAGS = flags.FLAGS

def class_text_to_int(row_label):
    if row_label == 'Citrus':
        return 1
    else:
        return 0
    """
    elif row_label == 'Apple':
        return 2
    elif row_label == 'Orange':
        return 3
    elif row_label == 'Watermelon':
        return 4
    """
    
def split(df, group):
    data = namedtuple('data', ['filename', 'object'])
    gb = df.groupby(group)
    return [data(filename, gb.