验证码cnn模型

2019独角兽企业重金招聘Python工程师标准>>>

"""基于切割的识别模型"""

import numpy as np
from PIL import Image
from keras import backend as K
from keras import layers
from keras.models import Sequential

from . import img_util

K.set_image_dim_ordering('tf')

img_width, img_height = 50, 70
chars = "23456789abcdefghiklmnpqrstuvwxyz"


class ZXRModel:
    def __init__(self):
        raise NotImplemented

    def inspect(self, file):
        from keras.utils import plot_model
        plot_model(self.model, to_file=file, show_shapes=True)

    def load_weights(self, weight_filepath):
        self.model.load_weights(weight_filepath)

    def predict(self, img: Image.Image):
        data = self.load_img(img)
        indices = self.model.predict_classes(np.asarray(data), batch_size=len(data), verbose=0)
        return [chars[i] for i in indices]

    def train(self, glob_img_path, save_weights_to=None):
        data, label = self._load_data(glob_img_path)
        self._train_with_data(data, label)
        if save_weights_to:
            self.model.save_weights(save_weights_to)

    def _load_data(self, glob_path):
        from glob import iglob
        from os import path

        img_files = iglob(glob_path)

        print('开始加载训练集:', glob_path)
        data, label = list(), list()
        for file in img_files:
            file_name = path.splitext(path.split(file)[-1])[0].split('.')[0].lower()
            try:
                t = self.load_img(Image.open(file))
                if len(file_name) == len(t):
                    data += t
            except Exception as ex:
                from sys import stderr
                print(ex.args[0], file=stderr)
                continue
            else:
                label += [chars.index(i) for i in file_name]
        return data, label

    def train_with_data(self, data, labels):
        from keras.utils import np_utils
        print('训练集: %d, label: %d' % (len(data), len(labels)))
        self.model.compile(loss='categorical_crossentropy', optimizer='adadelta', metrics=['accuracy'])
        self.model.fit(np.asarray(data),
                       np_utils.to_categorical(labels, len(chars)),
                       batch_size=500, epochs=150, shuffle=True,
                       verbose=1)

    def load_img(self, img):
        segs = img_util.split_zxr(img if img.mode == 'RGB' else img.convert('RGB'))
        return [255 - np.asarray(seg, dtype='int32') for seg in segs]


class SimpleModel(ZXRModel):
    def __init__(self):
        """
        vggnet简化版
        """
        _model = Sequential()
        _model.add(layers.InputLayer(input_shape=(img_height, img_width, 1)))
        for i in range(2):
            _model.add(layers.Convolution2D(9 * 3 ** i, (3, 3), border_mode='valid', activation="relu"))
            _model.add(layers.Convolution2D(9 * 3 ** i, (3, 3), border_mode='valid', activation="relu"))
            _model.add(layers.MaxPooling2D((2, 2)))

        _model.add(layers.Flatten())
        _model.add(layers.Dense(output_dim=128, activation='tanh'))
        _model.add(layers.Dense(len(chars), init='normal', activation='softmax'))
        self.model = _model


class VGGModel(ZXRModel):
    def __init__(self, training=False):
        """定制的vggnet模型"""
        drop_rate = 0.5 if training else 1.0
        _model = Sequential()
        _model.add(layers.InputLayer(input_shape=(img_height, img_width, 3)))
        for i in range(1, 3):
            _model.add(layers.Convolution2D(16 * 2 ** i, (3, 3), border_mode='valid', activation="relu"))
            _model.add(layers.Convolution2D(16 * 2 ** i, (3, 3), border_mode='valid', activation="relu"))
            _model.add(layers.MaxPooling2D((2, 2), strides=(2, 2)))

        _model.add(layers.Flatten())
        _model.add(layers.Dense(1024, activation='relu'))
        _model.add(layers.Dropout(drop_rate))
        _model.add(layers.Dense(1024, activation='relu'))
        _model.add(layers.Dropout(drop_rate))
        _model.add(layers.Dense(output_dim=256, activation='relu'))
        _model.add(layers.Dense(len(chars), init='normal', activation='softmax'))
        self.model = _model


class AlexNet(ZXRModel):
    def __init__(self, training=False):
        drop_rate = 0.5 if training else 1.0
        _model = Sequential()
        _model.add(layers.InputLayer((img_height, img_width, 3)))
        _model.add(layers.Convolution2D(16, (5, 5), strides=(2, 2), activation='relu'))
        _model.add(layers.MaxPool2D((2, 2), strides=(2, 2)))
        _model.add(layers.Convolution2D(32, (3, 3), strides=(1, 1), activation='relu'))
        _model.add(layers.MaxPool2D((2, 2), strides=(2, 2)))

        # for _ in range(3):
        _model.add(layers.Convolution2D(96, (2, 2), strides=(1, 1), activation='relu'))
        # _model.add(layers.Convolution2D(64, (3, 3), strides=(1, 1), activation='relu'))

        _model.add(layers.Flatten())
        _model.add(layers.Dense(512, activation='relu'))
        _model.add(layers.Dropout(drop_rate))
        _model.add(layers.Dense(512, activation='relu'))
        _model.add(layers.Dropout(drop_rate))
        _model.add(layers.Dense(output_dim=256, activation='relu'))
        _model.add(layers.Dense(len(chars), init='normal', activation='softmax'))
        self.model = _model

转载于:https://my.oschina.net/wisedream/blog/1023369