验证码识别2(使用整体训练法)——验证码数据转换为tfrecord格式

上一篇文章我们生成了验证码,而我们知道一张图片的直接喂入到网络中会使得训练速度变得很慢,这时候我们就需要把图片转换成其他方便读取快速的格式,tfrecord格式是Google提供的一种格式,速度很快。那么接下来,我们需要把图片转换格式了。

代码如下,你需要修改的是图片给出的路径,还有tfrecord文件存放的路径。

import tensorflow as tf
import os
import random
import math
import sys
from PIL import Image
import numpy as np

# %%
# 验证集数量
_NUM_TEST = 500

# 随机种子
_RANDOM_SEED = 0

# 数据集路径
DATASET_DIR = "D:/验证码识别/images/"

# tfrecord文件存放路径
TFRECORD_DIR = "D:/验证码识别/使用一个标签来训练/"


# 判断tfrecord文件是否存在
def _dataset_exists(dataset_dir):
    for split_name in ['train', 'test']:
        output_filename = os.path.join(dataset_dir, split_name + '.tfrecords')
        if not tf.gfile.Exists(output_filename):
            return False
    return True


# 获取所有验证码图片
def _get_filenames_and_classes(dataset_dir):
    photo_filenames = []
    for filename in os.listdir(dataset_dir):
        # 获取文件路径
        path = os.path.join(dataset_dir, filename)
        photo_filenames.append(path)
    return photo_filenames


def int64_feature(values):
    if not isinstance(values, (tuple, list)):
        values = [values]
    return tf.train.Feature(int64_list=tf.train.Int64List(value=values))


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


def image_to_tfexample(image_data, label):
    # Abstract base class for protocol messages.
    return tf.train.Example(features=tf.train.Features(feature={
        'image': bytes_feature(image_data),
        'label': int64_feature(label)
    }))


# 把数据转为TFRecord格式
def _convert_dataset(split_name, filenames, dataset_dir):
    assert split_name in ['train', 'test']

    with tf.Session() as sess:
        # 定义tfrecord文件的路径+名字
        output_filename = os.path.join(TFRECORD_DIR, split_name + '.tfrecords')
        with tf.python_io.TFRecordWriter(output_filename) as tfrecord_writer:
            for i, filename in enumerate(filenames):
                try:
                    sys.stdout.write('\r>> Converting image %d/%d' % (i + 1, len(filenames)))
                    sys.stdout.flush()

                    # 读取图片
                    image_data = Image.open(filename)
                    # 根据模型的结构resize
                    image_data = image_data.resize((224, 224))
                    # 灰度化
                    image_data = np.array(image_data.convert('L'))
                    # 将图片转化为bytes
                    image_data = image_data.tobytes()

                    # 获取label
                    labels = filename.split('/')[-1]
                    num_labels = []
                    num_labels.append(int(labels[0:4]))

                    # 生成protocol数据类型
                    example = image_to_tfexample(image_data, num_labels)
                    tfrecord_writer.write(example.SerializeToString())
                except IOError as e:
                    print('Could not read:', filename)
                    print('Error:', e)
                    print('Skip it\n')
        print(labels)
        print(num_labels)
    sys.stdout.write('\n')
    sys.stdout.flush()


# 判断tfrecord文件是否存在
if _dataset_exists(TFRECORD_DIR):
    print('tfcecord文件已存在')
else:
    # 获得所有图片
    photo_filenames = _get_filenames_and_classes(DATASET_DIR)

    # 把数据切分为训练集和测试集,并打乱
    random.seed(_RANDOM_SEED)
    random.shuffle(photo_filenames)
    training_filenames = photo_filenames[_NUM_TEST:]
    testing_filenames = photo_filenames[:_NUM_TEST]

    # 数据转换
    _convert_dataset('train', training_filenames, DATASET_DIR)
    _convert_dataset('test', testing_filenames, DATASET_DIR)
    print('生成tfcecord文件')

运行程序后在对应的目录下会生成两个文件,如下:

train.tfrecord文件是用来训练的

test.tfrecord文件是用来测试的

 

生成文件的部分就到这里结束了,有兴趣的童鞋如果想看生成的tf.record文件怎么恢复到原来的图片请接着往下看,否则请看下一篇文章验证码识别3——训练部分

我们想看看生成的tfrecord文件是不是正确的,下面我们生成文件来看看,代码如下(你需要在路径里添加你的tfrecord文件):

import tensorflow as tf
from PIL import Image
#tfrecord文件路径
filename_queue = tf.train.string_input_producer(["D:/验证码识别/使用一个标签来训练/train.tfrecords"])
reader = tf.TFRecordReader()
# 返回文件名和文件
_, serialized_example = reader.read(filename_queue)
features = tf.parse_single_example(serialized_example,
                                   features={
                                       'image': tf.FixedLenFeature([], tf.string),
                                       'label': tf.FixedLenFeature([], tf.int64)
                                   })
image = tf.decode_raw(features['image'],tf.uint8)
image = tf.reshape(image,[224,224])
label = tf.cast(features['label'], tf.int64)
with tf.Session() as sess:
    init_op = tf.global_variables_initializer()
    sess.run(init_op)
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(coord=coord)
    for i in range(10):
        example, label0 = sess.run([image, label])  # 在会话中取出image和label
        print(example, label0)
        print(example.shape)
        img = Image.fromarray(example, 'L')  # 这里Image是之前提到的
        img.save('./' + str(i) + '_''Label_' + str(label0) + '.jpg')  # 存下图片
    coord.request_stop()
    coord.join(threads)

 

我们随机的生成了10张图片来看看是否准确,运行上面的程序后会在你的当前目录下生成一些图片,好了看看是否正确把!

接下来就是训练部分了,请看我的下一篇文章验证码识别3——训练部分

 

 

 

  • 1
    点赞
  • 0
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值