test.py
#coding=utf-8import os#图像读取库from PIL import Image#矩阵运算库import numpy as npimport tensorflow as tf# 数据文件夹data_dir = "E:/Tensorflow-CNN-VehicleIdentification/test"# 训练还是测试True Falsetrain = False# 模型文件路径model_path = "model/image_model"# 从文件夹读取图片和标签到numpy数组中# 标签信息在文件名中，例如1_40.jpg表示该图片的标签为1def read_data(data_dir):datas = []labels = []fpaths = []for fname in os.listdir(data_dir):fpath = os.path.join(data_dir, fname)fpaths.append(fpath)image = Image.open(fpath)data = http://www.kingceram.com/post/np.array(image) / 255.0label = int(fname.split("_")[0])datas.append(data)labels.append(label)datas = np.array(datas)labels = np.array(labels)print("shape of datas: {}\tshape of labels: {}".format(datas.shape, labels.shape))return fpaths, datas, labelsfpaths, datas, labels = read_data(data_dir)# 计算有多少类图片num_classes = len(set(labels))# 定义Placeholder，存放输入和标签datas_placeholder = tf.placeholder(tf.float32, [None, 32, 32, 3])labels_placeholder = tf.placeholder(tf.int32, [None])# 存放DropOut参数的容器，训练时为0.25，测试时为0dropout_placeholdr = tf.placeholder(tf.float32)# 定义卷积层, 20个卷积核, 卷积核大小为5，用Relu激活conv0 = tf.layers.conv2d(datas_placeholder, 20, 5, activation=tf.nn.relu)# 定义max-pooling层，pooling窗口为2x2，步长为2x2pool0 = tf.layers.max_pooling2d(conv0, [2, 2], [2, 2])# 定义卷积层, 40个卷积核, 卷积核大小为4，用Relu激活conv1 = tf.layers.conv2d(pool0, 40, 4, activation=tf.nn.relu)# 定义max-pooling层，pooling窗口为2x2，步长为2x2pool1 = tf.layers.max_pooling2d(conv1, [2, 2], [2, 2])# 定义卷积层, 60个卷积核, 卷积核大小为3，用Relu激活conv2 = tf.layers.conv2d(pool0, 60, 3, activation=tf.nn.relu)# 定义max-pooling层，pooling窗口为2x2，步长为2x2pool2 = tf.layers.max_pooling2d(conv2, [2, 2], [2, 2])# 将3维特征转换为1维向量flatten = tf.layers.flatten(pool2)# 全连接层，转换为长度为100的特征向量fc = tf.layers.dense(flatten, 400, activation=tf.nn.relu)# 加上DropOut，防止过拟合dropout_fc = tf.layers.dropout(fc, dropout_placeholdr)# 未激活的输出层logits = tf.layers.dense(dropout_fc, num_classes)predicted_labels = tf.arg_max(logits, 1)# 利用交叉熵定义损失losses = tf.nn.softmax_cross_entropy_with_logits(labels=tf.one_hot(labels_placeholder, num_classes),logits=logits)# 平均损失mean_loss = tf.reduce_mean(losses)# 定义优化器，指定要优化的损失函数optimizer = tf.train.AdamOptimizer(learning_rate=1e-2).minimize(losses)# 用于保存和载入模型saver = tf.train.Saver()with tf.Session() as sess:if train:print("训练模式")# 如果是训练，初始化参数sess.run(tf.global_variables_initializer())# 定义输入和Label以填充容器，训练时dropout为0.25train_feed_dict = {datas_placeholder: datas,labels_placeholder: labels,dropout_placeholdr: 0.5}for step in range(150):_, mean_loss_val = sess.run([optimizer, mean_loss], feed_dict=train_feed_dict)if step % 10 == 0:print("step = {}\tmean loss = {}".format(step, mean_loss_val))saver.save(sess, model_path)print("训练结束，保存模型到{}".format(model_path))else:print("测试模式")# 如果是测试，载入参数saver.restore(sess, model_path)print("从{}载入模型".format(model_path))# label和名称的对照关系label_name_dict = {0: "巴士",1: "出租车",2: "货车",3: "家用轿车",4: "面包车",5: "吉普车",6: "运动型多功能车",7: "重型货车",8: "赛车",9: "消防车"}# 定义输入和Label以填充容器，测试时dropout为0test_feed_dict = {datas_placeholder: datas,labels_placeholder: labels,dropout_placeholdr: 0}predicted_labels_val = sess.run(predicted_labels, feed_dict=test_feed_dict)# 真实label与模型预测labelfor fpath, real_label, predicted_label in zip(fpaths, labels, predicted_labels_val):# 将label id转换为label名real_label_name = label_name_dict[real_label]predicted_label_name = label_name_dict[predicted_label]print("{}\t{} => {}".format(fpath, real_label_name, predicted_label_name))# 正确次数correct_number=0# 计算正确率forfpath, real_label, predicted_label in zip(fpaths, labels, predicted_labels_val):if real_label==predicted_label:correct_number+=1correct_rate=correct_number/200print('正确率: {:.2%}'.format(correct_rate))