使用Tensorflow構建和訓練自己的CNN來做簡單的驗證碼識別
阿新 • • 發佈:2019-02-13
Tensorflow是目前最流行的深度學習框架,我們可以用它來搭建自己的卷積神經網路並訓練自己的分類器,本文介紹怎樣使用Tensorflow構建自己的CNN,怎樣訓練用於簡單的驗證碼識別的分類器。本文假設你已經安裝好了Tensorflow,瞭解過CNN的一些知識。
下面將分步介紹怎樣獲得訓練資料,怎樣使用tensorflow構建卷積神經網路,怎樣訓練,以及怎樣測試訓練出來的分類器
1. 準備訓練樣本
使用Python的庫captcha來生成我們需要的訓練樣本,程式碼如下:
import sys
import os import shutil import random import time #captcha是用於生成驗證碼圖片的庫,可以 pip install captcha 來安裝它 from captcha.image import ImageCaptcha #用於生成驗證碼的字符集 CHAR_SET = ['0','1','2','3','4','5','6','7','8','9'] #字符集的長度 CHAR_SET_LEN = 10 #驗證碼的長度,每個驗證碼由4個數字組成 CAPTCHA_LEN = 4 #驗證碼圖片的存放路徑 CAPTCHA_IMAGE_PATH = 'E:/Tensorflow/captcha/images/' #用於模型測試的驗證碼圖片的存放路徑,它裡面的驗證碼圖片作為測試集 TEST_IMAGE_PATH = 'E:/Tensorflow/captcha/test/' #用於模型測試的驗證碼圖片的個數,從生成的驗證碼圖片中取出來放入測試集中 TEST_IMAGE_NUMBER = 50 #生成驗證碼圖片,4位的十進位制數字可以有10000種驗證碼 def generate_captcha_image(charSet = CHAR_SET, charSetLen=CHAR_SET_LEN, captchaImgPath=CAPTCHA_IMAGE_PATH): k = 0 total = 1 for i in range(CAPTCHA_LEN): total *= charSetLen for i in range(charSetLen): for j in range(charSetLen): for m in range(charSetLen): for n in range(charSetLen): captcha_text = charSet[i] + charSet[j] + charSet[m] + charSet[n] image = ImageCaptcha() image.write(captcha_text, captchaImgPath + captcha_text + '.jpg') k += 1 sys.stdout.write("\rCreating %d/%d" % (k, total)) sys.stdout.flush() #從驗證碼的圖片集中取出一部分作為測試集,這些圖片不參加訓練,只用於模型的測試 def prepare_test_set(): fileNameList = [] for filePath in os.listdir(CAPTCHA_IMAGE_PATH): captcha_name = filePath.split('/')[-1] fileNameList.append(captcha_name) random.seed(time.time()) random.shuffle(fileNameList) for i in range(TEST_IMAGE_NUMBER): name = fileNameList[i] shutil.move(CAPTCHA_IMAGE_PATH + name, TEST_IMAGE_PATH + name) if __name__ == '__main__': generate_captcha_image(CHAR_SET, CHAR_SET_LEN, CAPTCHA_IMAGE_PATH) prepare_test_set() sys.stdout.write("\nFinished") sys.stdout.flush()
執行上面的程式碼,可以生成驗證碼圖片,
生成的驗證碼圖片如下圖所示:
2. 構建CNN,訓練分類器
程式碼如下:
import tensorflow as tf import numpy as np from PIL import Image import os import random import time #驗證碼圖片的存放路徑 CAPTCHA_IMAGE_PATH = 'E:/Tensorflow/captcha/images/' #驗證碼圖片的寬度 CAPTCHA_IMAGE_WIDHT = 160 #驗證碼圖片的高度 CAPTCHA_IMAGE_HEIGHT = 60 CHAR_SET_LEN = 10 CAPTCHA_LEN = 4 #60%的驗證碼圖片放入訓練集中 TRAIN_IMAGE_PERCENT = 0.6 #訓練集,用於訓練的驗證碼圖片的檔名 TRAINING_IMAGE_NAME = [] #驗證集,用於模型驗證的驗證碼圖片的檔名
VALIDATION_IMAGE_NAME = []
#存放訓練好的模型的路徑 MODEL_SAVE_PATH = 'E:/Tensorflow/captcha/models/' def get_image_file_name(imgPath=CAPTCHA_IMAGE_PATH): fileName = [] total = 0 for filePath in os.listdir(imgPath): captcha_name = filePath.split('/')[-1] fileName.append(captcha_name) total += 1 return fileName, total #將驗證碼轉換為訓練時用的標籤向量,維數是 40 #例如,如果驗證碼是 ‘0296’ ,則對應的標籤是 # [1 0 0 0 0 0 0 0 0 0 # 0 0 1 0 0 0 0 0 0 0 # 0 0 0 0 0 0 0 0 0 1 # 0 0 0 0 0 0 1 0 0 0] def name2label(name): label = np.zeros(CAPTCHA_LEN * CHAR_SET_LEN) for i, c in enumerate(name): idx = i*CHAR_SET_LEN + ord(c) - ord('0') label[idx] = 1 return label #取得驗證碼圖片的資料以及它的標籤 def get_data_and_label(fileName, filePath=CAPTCHA_IMAGE_PATH): pathName = os.path.join(filePath, fileName) img = Image.open(pathName) #轉為灰度圖 img = img.convert("L") image_array = np.array(img) image_data = image_array.flatten()/255 image_label = name2label(fileName[0:CAPTCHA_LEN]) return image_data, image_label #生成一個訓練batch def get_next_batch(batchSize=32, trainOrTest='train', step=0): batch_data = np.zeros([batchSize, CAPTCHA_IMAGE_WIDHT*CAPTCHA_IMAGE_HEIGHT]) batch_label = np.zeros([batchSize, CAPTCHA_LEN * CHAR_SET_LEN]) fileNameList = TRAINING_IMAGE_NAME if trainOrTest == 'validate': fileNameList = VALIDATION_IMAGE_NAME totalNumber = len(fileNameList) indexStart = step*batchSize for i in range(batchSize): index = (i + indexStart) % totalNumber name = fileNameList[index] img_data, img_label = get_data_and_label(name) batch_data[i, : ] = img_data batch_label[i, : ] = img_label return batch_data, batch_label #構建卷積神經網路並訓練 def train_data_with_CNN(): #初始化權值 def weight_variable(shape, name='weight'): init = tf.truncated_normal(shape, stddev=0.1) var = tf.Variable(initial_value=init, name=name) return var #初始化偏置 def bias_variable(shape, name='bias'): init = tf.constant(0.1, shape=shape) var = tf.Variable(init, name=name) return var #卷積 def conv2d(x, W, name='conv2d'): return tf.nn.conv2d(x, W, strides=[1,1,1,1], padding='SAME', name=name) #池化 def max_pool_2X2(x, name='maxpool'): return tf.nn.max_pool(x, ksize=[1,2,2,1], strides=[1,2,2,1], padding='SAME', name=name) #輸入層 #請注意 X 的 name,在測試model時會用到它 X = tf.placeholder(tf.float32, [None, CAPTCHA_IMAGE_WIDHT * CAPTCHA_IMAGE_HEIGHT], name='data-input') Y = tf.placeholder(tf.float32, [None, CAPTCHA_LEN * CHAR_SET_LEN], name='label-input') x_input = tf.reshape(X, [-1, CAPTCHA_IMAGE_HEIGHT, CAPTCHA_IMAGE_WIDHT, 1], name='x-input') #dropout,防止過擬合 #請注意 keep_prob 的 name,在測試model時會用到它 keep_prob = tf.placeholder(tf.float32, name='keep-prob') #第一層卷積 W_conv1 = weight_variable([5,5,1,32], 'W_conv1') B_conv1 = bias_variable([32], 'B_conv1') conv1 = tf.nn.relu(conv2d(x_input, W_conv1, 'conv1') + B_conv1) conv1 = max_pool_2X2(conv1, 'conv1-pool') conv1 = tf.nn.dropout(conv1, keep_prob) #第二層卷積 W_conv2 = weight_variable([5,5,32,64], 'W_conv2') B_conv2 = bias_variable([64], 'B_conv2') conv2 = tf.nn.relu(conv2d(conv1, W_conv2,'conv2') + B_conv2) conv2 = max_pool_2X2(conv2, 'conv2-pool') conv2 = tf.nn.dropout(conv2, keep_prob) #第三層卷積 W_conv3 = weight_variable([5,5,64,64], 'W_conv3') B_conv3 = bias_variable([64], 'B_conv3') conv3 = tf.nn.relu(conv2d(conv2, W_conv3, 'conv3') + B_conv3) conv3 = max_pool_2X2(conv3, 'conv3-pool') conv3 = tf.nn.dropout(conv3, keep_prob) #全連結層 #每次池化後,圖片的寬度和高度均縮小為原來的一半,進過上面的三次池化,寬度和高度均縮小8倍 W_fc1 = weight_variable([20*8*64, 1024], 'W_fc1') B_fc1 = bias_variable([1024], 'B_fc1') fc1 = tf.reshape(conv3, [-1, 20*8*64]) fc1 = tf.nn.relu(tf.add(tf.matmul(fc1, W_fc1), B_fc1)) fc1 = tf.nn.dropout(fc1, keep_prob) #輸出層 W_fc2 = weight_variable([1024, CAPTCHA_LEN * CHAR_SET_LEN], 'W_fc2') B_fc2 = bias_variable([CAPTCHA_LEN * CHAR_SET_LEN], 'B_fc2') output = tf.add(tf.matmul(fc1, W_fc2), B_fc2, 'output') loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=Y, logits=output)) optimizer = tf.train.AdamOptimizer(0.001).minimize(loss) predict = tf.reshape(output, [-1, CAPTCHA_LEN, CHAR_SET_LEN], name='predict') labels = tf.reshape(Y, [-1, CAPTCHA_LEN, CHAR_SET_LEN], name='labels') #預測結果 #請注意 predict_max_idx 的 name,在測試model時會用到它 predict_max_idx = tf.argmax(predict, axis=2, name='predict_max_idx') labels_max_idx = tf.argmax(labels, axis=2, name='labels_max_idx') predict_correct_vec = tf.equal(predict_max_idx, labels_max_idx) accuracy = tf.reduce_mean(tf.cast(predict_correct_vec, tf.float32)) saver = tf.train.Saver() with tf.Session() as sess: sess.run(tf.global_variables_initializer()) steps = 0 for epoch in range(6000): train_data, train_label = get_next_batch(64, 'train', steps) sess.run(optimizer, feed_dict={X : train_data, Y : train_label, keep_prob:0.75}) if steps % 100 == 0: test_data, test_label = get_next_batch(100, 'validate', steps) acc = sess.run(accuracy, feed_dict={X : test_data, Y : test_label, keep_prob:1.0}) print("steps=%d, accuracy=%f" % (steps, acc)) if acc > 0.99: saver.save(sess, MODEL_SAVE_PATH+"crack_captcha.model", global_step=steps) break steps += 1 if __name__ == '__main__': image_filename_list, total = get_image_file_name(CAPTCHA_IMAGE_PATH) random.seed(time.time()) #打亂順序 random.shuffle(image_filename_list) trainImageNumber = int(total * TRAIN_IMAGE_PERCENT) #分成測試集 TRAINING_IMAGE_NAME = image_filename_list[ : trainImageNumber] #和驗證集 VALIDATION_IMAGE_NAME = image_filename_list[trainImageNumber : ] train_data_with_CNN() print('Training finished')
執行上面的程式碼,開始訓練,訓練要花些時間,如果沒有GPU的話,會慢些,
訓練完後,輸出如下結果,經過4100次的迭代,訓練出來的分類器模型在驗證集上識別的準確率為99.5%
生成的模型檔案如下,在模型測試時將用到這些檔案
3. 測試模型
編寫程式碼,對訓練出來的模型進行測試
import tensorflow as tf
import numpy as np
from PIL import Image
import os
import matplotlib.pyplot as plt
CAPTCHA_LEN = 4
MODEL_SAVE_PATH = 'E:/Tensorflow/captcha/models/'
TEST_IMAGE_PATH = 'E:/Tensorflow/captcha/test/'
def get_image_data_and_name(fileName, filePath=TEST_IMAGE_PATH):
pathName = os.path.join(filePath, fileName)
img = Image.open(pathName)
#轉為灰度圖
img = img.convert("L")
image_array = np.array(img)
image_data = image_array.flatten()/255
image_name = fileName[0:CAPTCHA_LEN]
return image_data, image_name
def digitalStr2Array(digitalStr):
digitalList = []
for c in digitalStr:
digitalList.append(ord(c) - ord('0'))
return np.array(digitalList)
def model_test():
nameList = []
for pathName in os.listdir(TEST_IMAGE_PATH):
nameList.append(pathName.split('/')[-1])
totalNumber = len(nameList)
#載入graph
saver = tf.train.import_meta_graph(MODEL_SAVE_PATH+"crack_captcha.model-4100.meta")
graph = tf.get_default_graph()
#從graph取得 tensor,他們的name是在構建graph時定義的(檢視上面第2步裡的程式碼)
input_holder = graph.get_tensor_by_name("data-input:0")
keep_prob_holder = graph.get_tensor_by_name("keep-prob:0")
predict_max_idx = graph.get_tensor_by_name("predict_max_idx:0")
with tf.Session() as sess:
saver.restore(sess, tf.train.latest_checkpoint(MODEL_SAVE_PATH))
count = 0
for fileName in nameList:
img_data, img_name = get_image_data_and_name(fileName, TEST_IMAGE_PATH)
predict = sess.run(predict_max_idx, feed_dict={input_holder:[img_data], keep_prob_holder : 1.0})
filePathName = TEST_IMAGE_PATH + fileName
print(filePathName)
img = Image.open(filePathName)
plt.imshow(img)
plt.axis('off')
plt.show()
predictValue = np.squeeze(predict)
rightValue = digitalStr2Array(img_name)
if np.array_equal(predictValue, rightValue):
result = '正確'
count += 1
else:
result = '錯誤'
print('實際值:{}, 預測值:{},測試結果:{}'.format(rightValue, predictValue, result))
print('\n')
print('正確率:%.2f%%(%d/%d)' % (count*100/totalNumber, count, totalNumber))
if __name__ == '__main__':
model_test()
對模型的測試結果如下,在測試集上識別的準確率為 94%
下面是兩個識別錯誤的驗證碼
訓練出的模型放在了下面的雲盤裡,有興趣的同學可以用它做下驗證碼的識別