基于TensorFlow图像分类实现

train.py

训练集图片文件名称中存在类型即可。根据需分类类型修改# 训练集生成和# 测试集生成代码块中的读取方式。

import os import numpy as np import tensorflow as tf from network import Network from datagenerator import ImageDataGenerator from datetime import datetime import glob from tensorflow.contrib.data import Iteratorlearning_rate &＃61; 1e-4 num_epochs &＃61; 1 # 迭代次数 batch_size &＃61; 50 dropout_rate &＃61; 0.5 num_classes &＃61; 5 # 类别数量 display_step &＃61; 5filewriter_path &＃61; "tmp/tensorboard_test" # tensorboard文件路径 checkpoint_path &＃61; "tmp/checkpoints_test" # 模型和参数路径if not os.path.isdir(checkpoint_path):os.mkdir(checkpoint_path)train_image_path &＃61; &＃39;train/&＃39; # 训练集数据路径 test_image_path &＃61; &＃39;test/&＃39; # 测试集数据路径label_path &＃61; [] test_label &＃61; []# 训练集生成 image_path &＃61; np.array(glob.glob(train_image_path &＃43; &＃39;*.jpg&＃39;)).tolist() for i in range(len(image_path)):if &＃39;Bus&＃39; in image_path[i]:label_path.append(0)elif &＃39;Microbus&＃39; in image_path[i]:label_path.append(1)elif &＃39;Sedan&＃39; in image_path[i]:label_path.append(2)elif &＃39;SUV&＃39; in image_path[i]:label_path.append(3)elif &＃39;Truck&＃39; in image_path[i]:label_path.append(4)# 测试集生成 test_image &＃61; np.array(glob.glob(test_image_path &＃43; &＃39;*.jpg&＃39;)).tolist() for i in range(len(test_image)):if &＃39;Bus&＃39; in image_path[i]:test_label.append(0)elif &＃39;Microbus&＃39; in image_path[i]:test_label.append(1)elif &＃39;Sedan&＃39; in image_path[i]:test_label.append(2)elif &＃39;SUV&＃39; in image_path[i]:test_label.append(3)elif &＃39;Truck&＃39; in image_path[i]:test_label.append(4)# 调用图片生成器&＃xff0c;把训练集图片转换成三维数组 tr_data &＃61; ImageDataGenerator(images&＃61;image_path,labels&＃61;label_path,batch_size&＃61;batch_size,num_classes&＃61;num_classes)# 调用图片生成器&＃xff0c;把测试集图片转换成三维数组 test_data &＃61; ImageDataGenerator(images&＃61;test_image,labels&＃61;test_label,batch_size&＃61;batch_size,num_classes&＃61;num_classes,shuffle&＃61;False)with tf.name_scope(&＃39;input&＃39;):# 定义迭代器iterator &＃61; Iterator.from_structure(tr_data.data.output_types,tr_data.data.output_shapes)training_initalize&＃61;iterator.make_initializer(tr_data.data)testing_initalize&＃61;iterator.make_initializer(test_data.data)# 定义每次迭代的数据next_batch &＃61; iterator.get_next()x &＃61; tf.placeholder(tf.float32, [batch_size, 224, 224, 3]) y &＃61; tf.placeholder(tf.float32, [batch_size, num_classes]) keep_prob &＃61; tf.placeholder(tf.float32)# 图片数据通过网络处理 model &＃61; Network(x, keep_prob, num_classes)# 执行整个网络图 score &＃61; model.fc8with tf.name_scope(&＃39;loss&＃39;):# 损失函数loss &＃61; tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits&＃61;score, labels&＃61;y))tf.summary.scalar(&＃39;loss&＃39;, loss)with tf.name_scope(&＃39;optimizer&＃39;):# 优化器train_op &＃61; tf.train.AdamOptimizer(learning_rate).minimize(loss)# 定义网络精确度 with tf.name_scope("accuracy"):correct_pred &＃61; tf.equal(tf.argmax(score, 1), tf.argmax(y, 1))accuracy &＃61; tf.reduce_mean(tf.cast(correct_pred, tf.float32))tf.summary.scalar(&＃39;accuracy&＃39;, accuracy)# 把精确度加入到Tensorboardmerged_summary &＃61; tf.summary.merge_all() writer &＃61; tf.summary.FileWriter(filewriter_path) saver &＃61; tf.train.Saver()# 定义一代的迭代次数 train_batches_per_epoch &＃61; int(np.floor(tr_data.data_size / batch_size)) test_batches_per_epoch &＃61; int(np.floor(test_data.data_size / batch_size))with tf.Session() as sess:sess.run(tf.global_variables_initializer())#saver &＃61; tf.train.Saver()#saver.restore(sess, "./tmp/checkpoints_t18/model_epoch10.ckpt")# 把模型图加入Tensorboardwriter.add_graph(sess.graph)print("{} 训练开始".format(datetime.now()))print("{} Tensorboard at --logdir {}".format(datetime.now(), filewriter_path))# 迭代所有次数for epoch in range(num_epochs):sess.run(training_initalize)print("{} 迭代{}次开始".format(datetime.now(), epoch &＃43; 1))#开始训练每一代for step in range(train_batches_per_epoch):img_batch, label_batch &＃61; sess.run(next_batch)sess.run(train_op, feed_dict&＃61;{x: img_batch, y: label_batch, keep_prob: dropout_rate})if step % display_step &＃61;&＃61; 0:s &＃61; sess.run(merged_summary, feed_dict&＃61;{x: img_batch,y: label_batch,keep_prob: 1.})writer.add_summary(s, epoch * train_batches_per_epoch &＃43; step)# 测试模型精确度print("{} 测试精度".format(datetime.now()))sess.run(testing_initalize)test_acc &＃61; 0.test_count &＃61; 0for _ in range(test_batches_per_epoch):img_batch, label_batch &＃61; sess.run(next_batch)acc &＃61; sess.run(accuracy, feed_dict&＃61;{x: img_batch,y: label_batch,keep_prob: 1.0})test_acc &＃43;&＃61; acctest_count &＃43;&＃61; 1test_acc /&＃61; test_countprint("{} 精度 &＃61; {:.4f}".format(datetime.now(), test_acc))# 把训练好的模型存储起来print("{} 保存模型".format(datetime.now()))checkpoint_name &＃61; os.path.join(checkpoint_path, &＃39;model_epoch&＃39; &＃43; str(epoch &＃43; 1) &＃43; &＃39;.ckpt&＃39;)save_path &＃61; saver.save(sess, checkpoint_name)print("{} 迭代{}次结束".format(datetime.now(), epoch &＃43; 1), save_path)

