2.CNN圖片多標(biāo)簽分類(lèi)(基于TensorFlow實(shí)現(xiàn)驗(yàn)證碼識(shí)別OCR)

上一篇實(shí)現(xiàn)了圖片CNN單標(biāo)簽分類(lèi)(貓狗圖片分類(lèi)任務(wù))
(地址:http://www.itdecent.cn/p/47f0319028f2
預(yù)告:下一篇用LSTM+CTC實(shí)現(xiàn)不定長(zhǎng)文本的OCR,本質(zhì)上是一種不固定標(biāo)簽個(gè)數(shù)的多標(biāo)簽分類(lèi)問(wèn)題
(地址:http://www.itdecent.cn/p/7a2b227896a8

本文所用到的10w驗(yàn)證碼數(shù)據(jù)集百度網(wǎng)盤(pán)下載地址(也可使用下文代碼自行生成):
https://pan.baidu.com/s/1N7bDHxIM38Vu7x9Z2Qr0og

利用本文代碼訓(xùn)練并生成的模型(對(duì)應(yīng)項(xiàng)目中的model文件夾):
https://pan.baidu.com/s/1GyEpLdM5FSxLYHSc6nq06w

項(xiàng)目簡(jiǎn)介:
(需要預(yù)先安裝pip install captcha==0.1.1,pip install opencv-python,pip install flask, pip install tensorflow/pip install tensorflow-gpu)
本文采用CNN實(shí)現(xiàn)4位定長(zhǎng)驗(yàn)證碼圖片OCR(生成的驗(yàn)證碼固定由隨機(jī)的4位大寫(xiě)字母組成),本質(zhì)上是一張圖片多個(gè)標(biāo)簽的分類(lèi)問(wèn)題(數(shù)據(jù)如下圖所示)


1_CKAN.png
0_HZDZ.png

整體訓(xùn)練邏輯:
1,將圖像傳入到CNN中提取特征
2,將特征圖拉伸輸入到FC layer中得出分類(lèi)預(yù)測(cè)向量
3,通過(guò)sigmoid交叉熵函數(shù)對(duì)預(yù)測(cè)向量和標(biāo)簽向量進(jìn)行訓(xùn)練,得出最終模型(注意:多標(biāo)簽分類(lèi)任務(wù)采用sigmoid,單標(biāo)簽分類(lèi)采用softmax)

整體預(yù)測(cè)邏輯:
1,將圖像傳入到CNN(VGG16)中提取特征
2,將特征圖拉伸輸入到FC layer中得出分類(lèi)預(yù)測(cè)向量
3,將預(yù)測(cè)向量做sigmoid操作,由于驗(yàn)證碼固定是4位,所以將向量切分成4條,從每條中找到最大值,并映射到對(duì)應(yīng)的字母上

制作成web服務(wù):
利用flask框架將整個(gè)項(xiàng)目啟動(dòng)成web服務(wù),使得項(xiàng)目支持http方式調(diào)用
啟動(dòng)服務(wù)后調(diào)用以下地址測(cè)試
http://127.0.0.1:5050/captchaOcr?img_path=./dataset/test/0_HZDZ.png
http://127.0.0.1:5050/captchaOcr?img_path=./dataset/test/1_CKAN.png
后續(xù)優(yōu)化邏輯:
提取特征部分的CNN可以用RNN取代
本方案只能OCR固定長(zhǎng)度文本,后續(xù)采用LSTM+CTC的方式來(lái)OCR非定長(zhǎng)文本

運(yùn)行命令:
自行生成驗(yàn)證碼訓(xùn)練寄(本文生成了10w張,修改self.im_total_num變量):python CnnOcr.py create_dataset
對(duì)數(shù)據(jù)集進(jìn)行訓(xùn)練:python CnnOcr.py train
對(duì)新的圖片進(jìn)行測(cè)試:python CnnOcr.py test
啟動(dòng)成http服務(wù):python CnnOcr.py start

項(xiàng)目目錄結(jié)構(gòu):


目錄結(jié)構(gòu)

訓(xùn)練過(guò)程:


epoch-1

epoch-2

epoch-3

epoch-4

epoch-5

epoch-6

epoch-6-結(jié)束

整體代碼如下:

# coding:utf-8

from captcha.image import ImageCaptcha
import numpy as np
import cv2
import tensorflow as tf
import random, os, sys

from flask import request
from flask import Flask
import json
app = Flask(__name__)


class CnnOcr:
    def __init__(self):
        self.epoch_max = 6  # 最大迭代epoch次數(shù)
        self.batch_size = 64  # 訓(xùn)練時(shí)每個(gè)批次參與訓(xùn)練的圖像數(shù)目,顯存不足的可以調(diào)小
        self.lr = 1e-3  # 初始學(xué)習(xí)率
        self.save_epoch = 1  # 每相隔多少個(gè)epoch保存一次模型


        self.im_width = 128
        self.im_height = 64
        self.im_total_num = 100000  # 總共生成的驗(yàn)證碼圖片數(shù)量
        self.train_max_num = self.im_total_num  # 訓(xùn)練時(shí)讀取的最大圖片數(shù)目
        self.val_num = 50 * self.batch_size  # 不能大于self.train_max_num  做驗(yàn)證集用
        self.words_num = 4  # 每張驗(yàn)證碼圖片上的數(shù)字個(gè)數(shù)
        self.words = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'
        self.label_num = self.words_num * len(self.words)
        self.keep_drop = tf.placeholder(tf.float32)
        self.x = None
        self.y = None


    def captchaOcr(self, img_path):
        """
        驗(yàn)證碼識(shí)別
        :param img_path:
        :return:
        """
        im = cv2.imread(img_path)
        im = cv2.resize(im, (self.im_width, self.im_height))
        im = [im]
        im = np.array(im, dtype=np.float32)
        im -= 147
        output = self.sess.run(self.max_idx_p, feed_dict={self.x: im, self.keep_drop: 1.})
        ret = ''
        for i in output.tolist()[0]:
            ret = ret + self.words[int(i)]
        return ret


    def test(self, img_path):
        """
        測(cè)試接口
        :param img_path:
        :return:
        """
        self.x = tf.placeholder(tf.float32, [None, self.im_height, self.im_width, 3])  # 輸入數(shù)據(jù)
        self.pred = self.cnnNet()
        self.output = tf.nn.sigmoid(self.pred)
        self.predict = tf.reshape(self.pred, [-1, self.words_num, len(self.words)])
        self.max_idx_p = tf.argmax(self.predict, 2)

        saver = tf.train.Saver()
        # tfconfig = tf.ConfigProto(allow_soft_placement=True)
        # tfconfig.gpu_options.per_process_gpu_memory_fraction = 0.3  # 占用顯存的比例
        # self.ses = tf.Session(config=tfconfig)
        self.sess = tf.Session()
        self.sess.run(tf.global_variables_initializer())  # 全局tf變量初始化

        # 加載w,b參數(shù)
        saver.restore(self.sess, './model/CnnOcr-6')
        im = cv2.imread(img_path)
        im = cv2.resize(im, (self.im_width, self.im_height))
        im = [im]
        im = np.array(im, dtype=np.float32)
        im -= 147
        output = self.sess.run(self.max_idx_p, feed_dict={self.x: im, self.keep_drop: 1.})
        ret = ''
        for i in output.tolist()[0]:
            ret = ret + self.words[int(i)]
        print(ret)


    def train(self):
        x_train_list, y_train_list, x_val_list, y_val_list = self.getTrainDataset()

        print('開(kāi)始轉(zhuǎn)換tensor隊(duì)列')
        x_train_list_tensor = tf.convert_to_tensor(x_train_list, dtype=tf.string)
        y_train_list_tensor = tf.convert_to_tensor(y_train_list, dtype=tf.float32)

        x_val_list_tensor = tf.convert_to_tensor(x_val_list, dtype=tf.string)
        y_val_list_tensor = tf.convert_to_tensor(y_val_list, dtype=tf.float32)

        x_train_queue = tf.train.slice_input_producer(tensor_list=[x_train_list_tensor], shuffle=False)
        y_train_queue = tf.train.slice_input_producer(tensor_list=[y_train_list_tensor], shuffle=False)

        x_val_queue = tf.train.slice_input_producer(tensor_list=[x_val_list_tensor], shuffle=False)
        y_val_queue = tf.train.slice_input_producer(tensor_list=[y_val_list_tensor], shuffle=False)

        train_im, train_label = self.dataset_opt(x_train_queue, y_train_queue)
        train_batch = tf.train.batch(tensors=[train_im, train_label], batch_size=self.batch_size, num_threads=2)

        val_im, val_label = self.dataset_opt(x_val_queue, y_val_queue)
        val_batch = tf.train.batch(tensors=[val_im, val_label], batch_size=self.batch_size, num_threads=2)

        print('開(kāi)啟訓(xùn)練')
        self.learning_rate = tf.placeholder(dtype=tf.float32)  # 動(dòng)態(tài)學(xué)習(xí)率
        self.x = tf.placeholder(tf.float32, [None, self.im_height, self.im_width, 3])  # 訓(xùn)練數(shù)據(jù)
        self.y = tf.placeholder(tf.float32, [None, self.label_num])  # 標(biāo)簽
        self.pred = self.cnnNet()
        self.loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=self.pred, labels=self.y))
        self.optimizer = tf.train.AdamOptimizer(learning_rate=self.learning_rate).minimize(self.loss)

        self.predict = tf.reshape(self.pred, [-1, self.words_num, len(self.words)])
        self.max_idx_p = tf.argmax(self.predict, 2)

        self.y_predict = tf.reshape(self.y, [-1, self.words_num, len(self.words)])
        self.max_idx_l = tf.argmax(self.y_predict, 2)

        self.correct_pred = tf.equal(self.max_idx_p, self.max_idx_l)
        self.accuracy = tf.reduce_mean(tf.cast(self.correct_pred, tf.float32))

        with tf.Session() as self.sess:
            # 全局tf變量初始化
            self.sess.run(tf.global_variables_initializer())
            coordinator = tf.train.Coordinator()
            threads = tf.train.start_queue_runners(sess=self.sess, coord=coordinator)

            # 模型保存
            saver = tf.train.Saver()

            batch_max = len(x_train_list) // self.batch_size
            total_step = 1
            for epoch_num in range(self.epoch_max):
                lr = self.lr * (1 - (epoch_num/self.epoch_max) ** 2)  # 動(dòng)態(tài)學(xué)習(xí)率
                for batch_num in range(batch_max):
                    x_train_tmp, y_train_tmp = self.sess.run(train_batch)
                    # print(x_train_tmp.shape, y_train_tmp.shape)
                    # sys.exit()

                    self.sess.run(self.optimizer, feed_dict={self.x: x_train_tmp, self.y: y_train_tmp, self.learning_rate: lr, self.keep_drop: .5})

                    # 輸出評(píng)價(jià)標(biāo)準(zhǔn)
                    if total_step % 50 == 0 or total_step == 1:
                        print()
                        print('epoch:%d/%d batch:%d/%d step:%d lr:%.10f' % ((epoch_num + 1), self.epoch_max, (batch_num + 1), batch_max, total_step, lr))

                        # 輸出訓(xùn)練集評(píng)價(jià)
                        train_loss, train_acc = self.sess.run([self.loss, self.accuracy], feed_dict={self.x: x_train_tmp, self.y: y_train_tmp, self.keep_drop: 1.})
                        print('train_loss:%.10f  train_acc:%.10f' % (np.mean(train_loss), train_acc))

                        # 輸出驗(yàn)證集評(píng)價(jià)
                        val_loss_list, val_acc_list = [], []
                        for i in range(int(self.val_num/self.batch_size)):
                            x_val_tmp, y_val_tmp = self.sess.run(val_batch)
                            val_loss, val_acc = self.sess.run([self.loss, self.accuracy], feed_dict={self.x: x_val_tmp, self.y: y_val_tmp, self.keep_drop: 1.})
                            val_loss_list.append(np.mean(val_loss))
                            val_acc_list.append(np.mean(val_acc))
                        print('  val_loss:%.10f    val_acc:%.10f' % (np.mean(val_loss), np.mean(val_acc)))

                    total_step += 1

                # 保存模型
                if (epoch_num + 1) % self.save_epoch == 0:
                    print('正在保存模型:')
                    saver.save(self.sess, './model/CnnOcr', global_step=(epoch_num + 1))
            coordinator.request_stop()
            coordinator.join(threads)



    def cnnNet(self):
        """
        cnn網(wǎng)絡(luò)
        :return:
        """
        weight = {
            # 輸入 128*64*3

            # 第一層
            'wc1_1': tf.get_variable('wc1_1', [5, 5, 3, 32]),  # 卷積 輸出:128*64*32
            'wc1_2': tf.get_variable('wc1_2', [5, 5, 32, 32]),  # 卷積 輸出:128*64*32
            # 池化 輸出:64*32*32

            # 第二層
            'wc2_1': tf.get_variable('wc2_1', [5, 5, 32, 64]),  # 卷積 輸出:64*32*64
            'wc2_2': tf.get_variable('wc2_2', [5, 5, 64, 64]),  # 卷積 輸出:64*32*64
            # 池化 輸出:32*16*64

            # 第三層
            'wc3_1': tf.get_variable('wc3_1', [3, 3, 64, 64]),  # 卷積 輸出:32*16*256
            'wc3_2': tf.get_variable('wc3_2', [3, 3, 64, 64]),  # 卷積 輸出:32*16*256
            # 池化 輸出:16*8*256

            # 第四層
            'wc4_1': tf.get_variable('wc4_1', [3, 3, 64, 64]),  # 卷積 輸出:16*8*64
            'wc4_2': tf.get_variable('wc4_2', [3, 3, 64, 64]),  # 卷積 輸出:16*8*64
            # 池化 輸出:8*4*64

            # 全鏈接第一層
            'wfc_1': tf.get_variable('wfc_1', [8*4*64, 2048]),

            # 全鏈接第二層
            'wfc_2': tf.get_variable('wfc_2', [2048, 2048]),

            # 全鏈接第三層
            'wfc_3': tf.get_variable('wfc_3', [2048, self.label_num]),
        }

        biase = {
            # 第一層
            'bc1_1': tf.get_variable('bc1_1', [32]),
            'bc1_2': tf.get_variable('bc1_2', [32]),

            # 第二層
            'bc2_1': tf.get_variable('bc2_1', [64]),
            'bc2_2': tf.get_variable('bc2_2', [64]),

            # 第三層
            'bc3_1': tf.get_variable('bc3_1', [64]),
            'bc3_2': tf.get_variable('bc3_2', [64]),

            # 第四層
            'bc4_1': tf.get_variable('bc4_1', [64]),
            'bc4_2': tf.get_variable('bc4_2', [64]),

            # 全鏈接第一層
            'bfc_1': tf.get_variable('bfc_1', [2048]),

            # 全鏈接第二層
            'bfc_2': tf.get_variable('bfc_2', [2048]),

            # 全鏈接第三層
            'bfc_3': tf.get_variable('bfc_3', [self.label_num]),
        }

        # 第一層
        net = tf.nn.conv2d(self.x, weight['wc1_1'], [1, 1, 1, 1], 'SAME')  # 卷積
        net = tf.nn.bias_add(net, biase['bc1_1'])
        net = tf.nn.relu(net)  # 加b 然后 激活
        print('conv1', net)
        net = tf.nn.max_pool(net, [1, 2, 2, 1], [1, 2, 2, 1], padding='VALID')  # 池化
        print('pool1', net)

        # 第二層
        net = tf.nn.conv2d(net, weight['wc2_1'], [1, 1, 1, 1], padding='SAME')  # 卷積
        net = tf.nn.bias_add(net, biase['bc2_1'])
        net = tf.nn.relu(net)  # 加b 然后 激活
        print('conv2', net)
        net = tf.nn.max_pool(net, [1, 2, 2, 1], [1, 2, 2, 1], padding='VALID')  # 池化
        print('pool2', net)

        # 第三層
        net = tf.nn.conv2d(net, weight['wc3_1'], [1, 1, 1, 1], padding='SAME')  # 卷積
        net = tf.nn.bias_add(net, biase['bc3_1'])
        net = tf.nn.relu(net)  # 加b 然后 激活
        print('conv3', net)
        net = tf.nn.max_pool(net, [1, 2, 2, 1], [1, 2, 2, 1], padding='VALID')  # 池化
        print('pool3', net)

        # 第四層
        net = tf.nn.conv2d(net, weight['wc4_1'], [1, 1, 1, 1], padding='SAME')  # 卷積
        net = tf.nn.bias_add(net, biase['bc4_1'])
        net = tf.nn.relu(net)  # 加b 然后 激活
        print('conv4', net)
        net = tf.nn.max_pool(net, [1, 2, 2, 1], [1, 2, 2, 1], padding='VALID')  # 池化
        print('pool4', net)

        # 拉伸flatten,把多個(gè)圖片同時(shí)分別拉伸成一條向量
        net = tf.reshape(net, shape=[-1, weight['wfc_1'].get_shape()[0]])
        print('拉伸flatten', net)

        # 全鏈接層
        # fc第一層
        net = tf.matmul(net, weight['wfc_1']) + biase['bfc_1']
        net = tf.nn.dropout(net, self.keep_drop)
        net = tf.nn.relu(net)

        print('fc第一層', net)
        # fc第二層
        net = tf.matmul(net, weight['wfc_2']) + biase['bfc_2']
        net = tf.nn.dropout(net, self.keep_drop)
        net = tf.nn.relu(net)

        print('fc第二層', net)
        # fc第三層
        net = tf.matmul(net, weight['wfc_3']) + biase['bfc_3']
        print('fc第三層', net)
        return net


    def getTrainDataset(self):
        """
        整理數(shù)據(jù)集,把圖像resize為128*64*3,訓(xùn)練集做成self.im_total_num*128*64*3,把label做成0,1向量形式
        :return:
        """
        train_data_list = os.listdir('./dataset/train/')
        print('共有%d張訓(xùn)練圖片, 讀取%d張:' % (len(train_data_list), self.train_max_num))
        random.shuffle(train_data_list)  # 打亂順序

        y_val_list, y_train_list = [], []
        x_val_list = train_data_list[:self.val_num]
        for x_val in x_val_list:
            words_tmp = x_val.split('.')[0].split('_')[1]
            y_val_list.append([1 if _w == w else 0 for w in words_tmp for _w in self.words])

        x_train_list = train_data_list[self.val_num:self.train_max_num]
        for x_train in x_train_list:
            words_tmp = x_train.split('.')[0].split('_')[1]
            y_train_list.append([1 if _w == w else 0 for w in words_tmp for _w in self.words])

        return x_train_list, y_train_list, x_val_list, y_val_list


    def createCaptchaDataset(self):
        """
        生成訓(xùn)練用圖片數(shù)據(jù)集
        :return:
        """
        image = ImageCaptcha(width=self.im_width, height=self.im_height, font_sizes=(56,))
        for i in range(self.im_total_num):
            words_tmp = ''
            for j in range(self.words_num):
                words_tmp = words_tmp + random.choice(self.words)
            print(words_tmp, type(words_tmp))
            im_path = './dataset/train/%d_%s.png' % (i, words_tmp)
            print(im_path)
            image.write(words_tmp, im_path)
        return True


    def dataset_opt(self, x_train_queue, y_train_queue):
        """
        處理圖片和標(biāo)簽
        :param queue:
        :return:
        """
        queue = x_train_queue[0]
        contents = tf.read_file('./dataset/train/' + queue)
        im = tf.image.decode_jpeg(contents)
        im = tf.image.resize_images(images=im, size=[self.im_height, self.im_width])
        im = tf.reshape(im, tf.stack([self.im_height, self.im_width, 3]))
        im -= 147  # 去均值化
        # im /= 255  # 將像素處理在0~1之間,加速收斂
        # im -= 0.5  # 將像素處理在-0.5~0.5之間
        return im, y_train_queue[0]




