TensorFlow — 基于CNN数字识别

环境:Ubuntu 14.04.1 LTS (GNU/Linux 3.13.0-105-generic x86_64)

1 安装captcha库

sudo pip install captcha

2 生成验证码训练数据

2.1 验证码生成器

采用 python 中生成器方式来生成我们的训练数据,这样的好处是,不需要提前生成大量的数据,训练过程中生成数据,并且可以无限生成数据。

2.1.1 示例代码

现在在 /home/ubuntu 目录下创建源文件 generate_captcha.py,内容可参考:

generate_captcha.py

#!/usr/bin/python

# -*- coding: utf-8 -*

from captcha.image import ImageCaptcha

from PIL import Image

import numpy as np

import random

import string

class generateCaptcha():

    def __init__(self,

        width = 160,#验证码图片的宽

        height = 60,#验证码图片的高

        char_num = 4,#验证码字符个数

characters = string.digits +string.ascii_uppercase + string.ascii_lowercase):

#验证码组成,数字+大写字母+小写字母

        self.width = width

        self.height = height

        self.char_num = char_num

        self.characters = characters

        self.classes = len(characters)

    def gen_captcha(self,batch_size = 50):

        X =np.zeros([batch_size,self.height,self.width,1])

        img =np.zeros((self.height,self.width),dtype=np.uint8)

        Y =np.zeros([batch_size,self.char_num,self.classes])

        image = ImageCaptcha(width =self.width,height = self.height)

        while True:

            for i in range(batch_size):

            captcha_str =''.join(random.sample(self.characters,self.char_num))

            img =image.generate_image(captcha_str).convert('L')

            img = np.array(img.getdata())

            X[i] =np.reshape(img,[self.height,self.width,1])/255.0

            for j,ch inenumerate(captcha_str):

                Y[i,j,self.characters.find(ch)] = 1

        Y =np.reshape(Y,(batch_size,self.char_num*self.classes))

yield X,Y

    def decode_captcha(self,y):

        y =np.reshape(y,(len(y),self.char_num,self.classes))

        return ''.join(self.characters[x] for xin np.argmax(y,axis = 2)[0,:])

    def get_parameter(self):

        return self.width, self.height, self.char_num, self.characters, self.classes

    def gen_test_captcha(self):

        image = ImageCaptcha(width =self.width,height = self.height)

        captcha_str =''.join(random.sample(self.characters,self.char_num))

        img = image.generate_image(captcha_str)

        img.save(captcha_str + '.jpg')

2.1.2 然后执行

cd /home/ubuntu;

python

import generate_captcha

g = generate_captcha.generateCaptcha()

g.gen_test_captcha()

2.1.3 执行结果

在 /home/ubuntu 目录下查看生成的验证码,jpg 格式的图片可以点击查看。

如:

Gxdl
KrO7

3 验证码识别模型

将验证码识别问题转化为分类问题,总共62^4 种类型,采用 4 个 one-hot 编码分别表示 4 个字符取值。

3.1 cnn 验证码识别模型

3 层隐藏层、2 层全连接层,对每层都进行dropout。input——>conv——>pool——>dropout——>conv——>pool——>dropout——>conv——>pool——>dropout——>fullyconnected layer——>dropout——>fully connected layer——>output

3.1.1 示例代码

现在在 /home/ubuntu 目录下创建源文件 captcha_model.py,内容可参考:

captcha_model.py

#!/usr/bin/python

# -*- coding: utf-8 -*

import tensorflow as tf

import math

