本文介绍如何使用ResNet网络进行图像分类，并通过OpenCV调用.pb模型文件进行预测。主要内容包括网络搭建、模型训练、pb文件生成及Python与C++调用测试。
原理这边就不多说了，主要流程就是：
1，搭建网络；
2，训练网络模型，生成ckpt文件
3，保存graph，并将训练好的参数固化到graph中，生成.pb文件
4，OpenCV调用.pb文件预测
需要注意的几点：
1，搭建网络时注意网络节点的命名，后续会用到；
2，TensorFlow有的算子OpenCV不支持，如果搭建网络时用到了OpenCV不支持的算子，进行预测时，会报错（没有找到OpenCV支持/不支持的算子文档，找到的同学可以留言告诉我噢）；
3，一定要有freeze这一步骤，如果直接保存，同样的输入，每次预测的结果都不一样，一开始我是参考网上的资料训练完成之后直接用graph_util.convert_variables_to_constants保存，调用时发现每次结果都随机；
4，生成.pb文件后，在用OpenCV调用前可以先Python调用预测一下保证pb模型文件的正确性
好了，不多说，直接上代码
一. ResNet网络搭建
def Block(X_input,kernel_size,in_filter,out_filters,stride,choice):
X_sortcut = X_input
f1, f2, f3 = out_filters
conv_w1 = tf.get_variable('weight1', [1, 1, in_filter, f1], initializer=tf.truncated_normal_initializer(stddev=0.1))
conv_b1 = tf.get_variable('bias1', [f1],initializer=tf.constant_initializer(0.0))
conv_w2 = tf.get_variable('weight2', [kernel_size, kernel_size, f1, f2], initializer=tf.truncated_normal_initializer(stddev=0.1))
conv_b2 = tf.get_variable('bias2', [f2], initializer=tf.constant_initializer(0.0))
conv_w3 = tf.get_variable('weight3', [1, 1, f2, f3], initializer=tf.truncated_normal_initializer(stddev=0.1))
conv_b3 = tf.get_variable('bias3', [f3], initializer=tf.constant_initializer(0.0))
if choice:
# frist
y = tf.nn.relu(tf.nn.conv2d(X_input, conv_w1, strides=[1, 1, 1, 1], padding="SAME") + conv_b1)
# second
y = tf.nn.relu(tf.nn.conv2d(y, conv_w2, strides=[1, 1, 1, 1], padding="SAME") + conv_b2)
# third
y = tf.nn.relu(tf.nn.conv2d(y, conv_w3, strides=[1, 1, 1, 1], padding="SAME") + conv_b3)
add_result = tf.nn.relu(y + X_sortcut)
return add_result
else:
# frist
y = tf.nn.relu(
tf.nn.conv2d(X_input, conv_w1, strides=[1, stride, stride, 1], padding="SAME") + conv_b1)
# second
y = tf.nn.relu(tf.nn.conv2d(y, conv_w2, strides=[1, 1, 1, 1], padding="SAME") + conv_b2)
# third
y = tf.nn.relu(tf.nn.conv2d(y, conv_w3, strides=[1, 1, 1, 1], padding="SAME") + conv_b3)
# final steap
add = tf.nn.conv2d(X_sortcut, conv_w3, strides=[1, 1, 1, 1], padding="SAME")
add_result = tf.nn.relu(y + add)
return add_result
def ResNet_18(input_tensor):
with tf.variable_scope('conv1'):
conv1_weights = tf.get_variable('weight', [3, 3, 1, 16],
initializer=tf.truncated_normal_initializer(stddev=0.1))
conv1_biases = tf.get_variable('bias', [16],
initializer=tf.constant_initializer(0.0))
conv1 = tf.nn.conv2d(input_tensor, conv1_weights, strides=[1, 1, 1, 1], padding='SAME')
relu1 = tf.nn.relu(tf.nn.bias_add(conv1, conv1_biases))
pool1 = tf.nn.max_pool(relu1, ksize=[1,2,2,1],strides=[1,2,2,1],padding='SAME')
with tf.variable_scope('block1'):
Rt1 = Block(pool1, 3, 16, [8,8,16],1,True)
with tf.variable_scope('block2'):
Rt2 = Block(Rt1, 3, 16, [8,8,16],1,True)
with tf.variable_scope('block3'):
Rt3 = Block(Rt2, 3, 16, [8,8,16],1,True)
with tf.variable_scope('block4'):
Rt4 = Block(Rt3, 3, 16, [8,8,16],1,True)
with tf.variable_scope('block5'):
Rt5 = Block(Rt4, 3, 16, [16,16,32],1,False)
with tf.variable_scope('Flat'):
pool2 = tf.nn.max_pool(Rt5, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding="SAME")
flat = tf.reshape(pool2, [-1, 7 * 7 * 32])
# ---------------------------------全链接层------------------------------------
with tf.variable_scope('Dense'):
W = tf.Variable(tf.random_normal([7 * 7 * 32, 128], dtype=tf.float32, stddev=0.1))
B = tf.Variable(tf.zeros([128]))
W1 = tf.Variable(tf.random_normal([128, 10], dtype=tf.float32, stddev=0.1))
B1 = tf.Variable(tf.zeros([10]))
y0 = tf.nn.relu(tf.matmul(flat, W) + B)
logit = tf.nn.softmax(tf.matmul(y0, W1) + B1, name='soft')
return logit
二. 训练模型，保存pb文件
def train(mnist):
x = tf.placeholder(tf.float32, [None, 28, 28, 1], name='x-input')
y_ = tf.placeholder(tf.float32, [None, 10], name='y-input')
y = ResNet_18(x)
global_step = tf.Variable(0, trainable=False)
variable_average = tf.train.ExponentialMovingAverage(MOVING_AVERAGE_DECAY, global_step)
variable_average_op = variable_average.apply(
tf.trainable_variables())
cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=tf.argmax(y_, 1), logits=y)
cross_entropy_mean = tf.reduce_mean(cross_entropy)
loss = cross_entropy_mean
learning_rate = tf.train.exponential_decay(LEARNING_RATE_BASE,
global_step=global_step,
decay_steps=mnist.train.num_examples / BATCH_SIZE,
decay_rate=LEARNING_RATE_DECAY)
train_step = tf.train.GradientDescentOptimizer(learning_rate).minimize(loss, global_step=global_step)
with tf.control_dependencies([train_step, variable_average_op]):
train_op = tf.no_op(name='train')
saver = tf.train.Saver()
with tf.Session() as sess:
tf.train.write_graph(sess.graph_def, "./model", 'test_graph.pb')
tf.global_variables_initializer().run()
for i in range(TRAIN_STEPS):
xs, ys = mnist.train.next_batch(BATCH_SIZE)
xs = np.reshape(xs, [-1, 28,28,1])
_, loss_value, step = sess.run([train_op, loss, global_step], feed_dict={x: xs, y_: ys})
if i % 1000 == 0:
print("After %d training steps, loss on training"
"batch is %g" % (step, loss_value))
saver.save(sess, os.path.join(MODEL_SAVE_PATH, MODEL_NAME), global_step=global_step)
ckpt = tf.train.get_checkpoint_state(MODEL_SAVE_PATH)
if ckpt and ckpt.model_checkpoint_path:
saver.restore(sess, ckpt.model_checkpoint_path)
tf.train.write_graph(sess.graph_def, './model', 'model.pb')
freeze_graph.freeze_graph('./model/model.pb', '', False, ckpt.model_checkpoint_path, 'Dense/soft',
'', '', './model/frozen_model.pb',
False, "")
三. 调用pb模型预测
Python调用测试代码：
import tensorflow as tf
from tensorflow.python.platform import gfile
import numpy as np
from PIL import Image
def main():
with tf.Session() as sess:
with gfile.FastGFile("model/frozen_model.pb", 'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
sess.graph.as_default()
tf.import_graph_def(graph_def, name='')
sess.run(tf.global_variables_initializer())
input_x = sess.graph.get_tensor_by_name('x-input:0')
op = sess.graph.get_tensor_by_name('Dense/soft:0')
num = tf.argmax(op, 1)
img = Image.open('3.bmp')
x_img_array = np.float32(np.asarray(img)) / 255.0
reshape_img = np.reshape(x_img_array, [-1, 28, 28, 1])
validate_feed = {input_x: reshape_img}
y_value,result = sess.run([op,num], feed_dict=validate_feed)
print(result[0])
main()
C++ OpenCV调用测试代码：
#include "stdafx.h"
#include <opencv2/dnn.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>
#include <Windows.h>
#include <fstream>
#include <iostream>
#include <cstdlib>
using namespace cv;
using namespace cv::dnn;
using namespace std;
//自己新建一个txt文件，写入分类的标签（一行写一个标签，例如二分类，第一行写good，第二行bad）
String labels_txt_file = "G:\\CODE\\Lesson\\lesson_leNet5\\model\\label.txt";
String tf_pb_file = "G:\\CODE\\Lesson\\lesson_ResNet\\model\\frozen_model.pb";
//String tf_pbtxt_file = "G:\\CODE\\Lesson\\lesson_leNet5\\model\\protobuf.pbtxt";
vector <String> readClassNames();
void main()
{
Mat src = imread("D:\\2.bmp");
if (src.empty())
{
cout << "error:no img" << endl;
}
vector <String> labels = readClassNames();
Mat rgb;
int w = 28;
int h = 28;
resize(src, src, Size(w, h));
cvtColor(src, rgb, COLOR_BGR2GRAY);
int n = src.channels();
Net net = readNetFromTensorflow(tf_pb_file);
DWORD timestart = GetTickCount();
if (net.empty())
{
cout << "error:no model" << endl;
}
Mat inputBlob = blobFromImage(rgb, 0.003921, Size(w, h), Scalar(), false, false);
//执行图像分类
Mat prob;
net.setInput(inputBlob, "x-input");
prob = net.forward("Dense/soft");
cout << prob << endl;
//得到最大分类概率
Mat probMat = prob.reshape(1, 1);
Point classNumber;
double classProb;
minMaxLoc(probMat, NULL, &classProb, NULL, &classNumber);
DWORD timeend = GetTickCount();
int classidx = classNumber.x;
printf("\n current image classification : %s, possible : %.2f\n", labels.at(classidx).c_str(), classProb);
cout << "用时(毫秒):" << timeend - timestart << endl;
// 显示文本
putText(src, labels.at(classidx), Point(20, 20), FONT_HERSHEY_SIMPLEX, 1.0, Scalar(0, 0, 255), 2, 8);
imshow("Image Classfication", src);
waitKey(0);
}
vector <String>readClassNames()
{
vector <String>classNames;
fstream fp(labels_txt_file);
if (!fp.is_open())
{
cout << "does not open" << endl;
exit(-1);
}
string name;
while (!fp.eof())
{
getline(fp, name);
if (name.length())
classNames.push_back(name);
}
fp.close();
return classNames;
}
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
