"""基于切割的识别模型"""
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