class captchaModel():

    def __init__(self,

                width = 160,

                height = 60,

                char_num = 4,

                classes = 62):

        self.width = width

        self.height = height

        self.char_num = char_num

        self.classes = classes

    def conv2d(self,x, W):

        return tf.nn.conv2d(x, W, strides=[1,1, 1, 1], padding='SAME')

    def max_pool_2x2(self,x):

        return tf.nn.max_pool(x, ksize=[1, 2,2, 1],

strides=[1, 2, 2,1], padding='SAME')

    def weight_variable(self,shape):

        initial = tf.truncated_normal(shape,stddev=0.1)

        return tf.Variable(initial)

    def bias_variable(self,shape):

        initial = tf.constant(0.1, shape=shape)

        return tf.Variable(initial)

    def create_model(self,x_images,keep_prob):

        #first layer

        w_conv1 = self.weight_variable([5, 5,1, 32])

        b_conv1 = self.bias_variable([32])

        h_conv1 = tf.nn.relu(tf.nn.bias_add(self.conv2d(x_images, w_conv1), b_conv1))

        h_pool1 = self.max_pool_2x2(h_conv1)

        h_dropout1 = tf.nn.dropout(h_pool1,keep_prob)

        conv_width = math.ceil(self.width/2)

        conv_height = math.ceil(self.height/2)

        #second layer

        w_conv2 = self.weight_variable([5, 5,32, 64])

        b_conv2 = self.bias_variable([64])

        h_conv2 = tf.nn.relu(tf.nn.bias_add(self.conv2d(h_dropout1, w_conv2), b_conv2))

        h_pool2 = self.max_pool_2x2(h_conv2)

        h_dropout2 = tf.nn.dropout(h_pool2,keep_prob)

        conv_width = math.ceil(conv_width/2)

        conv_height = math.ceil(conv_height/2)

        #third layer

        w_conv3 = self.weight_variable([5, 5,64, 64])

        b_conv3 = self.bias_variable([64])

        h_conv3 = tf.nn.relu(tf.nn.bias_add(self.conv2d(h_dropout2, w_conv3), b_conv3))

        h_pool3 = self.max_pool_2x2(h_conv3)

        h_dropout3 =tf.nn.dropout(h_pool3,keep_prob)

        conv_width = math.ceil(conv_width/2)

        conv_height = math.ceil(conv_height/2)

        #first fully layer

        conv_width = int(conv_width)

        conv_height = int(conv_height)

        w_fc1 = self.weight_variable([64*conv_width*conv_height,1024])

        b_fc1 = self.bias_variable([1024])

        h_dropout3_flat = tf.reshape(h_dropout3,[-1,64*conv_width*conv_height])

        h_fc1 = tf.nn.relu(tf.nn.bias_add(tf.matmul(h_dropout3_flat, w_fc1), b_fc1))

        h_fc1_drop = tf.nn.dropout(h_fc1,keep_prob)

        #second fully layer

        w_fc2 = self.weight_variable([1024,self.char_num*self.classes])

        b_fc2 = self.bias_variable([self.char_num*self.classes])

        y_conv = tf.add(tf.matmul(h_fc1_drop,w_fc2), b_fc2)

        return y_conv

3.2 训练cnn 验证码识别模型

每批次采用 64 个训练样本,每 100 次循环采用 100 个测试样本检查识别准确度,当准确度大于 99% 时,训练结束,采用 GPU 需要 5-6个小时左右,CPU 大概需要 20 个小时左右。

注:作为实验,你可以通过调整train_captcha.py文件中if acc > 0.99:代码行的准确度节省训练时间(比如将0.99 为 0.01);同时,我们已经通过长时间的训练得到了一个训练集,可以通过如下命令将训练集下载到本地。

wget http://tensorflow-1253902462.cosgz.myqcloud.com/captcha/capcha_model.zip

unzip capcha_model.zip

3.2.1 示例代码

现在在 /home/ubuntu 目录下创建源文件 train_captcha.py,内容可参考:

train_captcha.py

#!/usr/bin/python

import tensorflow as tf

import numpy as np

import string

import generate_captcha

import captcha_model

