Opencv调用tensorflow训练自己的数据集生成的模型
要实现opencv调用tensorflow训练的模型,主要分为两步,第一步是训练模型,将模型保存成model.pb格式,然后利用opencv的readNetFromTensorflow方法调用model.pb
一、训练生成模型:
训练生成模型的关键点:
1.将图片生成自己的数据集
2.将标签转换成独热编码
3.利用cnn训练模型,要获取输出的节点名称,最后是自己命名。
4.保存成pb模型,在这里会用到输出的节点名称,这一步是关键,否则opencv调用时会报错。
二、Opencv调用
这一步比较简单,直接调用即可,下面附代码,具体内容会写在注释中
图片的类型:如图是32x32的图像,像这样命名,标签在前,方便读取图像时获取标签
Tensorflow训练的代码:具体见注释
import tensorflow as tf
from tensorflow.python.framework import graph_util
import os
from PIL import Image
import numpy as np
#将标签进行独热编码
def getOneHot(labels, num):
a = []
for i in labels:
k = np.zeros(num)
k[i] = 1
a.append(k)
return np.array(a,dtype=np.float32)
#读取数据,返回datas为numpy的array格式,类型一定要设置为float32
def 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 = np.array(image, dtype=np.float32) / 255.0
label = 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, labels
data_dir = './data'
fpaths, datas, labels = read_data(data_dir)
labels = getOneHot(labels, 3)
#生成权重
def weight_variable(shape):
return tf.Variable(tf.truncated_normal(shape, stddev=0.1))
#生成偏置
def bias_variable(shape):
return tf.Variable(tf.constant(0.1, shape=shape))
#卷积
def conv2d(x, W):
return tf.nn.conv2d(x, W, strides=[1,1,1,1], padding='SAME')
#池化
def max_pool_2x2(x):
return tf.nn.max_pool(x, ksize=[1,2,2,1], strides=[1,2,2,1], padding='SAME')
#这里的32,32,3是图像的维度,None代表输入的图像数量可以是任意的
datas_placeholder = tf.placeholder(tf.float32, [None, 32, 32, 3])
#这里的3是类别,None的数目和输入的图像数目是相对应的
y_ = tf.placeholder(tf.float32, [None, 3])
#第一层卷积
#将过滤器设置成5×5×1的矩阵,
#其中5×5表示过滤器大小,3表示深度。
#32表示我们要创建32个大小5×5×1的过滤器,经过卷积后算出32个特征图(每个过滤器得到一个特征图),即输出深度为64
W_conv1 = weight_variable([5, 5, 3, 32])
#有多少个特征图就有多少个偏置
b_conv1 = bias_variable([32])
#使用conv2d函数进行卷积计算,然后再用ReLU作为激活函数
h_conv1 = tf.nn.relu(conv2d(datas_placeholder, W_conv1) + b_conv1)
# #卷积以后再经过池化操作
h_pool1 = max_pool_2x2(h_conv1)
#第二层卷积
#因为经过第一层卷积运算后,输出的深度为32,所以过滤器深度和下一层输出深度也做出改变
W_conv2 = weight_variable([5, 5, 32, 64])
b_conv2 = bias_variable([64])
h_conv2 = tf.nn.relu(conv2d(h_pool1, W_conv2) + b_conv2)
h_pool2 = max_pool_2x2(h_conv2)
print(h_pool2)
#第三层卷积
W_conv3 = weight_variable([3, 3, 64, 128])
b_conv3 = bias_variable([128])
h_conv3 = tf.nn.relu(conv2d(h_pool2, W_conv3) + b_conv3)
h_pool3 = max_pool_2x2(h_conv3)
print(h_pool3)
#全连接层
#经过两层卷积后,图片的大小为4×4(第一层池化后输出为(32/2)×(32/2),
#第二层池化后输出为(16/2)×(16/2)),深度为64,
#第三层输出后为(8/2)×(8/2)),深度为128,
#我们在这里加入一个有1024个神经元的全连接层,所以权重W的尺寸为[4 * 4 * 128, 1024]
W_fc1 = weight_variable([4 * 4 * 128, 1024])
#偏置的个数和权重的个数一致
b_fc1 = bias_variable([1024])
#这里将第二层池化后的张量(长:4 宽:4 深度:64) 变成向量(跟上一节的Softmax模型的输入一样了)
h_pool2_flat = tf.reshape(h_pool3, [-1, 4 * 4 * 128])
#使用ReLU激活函数
h_fc1 = tf.nn.relu(tf.matmul(h_pool2_flat, W_fc1) + b_fc1)
#输出层
#全连接层输入的大小为1024,而我们要得到的结果是三类(0~2),
#所以这里权重W的尺寸为[1024, 3]
W_fc2 = weight_variable([1024, 3])
b_fc2 = bias_variable([3])
#最后都要经过Softmax函数将输出转化为概率问题
y_mul = tf.matmul(h_fc1, W_fc2) + b_fc2
#保存为pb格式要用到输出的节点名称,因此最好自己设置
y_conv = tf.nn.softmax(y_mul, name='output')
#损失函数和损失优化
cross_entropy = tf.reduce_mean(-tf.reduce_sum(y_ * tf.log(y_conv)))
train_step = tf.train.AdamOptimizer(1e-4).minimize(cross_entropy)
#测试准确率,跟Softmax回归模型的一样
correct_prediction = tf.equal(tf.argmax(y_conv, 1), tf.argmax(y_, 1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
#将训练结果保存,如果不保存我们这次训练结束后的结果也随着程序运行结束而释放了
saveFile = 'mnist_model.ckpt'
saver = tf.train.Saver()
with tf.Session() as sess:
#初始化所有变量
sess.run(tf.global_variables_initializer())
for i in range(150):
if i % 10 == 0:
train_accuracy =sess.run(accuracy, feed_dict={datas_placeholder:datas, y_:labels})
print("step %d, training accuracy %g" % (i, train_accuracy))
sess.run(train_step, feed_dict={datas_placeholder:datas, y_:labels})
# 最后,将会话保存下来
saver.save(sess, saveFile)
#保存为pb文件
constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph_def, ['output'])
with tf.io.gfile.GFile('mnist.pb', mode='wb') as f:
f.write(constant_graph.SerializeToString())
Opencv调用代码:
#include <opencv.hpp>
#include <dnn.hpp>
#include <stdlib.h>
#include <iostream>
#include <string.h>
using namespace cv;
using namespace std;
using namespace dnn;
int main()
{
//读取图片
Mat img = imread("E:\\python\\deeplearn\\tensorflow\\AlexNetTest\\img\\1.jpg");
//pb文件路径
string pbfile = "E:\\python\\deeplearn\\tensorflow\\AlexNetTest\\mnist.pb";
//(3)读取Tensorflow模型文件
Net net = readNetFromTensorflow(pbfile);
//(4)创建输入数据
Mat inputBlob = blobFromImage(img, 1.0, Size(32, 32), Scalar(), false, false);
//(5)输入网络
net.setInput(inputBlob);
//(6)预测结果
Mat pred = net.forward();
//(7)输出结果
Point maxLoc;
minMaxLoc(pred, NULL, NULL, NULL, &maxLoc);
cout << maxLoc.x << endl;
//(8)显示图片
imshow("img", img);
waitKey(0);
}