学习tensorflow有一段时间了，想做点东西来练一下手。为了更有意思点，下面将搭建一个简单的验证码识别系统。
准备验证码数据
下面将生成一万张四位英文字母的验证码，验证码的大小是100 * 30的图片，只包含大写的英文字母，并将目标值保存到csv文件。
import random
import pandas as pd
from PIL import Image, ImageDraw, ImageFont
def generate_captcha(filename, format):
"""
生成四位验证码
:param filename: 要保存的文件名
:param format: 保存图片格式
:return: 验证码的值
"""
# 定义使用Image类实例化一个长为100px,宽为30px,基于RGB的(255,255,255)颜色的图片
img = Image.new(mode="RGB", size=(100, 30), color=(255, 255, 255))
# 实例化一支画笔
draw = ImageDraw.Draw(img, mode="RGB")
# 定义要使用的字体
font = ImageFont.truetype("arial", 28)
result = ""
for i in range(4):
# 每循环一次,从a到z中随机生成一个字母
# 65到90为字母的ASCII码,使用chr把生成的ASCII码转换成字符
# str把生成的数字转换成字符串
char = random.choice([chr(random.randint(65, 90))])
result += char
# 每循环一次重新生成随机颜色
color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
# 把生成的字母或数字添加到图片上
# 图片长度为100px,要生成4个数字或字母则每添加一个,其位置就要向后移动24px
draw.text([i * 24 + 3, 0], char, color, font=font)
# 保存生成的文件
with open(filename, "wb") as f:
img.save(f, format=format)
return result
if __name__ == "__main__":
data = []
# 生成10000张验证码图片，并将目标值存入csv文件
for j in range(10000):
val = generate_captcha("./pics/{}.png".format(j), "png")
data.append([val])
# 将验证码的值保存到csv文件
df = pd.DataFrame(data, columns=['label'])
df.to_csv('./pics/data.csv', header=False)
生成的验证码图片是这样子的，如下：
csv文件内容：
0,EFGQ
1,ZDKO
2,UWLD
3,CPDH
....
保存为tfrecords文件
上面生成的图片和其目标值是分开的，在进行训练时不太方便（训练时每次都要单独的读取图片和特征值）。保存为tfrecords文件，在训练时会方便很多，读取出来的每条记录既有图片特征值又有目标值。
import tensorflow as tf
import os
import numpy as np
class CaptchaInput(object):
def __init__(self, captcha_dir, letter, tfrecords_dir):
"""
:param captcha_dir: 验证码路径
:param letter: 验证码字符种类
:param tfrecords_dir: tfrecords文件保存的目录
"""
self.captcha_dir = captcha_dir
self.letter = letter
self.tfrecords_dir = tfrecords_dir
# 列出图片文件，并进行排序
self.file_list = os.listdir(self.captcha_dir)
self.file_list = [i for i in self.file_list if i.endswith(".png")]
self.file_list.sort(key=lambda x: int(x[0:-4]))
self.file_list = [os.path.join(self.captcha_dir, i) for i in self.file_list]
# 标签文件路径
self.labels_path = os.path.join(self.captcha_dir, "data.csv")
def read_captcha_image(self):
"""读取验证码图片数据"""
# 构造文件队列
file_queue = tf.train.string_input_producer(self.file_list, shuffle=False)
# 构建阅读器
reader = tf.WholeFileReader()
# 读取图片内容
key, value = reader.read(file_queue)
# 解码图片
image = tf.image.decode_png(value)
image.set_shape([30, 100, 3])
# 批量读取
image_batch = tf.train.batch([image], batch_size=len(self.file_list),
num_threads=1, capacity=len(self.file_list))
return image_batch
def read_captcha_label(self):
"""读取 验证码标签数据"""
# 构造文件队列
file_queue = tf.train.string_input_producer([self.labels_path], shuffle=False)
# 构建文件阅读器
reader = tf.TextLineReader()
# 读取标签内容
key, value = reader.read(file_queue)
records = [[0], [""]]
index, label = tf.decode_csv(value, record_defaults=records)
# 批量读取
label_batch = tf.train.batch([label], batch_size=len(self.file_list),
num_threads=1, capacity=len(self.file_list))
return label_batch
def process_labels(self, labels):
"""将标签字符转换成数字张量"""
# 构建字符索引
num_letter_dict = dict(enumerate(list(self.letter)))
letter_num_dict = dict(zip(num_letter_dict.values(), num_letter_dict.keys()))
ret = []
for label in labels:
arr = [letter_num_dict[i] for i in label.decode("utf-8")]
ret.append(arr)
return np.array(ret)
def write_to_tfrecords(self, images, labels):
"""
将图片和标签写入到tfrecords文件中
:param images: 特征值
:param labels: 目标值
:return:
"""
# labels = tf.cast(labels, tf.uint8)
# images = tf.cast(images, tf.uint8)
# 建立存储文件
fw = tf.python_io.TFRecordWriter(self.tfrecords_dir)
for i in range(len(self.file_list)):
# images[i]为numpy.ndarray
image_bytes = images[i].tobytes()
# labels[i]为numpy.ndarray
label_bytes = labels[i].tobytes()
example = tf.train.Example(features=tf.train.Features(feature={
"image": tf.train.Feature(bytes_list=tf.train.BytesList(value=[image_bytes])),
"label": tf.train.Feature(bytes_list=tf.train.BytesList(value=[label_bytes]))
}))
print("保存第%d张图片" % (i, ))
fw.write(example.SerializeToString())
# 关闭
fw.close()
def execute(self):
image_batch = self.read_captcha_image()
label_batch = self.read_captcha_label()
with tf.Session() as sess:
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess, coord=coord)
# [b'EFGQ' b'ZDKO' b'UWLD' ... b'TKPD' b'ZZEU' b'ATYA']
labels = sess.run(label_batch)
# labels为numpy.ndarray
labels = self.process_labels(labels)
# images为numpy.ndarray
images = sess.run(image_batch)
self.write_to_tfrecords(images, labels)
coord.request_stop()
coord.join(threads)
FLAGS = tf.app.flags.FLAGS
tf.app.flags.DEFINE_string("captcha_dir", "./pics", "验证码图片路径")
tf.app.flags.DEFINE_string("letter", "ABCDEFGHIJKLMNOPQRSTUVWXYZ", "验证码字符种类")
tf.app.flags.DEFINE_string("tfrecords_dir", "./tfrecords/captcha.tfrecords", "验证码tfrecords文件")
if __name__ == "__main__":
c = CaptchaInput(FLAGS.captcha_dir, FLAGS.letter, FLAGS.tfrecords_dir)
c.execute()
需要注意：
-
os.listdir返回的文件名称的顺序是按照ascii表的顺序（1.png, 10.png...）需要对其进行排序 - 使用tensorflow读取图片和标签文件时，需要加上
shuffle=False，避免文件乱序了，图片和目标值对应不上。
验证码训练
import tensorflow as tf
FLAGS = tf.app.flags.FLAGS
tf.app.flags.DEFINE_string("captcha_dir", "./tfrecords/captcha.tfrecords", "验证码数据文件")
tf.app.flags.DEFINE_integer("batch_size", 100, "每批次训练样本数")
def read_and_decode():
"""读取验证码数据
:return image_batch, label_batch
"""
# 文件队列
file_queue = tf.train.string_input_producer([FLAGS.captcha_dir])
# 文件读取器
reader = tf.TFRecordReader()
# 读取内容
key, value = reader.read(file_queue)
# 解析tfrecords
features = tf.parse_single_example(value, features={
"image": tf.FixedLenFeature([], tf.string),
"label": tf.FixedLenFeature([], tf.string)
})
# 解码
image = tf.decode_raw(features["image"], tf.uint8)
label = tf.decode_raw(features["label"], tf.uint8)
# print(image, label)
# 改变形状
image_reshape = tf.reshape(image, [30, 100, 3])
label_reshape = tf.reshape(label, [4])
# print(image_reshape, label_reshape)
# 批处理
image_batch, label_batch = tf.train.batch([image_reshape, label_reshape],
batch_size=FLAGS.batch_size, num_threads=1, capacity=FLAGS.batch_size)
return image_batch, label_batch
def weight_variables(shape):
"""权重初始化函数"""
w = tf.Variable(tf.random_normal(shape=shape, mean=0.0, stddev=1.0))
return w
def bias_variables(shape):
"""偏置初始化函数"""
b = tf.Variable(tf.constant(0.0, shape=shape))
return b
def fc_model(image):
"""全连接模型"""
with tf.variable_scope("fc_model"):
image_reshape = tf.reshape(image, [-1, 30 * 100 * 3])
# 随机初始化权重和偏重
weights = weight_variables([30 * 100 * 3, 4 * 26])
bias = bias_variables([4 * 26])
# 全连接计算
y_predict = tf.matmul(tf.cast(image_reshape, tf.float32), weights) + bias
return y_predict
def label_to_onehot(label):
"""目标值转换成one-hot编码"""
label_onehot = tf.one_hot(label, depth=26, on_value=1.0, axis=2)
return label_onehot
def captcharec():
"""验证码识别"""
image_batch, label_batch = read_and_decode()
# [100, 104]
y_predict = fc_model(image_batch)
y_true = label_to_onehot(label_batch)
# softmax计算，交叉熵损失计算
with tf.variable_scope("soft_cross"):
loss = tf.nn.softmax_cross_entropy_with_logits(
labels=tf.reshape(y_true, [-1, 4 * 26]),
logits=y_predict
)
# 梯度下降损失优化
with tf.variable_scope("optimizer"):
train_op = tf.train.GradientDescentOptimizer(0.01).minimize(loss)
# 准确率
with tf.variable_scope("acc"):
equal_list = tf.equal(tf.argmax(y_true, 2), tf.argmax(tf.reshape(y_predict, [-1, 4, 26]), 2))
accuracy = tf.reduce_mean(tf.cast(equal_list, tf.float32))
init_op = tf.global_variables_initializer()
with tf.Session() as sess:
sess.run(init_op)
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess, coord=coord)
for i in range(3000):
sess.run(train_op)
print("第%d次训练的准确率为：%f" % (i, accuracy.eval()))
coord.request_stop()
coord.join(threads)
if __name__ == '__main__':
captcharec()