def __call__(self, segments):
return \
C(
self.color,
' '.join(
[
color_symbol('IYellow') + segment[5:] + color_symbol(self.color)
if segment.startswith('tag: ')
else segment
for segment
in self.parse(segments)
]
)
)
def read_dir_to_list(p):
"""
加载目录下的全部文件，经测试在普通机械硬盘下加载 10W 张图片（平均每张1.8kb）耗时 500 s .... =.=
:param p:
:return:
"""
st = clock()
C.yellow("> start to load image from %s" % p)
imgs, captchas = [], []
for _, _, names in os.walk(p):
C.yellow("> get %d files..." % len(names))
for n in names:
captchas.append(n.replace(Config.SUFFIX, ""))
d = Image.open(os.path.join(p, n))
imgs.append(np.asarray(d))
C.yellow("> success load %d images, used %f seconds ..." %
(len(imgs), clock() - st))
return np.asarray(imgs, dtype=np.float32), np.asarray(captchas)
def main():
C.bgGreen(":::: Welcome to use %s (Version=%s) ::::" %
(Config.NAME, Config.VERSION))
print()
train()
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))
def __call__(self, segment):
return C(self.color, self.render(self.parse(segment)))