def predicted(test_dir=Config.DIR_TEST):
"""
对某个目录下的验证码文件进行预测（文件名即为验证码文本）
:param test_dir:
:return:
"""
C.bgYellow(">>>> Begin to prediction <<<<")
if not os.path.exists(test_dir) or os.path.isfile(test_dir):
return C.red("directory not exist: %s" % test_dir)
# 检查模型
model = tf.train.get_checkpoint_state(Config.DIR_MODEL)
if model:
data_x, data_y = load_data(test_dir)
print("total %d captcha to be predicted..." % len(data_x))
with tf.Session() as session:
saver.restore(session, model.model_checkpoint_path)
C.green("restore session from %s" % model.model_checkpoint_path)
y_pred, acc_val = session.run([prediction, accuracy],
feed_dict={
X: data_x,
Y: data_y,
keep_prob: Config.NN_KEEP_PROB
})
correct_count = 0
for i in range(len(data_y)):
origin_code, pred_code = vec2text(data_y[i]), vec2text(
y_pred[i])
correct = origin_code == pred_code
if correct:
correct_count += 1
C.green("{} predicted to {} : {}".format(
origin_code, pred_code, correct))
else:
C.red("{} predicted to {} : {}".format(
origin_code, pred_code, correct))
C.yellow("\nAccuracy = {:5.2f} % (Correct_count={})".format(
acc_val * 100, correct_count))
else:
C.red("Model not exist on %s" % Config.DIR_MODEL)
def train():
C.bgGreen(
">>>> Begin to train ( keep_prob={}, learning_rate={} ) <<<<".format(
Config.NN_KEEP_PROB, Config.NN_LEARNING_RATE))
# 读取数据
data_x, data_y = load_data(Config.DIR_TRAIN)
C.green("data_x shape=%s data_y shape=%s" % (data_x.shape, data_y.shape))
C.green("split data into Train 80%, Dev 10%, Test 10% ...")
train_x, test_x, train_y, test_y = train_test_split(data_x,
data_y,
test_size=0.2,
random_state=40)
dev_x, test_x, dev_y, test_y = train_test_split(test_x,
test_y,
test_size=0.5,
random_state=40)
step_train = int(len(train_x) / Config.NN_BATCH_SIZE)
print("train step = %d" % step_train)
startOn = clock()
with tf.Session() as session:
session.run(tf.global_variables_initializer())
for epoch in range(Config.NN_EPOCH):
C.yellow("\nbegin the {:<4} epoch training...".format(epoch))
for batch in range(step_train):
start = batch * Config.NN_BATCH_SIZE
end = start + Config.NN_BATCH_SIZE
dx, dy = train_x[start:end], train_y[start:end]
_, loss_val, acc_val = session.run([train_op, loss, accuracy],
feed_dict={
X:
dx,
Y:
dy,
keep_prob:
Config.NN_KEEP_PROB
})
if batch % 5 == 0:
print(
"\tepoch={:4} batch = {:3} loss= {:<25} accuracy = {:5.2f} %"
.format(epoch, batch, loss_val, acc_val * 100))
loss_val, acc_val = session.run([loss, accuracy],
feed_dict={
X: dev_x,
Y: dev_y,
keep_prob: Config.NN_KEEP_PROB
})
C.yellow(
"\t[Dev Test] epoch={:5} loss= {:<25} accuracy = {:5.2f} %".
format(epoch, loss_val, acc_val * 100))
if Config.NN_SAVE and epoch % 10 == 0:
saver.save(session, Config.DIR_MODEL + "/captcha")
C.green("Model saved on %s" % Config.DIR_MODEL)
C.bgYellow("---- Train Done (used time = %.4f seconds) ----" %
(clock() - startOn))
loss_val, acc_val = session.run([loss, accuracy],
feed_dict={
X: test_x,
Y: test_y,
keep_prob: Config.NN_KEEP_PROB
})
C.yellow("\t[Test] loss= {:<25} accuracy = {:5.2f} %".format(
loss_val, acc_val * 100))