if __name__ == '__main__':

    captcha =generate_captcha.generateCaptcha()

    width,height,char_num,characters,classes = captcha.get_parameter()

    x = tf.placeholder(tf.float32, [None,height,width,1])

    y_ = tf.placeholder(tf.float32, [None,char_num*classes])

    keep_prob = tf.placeholder(tf.float32)

    model = captcha_model.captchaModel(width,height,char_num,classes)

    y_conv = model.create_model(x,keep_prob)

    cross_entropy = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(labels=y_,logits=y_conv))

    train_step = tf.train.AdamOptimizer(1e-4).minimize(cross_entropy)

    predict = tf.reshape(y_conv, [-1,char_num,classes])

    real = tf.reshape(y_,[-1,char_num,classes])

    correct_prediction = tf.equal(tf.argmax(predict,2), tf.argmax(real,2))

    correct_prediction = tf.cast(correct_prediction, tf.float32)

    accuracy = tf.reduce_mean(correct_prediction)

    saver = tf.train.Saver()

    with tf.Session() as sess:

        sess.run(tf.global_variables_initializer())

        step = 0

        while True:

            batch_x,batch_y = next(captcha.gen_captcha(64))

            _,loss = sess.run([train_step,cross_entropy],feed_dict ={x: batch_x, y_: batch_y,keep_prob: 0.75})

            print ('step:%d,loss:%f' %(step,loss))

            if step % 100 == 0:

                batch_x_test,batch_y_test = next(captcha.gen_captcha(100))

                acc = sess.run(accuracy,feed_dict={x: batch_x_test, y_: batch_y_test, keep_prob: 1.})

                print('###############################################step:%d,accuracy:%f' %(step,acc))

if acc > 0.99:

saver.save(sess,"capcha_model.ckpt")

break

step += 1

3.2.2 然后执行

cd /home/ubuntu;

python train_captcha.py

3.2.3执行结果

step:75173,loss:0.010555

step:75174,loss:0.009410

step:75175,loss:0.009978

step:75176,loss:0.008089

step:75177,loss:0.009949

step:75178,loss:0.010126

step:75179,loss:0.009584

step:75180,loss:0.012272

step:75181,loss:0.010157

step:75182,loss:0.009529

step:75183,loss:0.007636

step:75184,loss:0.009058

step:75185,loss:0.010061

step:75186,loss:0.009941

step:75187,loss:0.009339

step:75188,loss:0.009685

step:75189,loss:0.009879

step:75190,loss:0.007799

step:75191,loss:0.010866

step:75192,loss:0.009838

step:75193,loss:0.010931

step:75194,loss:0.012859

step:75195,loss:0.008747

step:75196,loss:0.009147

step:75197,loss:0.009351

step:75198,loss:0.009746

step:75199,loss:0.010014

step:75200,loss:0.009024

###############################################step:75200,accuracy:0.992500

3.3 测试

cnn 验证码识别模型

3.3.1示例代码

现在在 /home/ubuntu 目录下创建源文件 predict_captcha.py,内容可参考:

predict_captcha.py

#!/usr/bin/python

from PIL import Image, ImageFilter

import tensorflow as tf

import numpy as np

import string

import sys

import generate_captcha

import captcha_model

if __name__ == '__main__':

    captcha = generate_captcha.generateCaptcha()

    width, height, char_num, characters, classes = captcha.get_parameter()

    gray_image = Image.open(sys.argv[1]).convert('L')

    img = np.array(gray_image.getdata())

    test_x = np.reshape(img,[height,width,1])/255.0

    x = tf.placeholder(tf.float32, [None,height,width,1])

    keep_prob = tf.placeholder(tf.float32)

    model = captcha_model.captchaModel(width,height,char_num,classes)

    y_conv = model.create_model(x,keep_prob)

    predict = tf.argmax(tf.reshape(y_conv,[-1,char_num, classes]),2)

    init_op = tf.global_variables_initializer()

    saver = tf.train.Saver()

    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.95)

    withtf.Session(config=tf.ConfigProto(log_device_placement=False,gpu_options=gpu_options))as sess:

        sess.run(init_op)

        saver.restore(sess,"capcha_model.ckpt")

        pre_list = sess.run(predict,feed_dict={x: [test_x],keep_prob: 1})

        for i in pre_list:

            s = ''

            for j in i:

                s += characters[j]

                print s

3.3.2 然后执行

cd /home/ubuntu;

python predict_captcha.py Gxdl.jpg

3.3.3 执行结果:

Gxdl

注:因时间的限制,我们可能调整了准确度导致执行结果不符合预期,这属于正常情况。

在训练时间足够长的情况下,可以采用验证码生成器生成测试数据,cnn训练出来的验证码识别模型还是很强大的,大小写的 z 都可以区分,甚至有时候人都无法区分,该模型也可以正确的识别。

4 完成

以上。

你可能感兴趣的:(TensorFlow — 基于CNN数字识别)