if __name__ == '__main__':
    opt_type = sys.argv[1:][0]

    instance = CnnOcr()

    if opt_type == 'create_dataset':
        instance.createCaptchaDataset()
    elif opt_type == 'train':
        instance.train()
    elif opt_type == 'test':
        instance.test('./dataset/test/0_HZDZ.png')
    elif opt_type == 'start':
        # 將session持久化到內(nèi)存中
        instance.test('./dataset/test/0_HZDZ.png')

        # 啟動(dòng)web服務(wù)
        # http://127.0.0.1:5050/captchaOcr?img_path=./dataset/test/2_SYVD.png
        @app.route('/captchaOcr', methods=['GET'])
        def captchaOcr():
            img_path = request.args.to_dict().get('img_path')
            print(img_path)
            ret = instance.captchaOcr(img_path)
            print(ret)
            return json.dumps({'img_path': img_path, 'ocr_ret': ret})

        app.run(host='0.0.0.0', port=5050, debug=False)
最后編輯于
?著作權(quán)歸作者所有,轉(zhuǎn)載或內(nèi)容合作請(qǐng)聯(lián)系作者
【社區(qū)內(nèi)容提示】社區(qū)部分內(nèi)容疑似由AI輔助生成,瀏覽時(shí)請(qǐng)結(jié)合常識(shí)與多方信息審慎甄別。
平臺(tái)聲明:文章內(nèi)容(如有圖片或視頻亦包括在內(nèi))由作者上傳并發(fā)布,文章內(nèi)容僅代表作者本人觀點(diǎn),簡(jiǎn)書(shū)系信息發(fā)布平臺(tái),僅提供信息存儲(chǔ)服務(wù)。

相關(guān)閱讀更多精彩內(nèi)容

友情鏈接更多精彩內(nèi)容