network.py

import tensorflow as tf import numpy as npclass Network(object):def __init__(self, x, keep_prob, num_classes):self.X &＃61; xself.NUM_CLASSES &＃61; num_classesself.KEEP_PROB &＃61; keep_probself.create()def create(self):#卷积层1conv1_1 &＃61; conv(self.X, 9, 9, 128, 4, 4, padding&＃61;&＃39;VALID&＃39;, name&＃61;&＃39;conv1_1&＃39;)pool1 &＃61; max_pool(conv1_1, 2, 2, 2, 2, padding&＃61;&＃39;SAME&＃39;, name&＃61;&＃39;pool1&＃39;)# 卷积层2conv2_1 &＃61; conv(pool1, 4, 4, 256, 1, 1, padding&＃61;&＃39;VALID&＃39;, name&＃61;&＃39;conv2_1&＃39;)pool2 &＃61; max_pool(conv2_1, 2, 2, 2, 2, padding&＃61;&＃39;SAME&＃39;, name&＃61;&＃39;pool2&＃39;)# 卷积层3conv3_1 &＃61; conv(pool2, 3, 3, 512, 1, 1, padding&＃61;&＃39;SAME&＃39;, name&＃61;&＃39;conv3_1&＃39;)conv3_2 &＃61; conv(conv3_1, 3, 3, 512, 1, 1, padding&＃61;&＃39;SAME&＃39;, name&＃61;&＃39;conv3_2&＃39;)pool3 &＃61; max_pool(conv3_2, 2, 2, 2, 2, padding&＃61;&＃39;SAME&＃39;, name&＃61;&＃39;pool3&＃39;)# 卷积层4conv4_1 &＃61; conv(pool3, 3, 3, 256, 1, 1, padding&＃61;&＃39;SAME&＃39;, name&＃61;&＃39;conv4_1&＃39;)#pool4 &＃61; max_pool(conv4_2, 2, 2, 2, 2, padding&＃61;&＃39;SAME&＃39;, name&＃61;&＃39;pool4&＃39;)flattened &＃61; tf.reshape(conv4_1, [-1, 6*6*256])# 全链接6fc6 &＃61; fc(flattened, 6*6*256, 4096, name&＃61;&＃39;fc6&＃39;)dropout6 &＃61; dropout(fc6, self.KEEP_PROB)# 全链接7fc7 &＃61; fc(dropout6, 4096, 4096, name&＃61;&＃39;fc7&＃39;)dropout7 &＃61; dropout(fc7, self.KEEP_PROB)# 全链接8self.fc8 &＃61; fc(dropout7, 4096, self.NUM_CLASSES, name&＃61;&＃39;fc8&＃39;, relu&＃61;False)def conv(x, filter_height, filter_width, num_filters, stride_y, stride_x, name,padding&＃61;&＃39;SAME&＃39;):input_channels &＃61; int(x.get_shape()[-1])convolve &＃61; lambda i, k: tf.nn.conv2d(i, k,strides&＃61;[1, stride_y, stride_x, 1],padding&＃61;padding)with tf.variable_scope(name) as scope:weights &＃61; tf.get_variable(&＃39;weights&＃39;, shape&＃61;[filter_height,filter_width,input_channels,num_filters])biases &＃61; tf.get_variable(&＃39;biases&＃39;, shape&＃61;[num_filters])conv &＃61; convolve(x, weights)bias &＃61; tf.reshape(tf.nn.bias_add(conv, biases), tf.shape(conv))relu &＃61; tf.nn.relu(bias, name&＃61;scope.name)return reludef fc(x, num_in, num_out, name, relu&＃61;True):with tf.variable_scope(name) as scope:weights &＃61; tf.get_variable(&＃39;weights&＃39;, shape&＃61;[num_in, num_out],trainable&＃61;True)biases &＃61; tf.get_variable(&＃39;biases&＃39;, [num_out], trainable&＃61;True)act &＃61; tf.nn.xw_plus_b(x, weights, biases, name&＃61;scope.name)if relu:relu &＃61; tf.nn.relu(act)return reluelse:return actdef max_pool(x, filter_height, filter_width, stride_y, stride_x, name,padding&＃61;&＃39;SAME&＃39;):return tf.nn.max_pool(x, ksize&＃61;[1, filter_height, filter_width, 1],strides&＃61;[1, stride_y, stride_x, 1],padding&＃61;padding, name&＃61;name)def lrn(x, radius, alpha, beta, name, bias&＃61;1.0):return tf.nn.local_response_normalization(x, depth_radius&＃61;radius,alpha&＃61;alpha, beta&＃61;beta,bias&＃61;bias, name&＃61;name)def dropout(x, keep_prob):return tf.nn.dropout(x, keep_prob)

