如何生成FSNS数据集结构的tfrecord数据

最新推荐文章于 2022-03-18 19:09:05 发布

qian99

最新推荐文章于 2022-03-18 19:09:05 发布

阅读量1.9k

点赞数

CC 4.0 BY-SA版权

分类专栏：人工智能深度学习

本文链接：https://blog.youkuaiyun.com/qian99/article/details/79051779

人工智能深度学习专栏收录该内容

32 篇文章

订阅专栏

本文介绍如何使用AttentionOCR进行光学字符识别，并详细解释了如何将自定义数据集转换为符合AttentionOCR要求的TFRecord格式。文章还提供了一段关键代码示例，帮助读者理解图像数据和标签数据的编码过程。

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

想要用一下attention ocr，github如下：
https://github.com/tensorflow/models/tree/master/research/attention_ocr

发现这里面并没有生成数据的样例代码，只能用fsns的格式，如果想要用自己的数据集要生成一个和FSNS一样结构的tfrecord数据，但是这部分没有仔细说，只给了一个stackoverflow的链接，更坑的是，这个链接里说的也不清楚，而且还有些错误：
https://stackoverflow.com/a/44461910/743658

FSNS的具体结构在这篇论文里有说：
https://arxiv.org/pdf/1702.03970.pdf

不过我们只需要关心这个图就行了。
这里写图片描述

image/format : 表示图像数据的格式，官方使用的png格式
image/encoded : 表示图像数据转码后的字符串数据
image/class : 长度为37的label list（37是label最长的长度），其实就是真实label的每个字符的id的list，后面长度不足37就添上null
image/unpadded_class : 没有填null的原始label list
image/width : 图像的宽度
image/orig-width : 切分图像的宽度。因为fsns数据集有4个视角，所以是图像宽度除4，如果只有一个视角，和上面宽度一样就行了
image/height : 图像高度，实际使用的时候高度数据没有使用，所以这个也可以不填
image/text : 原始的label字符串

下面贴一下关键代码：

def encode_utf8_string(text, length, charset, null_char_id=133):
    char_ids_padded = [null_char_id]*length
    char_ids_unpadded = [null_char_id]*len(text)
    for i in range(len(text)):
        hash_id = charset[text[i]]
        char_ids_padded[i] = hash_id
        char_ids_unpadded[i] = hash_id
    return char_ids_padded, char_ids_unpadded

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

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

'''生成tfrecord'''
img_data = img_data.resize((IMAGE_SIZE,IMAGE_SIZE),Image.ANTIALIAS)
np_data = np.array(img_data)
image = tf.image.convert_image_dtype(np_data, dtype=tf.uint8)  
image = tf.image.encode_png(image)
image_data = sess.run(image)
char_ids_padded, char_ids_unpadded = encode_utf8_string(
    text=text,
    charset = charset,
    length=MAX_LABEL_LENGTH,
    null_char_id=38
)
example = tf.train.Example(features=tf.train.Features(
  feature={
    'image/encoded': _bytes_feature(image_data),
    'image/format': _bytes_feature("PNG"),
    'image/width': _int64_feature([np_data.shape[1]]),
    'image/orig_width': _int64_feature([np_data.shape[1]]),
    'image/class': _int64_feature(char_ids_padded),
    'image/unpadded_class': _int64_feature(char_ids_unpadded),
    'image/text': _bytes_feature(str(text)),
    # 'height': _int64_feature([crop_data.shape[0]]),
  }
))
tfrecord_writer.write(example.SerializeToString())

大概就是这样了，因为中间要把numpy的矩阵转为tensorflow的png格式，还要再跑个session，之前按照stackoverflow里的写法，直接用numpy的格式，一直数据读取错误……我看了下源码，其实是可以不用png格式的：
这里写图片描述
只需要把’image/format’的值改成’raw’就行了。