這裏對之前訓練好的模型進行測試,代碼如下:
import os
import tensorflow as tf
from PIL import Image
from nets import nets_factory
import numpy as np
import matplotlib.pyplot as plt
# 不同字符數量
CHAR_SET_LEN = 10
# 圖片高度
IMAGE_HEIGHT = 60
# 圖片寬度
IMAGE_WIDTH = 160
# 批次
BATCH_SIZE = 1
# tfrecord文件存放路徑
TFRECORD_FILE = "D:/Tensorflow/captcha/test.tfrecords"
# placeholder
x = tf.placeholder(tf.float32, [None, 224, 224])
# 從tfrecord讀出數據
def read_and_decode(filename):
# 根據文件名生成一個隊列
filename_queue = tf.train.string_input_producer([filename])
reader = tf.TFRecordReader()
# 返回文件名和文件
_, serialized_example = reader.read(filename_queue)
features = tf.parse_single_example(serialized_example,
features={
'image' : tf.FixedLenFeature([], tf.string),
'label0': tf.FixedLenFeature([], tf.int64),
'label1': tf.FixedLenFeature([], tf.int64),
'label2': tf.FixedLenFeature([], tf.int64),
'label3': tf.FixedLenFeature([], tf.int64),
})
# 獲取圖片數據
image = tf.decode_raw(features['image'], tf.uint8)
# 沒有經過預處理的灰度圖
image_raw = tf.reshape(image, [224, 224])
# tf.train.shuffle_batch必須確定shape
image = tf.reshape(image, [224, 224])
# 圖片預處理
image = tf.cast(image, tf.float32) / 255.0
image = tf.subtract(image, 0.5)
image = tf.multiply(image, 2.0)
# 獲取label
label0 = tf.cast(features['label0'], tf.int32)
label1 = tf.cast(features['label1'], tf.int32)
label2 = tf.cast(features['label2'], tf.int32)
label3 = tf.cast(features['label3'], tf.int32)
return image, image_raw, label0, label1, label2, label3
# 獲取圖片數據和標籤
image, image_raw, label0, label1, label2, label3 = read_and_decode(TFRECORD_FILE)
#使用shuffle_batch可以隨機打亂
image_batch, image_raw_batch, label_batch0, label_batch1, label_batch2, label_batch3 = tf.train.shuffle_batch(
[image, image_raw, label0, label1, label2, label3], batch_size = BATCH_SIZE,
capacity = 50000, min_after_dequeue=10000, num_threads=1)
#定義網絡結構
train_network_fn = nets_factory.get_network_fn(
'alexnet_v2',
num_classes=CHAR_SET_LEN,
weight_decay=0.0005,
is_training=False)
with tf.Session() as sess:
# inputs: a tensor of size [batch_size, height, width, channels]
X = tf.reshape(x, [BATCH_SIZE, 224, 224, 1])
# 數據輸入網絡得到輸出值
logits0,logits1,logits2,logits3,end_points = train_network_fn(X)
# 預測值
predict0 = tf.reshape(logits0, [-1, CHAR_SET_LEN])
predict0 = tf.argmax(predict0, 1)
predict1 = tf.reshape(logits1, [-1, CHAR_SET_LEN])
predict1 = tf.argmax(predict1, 1)
predict2 = tf.reshape(logits2, [-1, CHAR_SET_LEN])
predict2 = tf.argmax(predict2, 1)
predict3 = tf.reshape(logits3, [-1, CHAR_SET_LEN])
predict3 = tf.argmax(predict3, 1)
# 初始化
sess.run(tf.global_variables_initializer())
# 載入訓練好的模型
saver = tf.train.Saver()
saver.restore(sess,'./captcha/models/crack_captcha.model-6000')
# 創建一個協調器,管理線程
coord = tf.train.Coordinator()
# 啓動QueueRunner, 此時文件名隊列已經進隊
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
for i in range(10):
# 獲取一個批次的數據和標籤
b_image, b_image_raw, b_label0, b_label1 ,b_label2 ,b_label3 = sess.run([image_batch,
image_raw_batch,
label_batch0,
label_batch1,
label_batch2,
label_batch3])
# 顯示圖片
img=Image.fromarray(b_image_raw[0],'L')
plt.imshow(img)
plt.axis('off')
plt.show()
# 打印標籤
print('label:',b_label0, b_label1 ,b_label2 ,b_label3)
# 預測
label0,label1,label2,label3 = sess.run([predict0,predict1,predict2,predict3], feed_dict={x: b_image})
# 打印預測值
print('predict:',label0,label1,label2,label3)
# 通知其他線程關閉
coord.request_stop()
# 其他所有線程關閉之後,這一函數才能返回
coord.join(threads)
效果如下: