TensorFlow入门教程（十四）：生成验证码图片并转化为tfrecoder文件

本实例将列举验证码识别实验，在实验前，可以自己生成验证码图片作为训练集，测试集。下面先讲述验证码生成部分。

本例子中，验证码为4位数，为数字0-9，后期可以自己更改个数，以及添加字母。

[code]from captcha.image import ImageCaptcha  # pip install captcha
import numpy as np
from PIL import Image
import random
import sys

number = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
# alphabet = ['a','b','c','d','e','f','g','h','i','j','k','l','m','n','o','p','q','r','s','t','u','v','w','x','y','z']
# ALPHABET = ['A','B','C','D','E','F','G','H','I','J','K','L','M','N','O','P','Q','R','S','T','U','V','W','X','Y','Z']

def random_captcha_text(char_set=number, captcha_size=4):  # 验证码备选数据列表，若想要加字母，可以用number + alphabet， 验证码个数
captcha_text = []
for i in range(captcha_size):  # 循环 captcha_size 次
# 从验证码列表中随机选择一个数据
c = random.choice(char_set)
# 加入验证码列表
captcha_text.append(c)
return captcha_text

# 生成字符对应的验证码
def gen_captcha_text_and_image():
image = ImageCaptcha()
# 获得随机生成的验证码
captcha_text = random_captcha_text()
# 把验证码列表转为字符串
captcha_text = ''.join(captcha_text)
# 生成验证码
captcha = image.generate(captcha_text)
image.write(captcha_text, 'images/' + captcha_text + '.jpg')  # 写到文件，路径+名称

# 数量会少于10000，因为会重名覆盖
num = 10000
if __name__ == '__main__':
for i in range(num):
gen_captcha_text_and_image()
sys.stdout.write('\r>> Creating image %d/%d' % (i + 1, num))  # 显示运行进程
sys.stdout.flush()
sys.stdout.write('\n')
sys.stdout.flush()

print("生成完毕")

运行结果：

每一个图片的名称为验证码图片上的数值

将图片生成tfcecord文件（注意先创建文件夹再生成）

[code]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 = "G:/Python_demo/Deep_learn/images/"

# tfrecord文件存放路径
TFRECORD_DIR = "G:/Python_demo/Deep_learn/captcha/"

# 判断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, label0, label1, label2, label3):
# Abstract base class for protocol messages.
return tf.train.Example(features=tf.train.Features(feature={
'image': bytes_feature(image_data),
'label0': int64_feature(label0),
'label1': int64_feature(label1),
'label2': int64_feature(label2),
'label3': int64_feature(label3),
}))

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

with tf.Session() as sess:
output_filename = os.path.join(TFRECORD_DIR, split_name + '.tfrecords')  # 定义tfrecord文件的路径+名字
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][0:4]
num_labels = []
for j in range(4):
num_labels.append(int(labels[j]))

# 生成protocol数据类型
example = image_to_tfexample(image_data, num_labels[0], num_labels[1], num_labels[2], num_labels[3])
tfrecord_writer.write(example.SerializeToString())

except IOError as e:
print('Could not read:', filename)
print('Error:', e)
print('Skip it\n')
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]  # 测试集，前500张图片

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