datagenerator.py

import tensorflow as tf import numpy as npfrom tensorflow.python.framework import dtypes from tensorflow.python.framework.ops import convert_to_tensor from tensorflow.contrib.data import DatasetVGG_MEAN &＃61; tf.constant([123.68, 116.779, 103.939], dtype&＃61;tf.float32)# 把图片数据转化为三维矩阵 class ImageDataGenerator(object):def __init__(self, images, labels, batch_size, num_classes, shuffle&＃61;True):self.img_paths &＃61; imagesself.labels &＃61; labelsself.num_classes &＃61; num_classesself.data_size &＃61; len(self.labels)self.pointer &＃61; 0if shuffle:self._shuffle_lists()self.img_paths &＃61; convert_to_tensor(self.img_paths, dtype&＃61;dtypes.string)self.labels &＃61; convert_to_tensor(self.labels, dtype&＃61;dtypes.int32)data &＃61; Dataset.from_tensor_slices((self.img_paths, self.labels))data &＃61; data.map(self._parse_function_train, num_threads&＃61;8,output_buffer_size&＃61;100 * batch_size)data &＃61; data.batch(batch_size)self.data &＃61; data# 打乱图片顺序def _shuffle_lists(self):path &＃61; self.img_pathslabels &＃61; self.labelspermutation &＃61; np.random.permutation(self.data_size)self.img_paths &＃61; []self.labels &＃61; []for i in permutation:self.img_paths.append(path[i])self.labels.append(labels[i])# 把图片生成三维数组&＃xff0c;以及把标签转化为向量def _parse_function_train(self, filename, label):one_hot &＃61; tf.one_hot(label, self.num_classes)img_string &＃61; tf.read_file(filename)img_decoded &＃61; tf.image.decode_png(img_string, channels&＃61;3)img_resized &＃61; tf.image.resize_images(img_decoded, [224, 224])img_centered &＃61; tf.subtract(img_resized, VGG_MEAN)img_bgr &＃61; img_centered[:, :, ::-1]return img_bgr, one_hot

