TensorFlow用训练好的CNN模型检测

已预先训练好了一个识别猫、狗的二分类CNN模型，持久化在了一个cat_vs_dog.pb的文件中。现需要用其进行图片检测。

1.代码

读取训练好的pb文件，用来进行图片的检测，代码为：

import tensorflow as tf
import numpy as np
from PIL import Image

FILEPATH = 'F:/PycharmProjects/deepLearning/cnn_cat&dog/data/image/check/2.jpg'  # 检测文件地址
MODELPATH = 'cat_vs_dog.pb'  # pb文件
SIZE = [224, 224]  # 检测图片规格
DIC = {0: 'cat', 1: 'dog'}  # 定义字典，保存标签和类别的映射关系


# 图片转化为数组
def get_check_data(path, size):
    img = Image.open(path)
    img = img.resize((size[1], size[0]), Image.ANTIALIAS)
    img = np.asarray(img, dtype="int32")
    img = img.astype(np.float32)
    img = img.reshape((img.shape[0], img.shape[1], 3))
    img = img * (1. / 255) - 0.5
    return np.array([img])


# 检测
def check(filename, modelname, size):
    data = get_check_data(filename, size)

    with tf.Graph().as_default():
        output_graph_def = tf.GraphDef()

        with open(modelname, "rb") as f:  # 读取pb文件
            output_graph_def.ParseFromString(f.read())
            _ = tf.import_graph_def(output_graph_def, name="")

        sess = tf.InteractiveSession()
        init = tf.global_variables_initializer()
        sess.run(init)

        input_x = sess.graph.get_tensor_by_name("input:0")
        keep_prob = sess.graph.get_tensor_by_name("keep_prob:0")
        out_y = sess.graph.get_tensor_by_name("fc3:0")
        softmax = tf.nn.softmax(out_y)

        y = sess.run(softmax, feed_dict={input_x: data, keep_prob: 1.})
        prediction_labels = np.argmax(y)
        print("label:", prediction_labels)
        print("this is a %s" % DIC[prediction_labels])
        sess.close()


def main(_):
    check(FILEPATH, MODELPATH, SIZE)


if __name__ == '__main__':
    main(None)

2.结果展示

要检测的图片为下面第二个图片：

结果为：