Keras TensorFlow 验证码识别（附数据集）
效果
先看效果图，左边是 两张测试数据，右边是 预测结果
标注数据集下载地址已更新
实现
1. 训练模型
bathcsize 为 700 轮次 50
#!/usr/bin/env python
# coding: utf-8
# # 训练模型
#
# ## 引入第三方包
# In[1]:
from PIL import Image
from keras import backend as K
from keras.utils.vis_utils import plot_model
from keras.models import *
from keras.layers import *
import glob
import pickle
import numpy as np
import tensorflow.gfile as gfile
import matplotlib.pyplot as plt
# ## 定义超参数和字符集
# In[2]:
NUMBER = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
CAPTCHA_CHARSET = NUMBER # 验证码字符集
CAPTCHA_LEN = 5 # 验证码长度
CAPTCHA_HEIGHT = 50 # 验证码高度
CAPTCHA_WIDTH = 200 # 验证码宽度
TRAIN_DATA_DIR = './train-data/' # 验证码数据集目录
TEST_DATA_DIR = './test-data/'
BATCH_SIZE = 700
EPOCHS = 50
OPT = 'adam' # adam
LOSS = 'binary_crossentropy'
MODEL_DIR = './model/train_demo/'
MODEL_FORMAT = '.h5'
HISTORY_DIR = './history/train_demo/'
HISTORY_FORMAT = '.history'
filename_str = "{}captcha_{}_{}_bs_{}_epochs_{}{}"
# 模型网络结构文件
MODEL_VIS_FILE = 'captcha_classfication' + '.png'
# 模型文件
MODEL_FILE = filename_str.format(MODEL_DIR, OPT, LOSS, str(BATCH_SIZE), str(EPOCHS), MODEL_FORMAT)
# 训练记录文件
HISTORY_FILE = filename_str.format(HISTORY_DIR, OPT, LOSS, str(BATCH_SIZE), str(EPOCHS), HISTORY_FORMAT)
# ## 将 RGB 验证码图像转为灰度图
# In[3]:
def rgb2gray(img):
# Y' = 0.299 R + 0.587 G + 0.114 B
# https://en.wikipedia.org/wiki/Grayscale#Converting_color_to_grayscale
return np.dot(img[..., :3], [0.299, 0.587, 0.114])
# ## 对验证码中每个字符进行 one-hot 编码
# In[4]:
def text2vec(text, length=CAPTCHA_LEN, charset=CAPTCHA_CHARSET):
text_len = len(text)
# 验证码长度校验
if text_len != length:
raise ValueError('Error: length of captcha should be {}, but got {}'.format(length, text_len))
# 生成一个形如（CAPTCHA_LEN*CAPTHA_CHARSET,) 的一维向量
# 例如，4个纯数字的验证码生成形如(4*10,)的一维向量
vec = np.zeros(length * len(charset))
for i in range(length):
# One-hot 编码验证码中的每个数字
# 每个字符的热码 = 索引 + 偏移量
vec[charset.index(text[i]) + i * len(charset)] = 1
return vec
# ## 将验证码向量解码为对应字符
# In[5]:
def vec2text(vector):
if not isinstance(vector, np.ndarray):
vector = np.asarray(vector)
vector = np.reshape(vector, [CAPTCHA_LEN, -1])
text = ''
for item in vector:
text += CAPTCHA_CHARSET[np.argmax(item)]
return text
# ## 适配 Keras 图像数据格式
# In[6]:
def fit_keras_channels(batch, rows=CAPTCHA_HEIGHT, cols=CAPTCHA_WIDTH):
if K.image_data_format() == 'channels_first':
batch = batch.reshape(batch.shape[0], 1, rows, cols)
input_shape = (1, rows, cols)
else:
batch = batch.reshape(batch.shape[0], rows, cols, 1)
input_shape = (rows, cols, 1)
return batch, input_shape
# ## 读取训练集
# In[7]:
X_train = []
Y_train = []
for filename in glob.glob(TRAIN_DATA_DIR + '*.jpg'):
X_train.append(np.array(Image.open(filename)))
Y_train.append(filename.lstrip(TRAIN_DATA_DIR).rstrip('.jpg'))
# ## 处理训练集图像
# In[8]:
# list -> rgb(numpy)
X_train = np.array(X_train, dtype=np.float32)
# rgb -> gray
X_train = rgb2gray(X_train)
# normalize
X_train = X_train / 255
# Fit keras channels
X_train, input_shape = fit_keras_channels(X_train)
print(X_train.shape, type(X_train))
print(input_shape)
# ## 处理训练集标签
# In[9]:
Y_train = list(Y_train)
for i in range(len(Y_train)):
Y_train[i] = text2vec(Y_train[i])
Y_train = np.asarray(Y_train)
print(Y_train.shape, type(Y_train))
# ## 读取测试集，处理对应图像和标签
# In[10]:
X_test = []
Y_test = []
for filename in glob.glob(TEST_DATA_DIR + '*.jpg'):
X_test.append(np.array(Image.open(filename)))
Y_test.append(filename.lstrip(TEST_DATA_DIR).rstrip('.jpg'))
# list -> rgb -> gray -> normalization -> fit keras
X_test = np.array(X_test, dtype=np.float32)
X_test = rgb2gray(X_test)
X_test = X_test / 255
X_test, _ = fit_keras_channels(X_test)
Y_test = list(Y_test)
for i in range(len(Y_test)):
Y_test[i] = text2vec(Y_test[i])
Y_test = np.asarray(Y_test)
print(X_test.shape, type(X_test))
print(Y_test.shape, type(Y_test))
# ## 创建验证码识别模型
# In[11]:
# 输入层
inputs = Input(shape=input_shape, name="inputs")
# 第1层卷积
conv1 = Conv2D(32, (3, 3), name="conv1")(inputs)
relu1 = Activation('relu', name="relu1")(conv1)
# 第2层卷积
conv2 = Conv2D(32, (3, 3), name="conv2")(relu1)
relu2 = Activation('relu', name="relu2")(conv2)
pool2 = MaxPooling2D(pool_size=(2, 2), padding='same', name="pool2")(relu2)
# 第3层卷积
conv3 = Conv2D(64, (3, 3), name="conv3")(pool2)
relu3 = Activation('relu', name="relu3")(conv3)
pool3 = MaxPooling2D(pool_size=(2, 2), padding='same', name="pool3")(relu3)
# 将 Pooled feature map 摊平后输入全连接网络
x = Flatten()(pool3)
# Dropout
x = Dropout(0.25)(x)
# 4个全连接层分别做10分类，分别对应4个字符。
x = [Dense(10, activation='softmax', name='fc%d' % (i + 1))(x) for i in range(5)]
# 4个字符向量拼接在一起，与标签向量形式一致，作为模型输出。
outs = Concatenate()(x)
# 定义模型的输入与输出
model = Model(inputs=inputs, outputs=outs)
model.compile(optimizer=OPT, loss=LOSS, metrics=['accuracy'])
# ## 查看模型摘要
# In[12]:
model.summary()
# ## 模型可视化
# In[13]:
plot_model(model, to_file=MODEL_VIS_FILE, show_shapes=True)
# ## 训练模型
# In[14]:
history = model.fit(X_train,
Y_train,
batch_size=BATCH_SIZE,
epochs=EPOCHS,
verbose=2,
validation_data=(X_test, Y_test))
# ## 预测样例
# In[47]:
print(vec2text(Y_test[9]))
# In[48]:
yy = model.predict(X_test[9].reshape(1, 50, 200, 1))
# In[49]:
print(vec2text(yy))
# ## 保存模型
# In[50]:
if not gfile.Exists(MODEL_DIR):
gfile.MakeDirs(MODEL_DIR)
model.save(MODEL_FILE)
print('Saved trained model at %s ' % MODEL_FILE)
# ## 保存训练过程记录
# In[51]:
print(history.history['acc'])
# In[52]:
history.history.keys()
# In[53]:
if gfile.Exists(HISTORY_DIR) == False:
gfile.MakeDirs(HISTORY_DIR)
# with open(HISTORY_FILE, 'wb') as f:
# pickle.dump(history.history, f)
# In[54]:
# print(HISTORY_FILE)
# In[ ]:
2 模型结构设计
3. 建立预测服务
import base64
import numpy as np
import tensorflow as tf
from io import BytesIO
from flask import Flask, request, jsonify
from keras.models import load_model
from PIL import Image
NUMBER = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
LOWERCASE = ['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']
UPPERCASE = ['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']
CAPTCHA_CHARSET = NUMBER # 验证码字符集
CAPTCHA_LEN = 5 # 验证码长度
CAPTCHA_HEIGHT = 50 # 验证码高度
CAPTCHA_WIDTH = 200 # 验证码宽度
# 10 个 Epochs 训练的模型 rmsprop adam
MODEL_FILE = './model/train_demo/captcha_adam_binary_crossentropy_bs_700_epochs_50.h5'
def vec2text(vector):
if not isinstance(vector, np.ndarray):
vector = np.asarray(vector)
vector = np.reshape(vector, [CAPTCHA_LEN, -1])
text = ''
for item in vector:
text += CAPTCHA_CHARSET[np.argmax(item)]
return text
def rgb2gray(img):
# Y' = 0.299 R + 0.587 G + 0.114 B
# https://en.wikipedia.org/wiki/Grayscale#Converting_color_to_grayscale
return np.dot(img[...,:3], [0.299, 0.587, 0.114])
app = Flask(__name__) # 创建 Flask 实例
# 测试 URL
@app.route('/ping', methods=['GET', 'POST'])
def hello_world():
return 'pong'
# 验证码识别 URL
@app.route('/predict', methods=['POST'])
def predict():
response = {'success': False, 'prediction': '', 'debug': 'error'}
received_image= False
if request.method == 'POST':
if request.files.get('image'): # 图像文件
image = request.files['image'].read()
# print("image is ")
# print(image)
received_image = True
response['debug'] = 'get image'
elif request.get_json(): # base64 编码的图像文件
encoded_image = request.get_json()['image']
image = base64.b64decode(encoded_image)
received_image = True
response['debug'] = 'get json'
if received_image:
image = np.array(Image.open(BytesIO(image)))
image = rgb2gray(image).reshape(1, 50, 200, 1).astype('float32') / 255
with graph.as_default():
pred = model.predict(image)
response['prediction'] = response['prediction'] + vec2text(pred)
response['success'] = True
response['debug'] = 'predicted'
else:
response['debug'] = 'No Post'
return jsonify(response)
model = load_model(MODEL_FILE) # 加载模型
graph = tf.get_default_graph() # 获取 TensorFlow 默认数据流图
# 启动命令
# export FLASK_ENV=development && flask run --host=0.0.0.0
# curl 127.0.0.1:5000/ping
# curl -X POST -F image=@/root/Workspace/leon/test-data/56497.jpg 'http://localhost:5000/predict'
4 数据集下载
train-data.zip