validate_image.py

import tensorflow as tf from network import Network import matplotlib.pyplot as plt import numpy as np import glob from tensorflow.python.framework import dtypes from tensorflow.python.framework.ops import convert_to_tensor from tensorflow.contrib.data import Dataset from tensorflow.contrib.data import IteratorVGG_MEAN &＃61; tf.constant([123.68, 116.779, 103.939], dtype&＃61;tf.float32) class_name &＃61; [&＃39;Bus&＃39;, &＃39;Microbus&＃39;, &＃39;Sedan&＃39;, &＃39;SUV&＃39;, &＃39;Truck&＃39;] validate_image_path &＃61; &＃39;validate/&＃39; # 指定验证集数据路径&＃xff08;根据实际情况指定验证数据集的路径&＃xff09;x &＃61; tf.placeholder(tf.float32, [1, 224, 224, 3]) model &＃61; Network(x, 1, 5) score &＃61; tf.nn.softmax(model.fc8) max &＃61; tf.arg_max(score, 1)t_num &＃61; 0 f_num &＃61; 0 image_path &＃61; np.array(glob.glob(validate_image_path &＃43; &＃39;*.jpg&＃39;)).tolist() fo &＃61; open("false.txt", "w")with tf.Session() as sess:sess.run(tf.global_variables_initializer())saver &＃61; tf.train.Saver()saver.restore(sess, "./tmp/checkpoints_t16/model_epoch7.ckpt")for i in range(len(image_path)):img_string &＃61; tf.read_file(image_path[i])img_decoded &＃61; tf.image.decode_png(img_string, channels&＃61;3)img_resized &＃61; tf.image.resize_images(img_decoded, [224, 224])img_resized &＃61; img_resized[:, :, ::-1]img_resized &＃61; np.asarray(img_resized.eval(), dtype&＃61;&＃39;uint8&＃39;)img_resized &＃61; img_resized.reshape((1, 224, 224, 3))prob &＃61; sess.run(max, feed_dict&＃61;{x: img_resized})[0]t &＃61; -1if &＃39;Bus&＃39; in image_path[i]:t &＃61; 0elif &＃39;Microbus&＃39; in image_path[i]:t &＃61; 1elif &＃39;Sedan&＃39; in image_path[i]:t &＃61; 2elif &＃39;SUV&＃39; in image_path[i]:t &＃61; 3elif &＃39;Truck&＃39; in image_path[i]:t &＃61; 4if t &＃61;&＃61; prob:t_num &＃43;&＃61; 1else:f_num &＃43;&＃61; 1fo.write(image_path[i] &＃43; &＃39;_Prediction:&＃39; &＃43; str(class_name[prob]) &＃43; &＃39;\n&＃39;)print(t_num/(t_num &＃43; f_num))