- leNet训练自制数据集并部署在OpenCV上
–0. 简介
–1. 数据集介绍
–2. 数据集读取
–3. 网络搭建
–4. 模型配置与部分超参数设置
–5. 训练结果
–6. 模型转换与部署
–7. 模型训练源码
在leNet的基础上修改网络部分超参数,训练自制数据集并保存模型为pb文件,最后部署在OpenCV的dnn模块上用于目标分类。
这是大疆举办的RoboMaster比赛中视觉组所要识别的一些装甲板。其中带数字的装甲板是比赛中要识别主要的目标,对应的标签为其数字;而一些比赛中可能误识别的目标,需要进行排除,其标签均设为0。故总共6类标签。
我随机挑选了1463张图片作为训练集,559张图片作为测试集。将训练好的模型保存
def makeDataset(path):
dataset_path = path
Images_Path = os.listdir(dataset_path)
Labels = []
Data = []
for image_name in Images_Path:
image_path = dataset_path + '/' + image_name
image = cv2.imread(image_path)
image = cv2.resize(image, (32, 32))
Data.append(image)
if image_name.find('hero') != -1:
Labels.append(1)
elif image_name.find('engineer') != -1:
Labels.append(2)
elif image_name.find('infantry3') != -1:
Labels.append(3)
elif image_name.find('infantry4') != -1:
Labels.append(4)
elif image_name.find('infantry5') != -1:
Labels.append(5)
else:
Labels.append(0)
Data = np.array(Data)
Data = Data / 255.0
Labels = np.array(Labels)
return Data, Labels
module = tf.keras.Sequential([
Input(shape=(32,32,3), dtype=tf.float32, name='Input'),
Conv2D(filters=5, kernel_size=(5,5), padding='valid'),
Activation('sigmoid'),
MaxPooling2D(pool_size=(2,2), strides=2, padding='valid'),
Conv2D(filters=16, kernel_size=(5,5), padding='valid'),
Activation('sigmoid'),
MaxPooling2D(pool_size=(2,2), strides=2, padding='valid'),
Flatten(),
Dense(120, activation='sigmoid'),
Dense(84, activation='sigmoid'),
Dense(6, activation='softmax', name='Output')
])
这里相较于原leNet,对第一层卷积层作出了一些修改,将卷积核数改为5个,输入图的通道改为3通道,同时也对最后一层全连接层的输出数修改为6以对应6种标签
module.compile(optimizer='adam',
loss=SparseCategoricalCrossentropy(from_logits=False),
metrics=['sparse_categorical_accuracy'])
history = module.fit(train_Data, train_Labels, batch_size=32, epochs=30,
validation_data=(test_Data, test_Labels),
validation_freq=1, callbacks=cp_callback)
在一开始的测试中发现训练5轮的效果非常差,正确率不到0.3,但差不多在10轮后,正确率开始陡增。训练100轮的正确率结果几乎为1,loss值也非常小数量级为10^-5,有可能过拟合。故设置训练30轮。
使用了回调函数对模型的ckpt文件进行保存。
网络在训练集上的正确率结果非常接近于1,但是在测试集上的效果一般,正确率约为0.9
通过tf.keras.models.save_model()得到的模型pb文件并不能直接用于OpenCV的dnn上,需要保存为frozen graph格式,因此需要对导出模型进行转换。出处: https://leimao.github.io/blog/Save-Load-Inference-From-TF2-Frozen-Graph/
import tensorflow as tf
from tensorflow import keras
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
import numpy as np
frozen_out_path = './Module'
frozen_graph_filename = 'frozen_graph'
model = tf.keras.models.load_model('./Module/leNet2')
full_model = tf.function(lambda x: model(x))
full_model = full_model.get_concrete_function(
tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype))
frozen_func = convert_variables_to_constants_v2(full_model)
frozen_func.graph.as_graph_def()
layers = [op.name for op in frozen_func.graph.get_operations()]
print("-" * 60)
print("Frozen model layers: ")
for layer in layers:
print(layer)
print("-" * 60)
print("Frozen model inputs: ")
print(frozen_func.inputs)
print("Frozen model outputs: ")
print(frozen_func.outputs)
tf.io.write_graph(graph_or_graph_def=frozen_func.graph,
logdir=frozen_out_path,
name=f"{frozen_graph_filename}.pb",
as_text=False)
tf.io.write_graph(graph_or_graph_def=frozen_func.graph,
logdir=frozen_out_path,
name=f"{frozen_graph_filename}.pbtxt",
as_text=True)
最后通过c++上的OpenCV进行部署:
#include
#include
using namespace std;
using namespace cv;
int main()
{
dnn::Net net = dnn::readNetFromTensorflow("/home/shanzoom/PycharmProjects/pythonProject/Module/frozen_graph.pb");
Mat frame = imread("/home/shanzoom/Robot_DataSet_3.0(balanced)/test2/engineer_ (2906).jpg");
imshow("frame",frame);
frame = dnn::blobFromImage(frame, 1.0/255.0, Size(32,32));
net.setInput(frame);
Mat score = net.forward();
Point maxclass;
minMaxLoc(score, NULL, NULL, NULL, &maxclass);
cout << "装甲板数字: " << maxclass.x << endl;
waitKey(0);
}
import tensorflow as tf
from tensorflow.keras.layers import *
from tensorflow.keras.losses import *
import matplotlib.pyplot as plt
import cv2
import numpy as np
import os
def makeDataset(path):
dataset_path = path
Images_Path = os.listdir(dataset_path)
Labels = []
Data = []
for image_name in Images_Path:
image_path = dataset_path + '/' + image_name
image = cv2.imread(image_path)
image = cv2.resize(image, (32, 32))
Data.append(image)
if image_name.find('hero') != -1:
Labels.append(1)
elif image_name.find('engineer') != -1:
Labels.append(2)
elif image_name.find('infantry3') != -1:
Labels.append(3)
elif image_name.find('infantry4') != -1:
Labels.append(4)
elif image_name.find('infantry5') != -1:
Labels.append(5)
else:
Labels.append(0)
Data = np.array(Data)
Data = Data / 255.0
Labels = np.array(Labels)
return Data, Labels
train_path = '/home/shanzoom/Robot_DataSet_3.0(balanced)/svm_train'
train_Data, train_Labels = makeDataset(train_path)
test_path = '/home/shanzoom/Robot_DataSet_3.0(balanced)/svm_test'
test_Data, test_Labels = makeDataset(test_path)
module = tf.keras.Sequential([
Input(shape=(32,32,3), dtype=tf.float32, name='Input'),
Conv2D(filters=5, kernel_size=(5,5), padding='valid'),
Activation('sigmoid'),
MaxPooling2D(pool_size=(2,2), strides=2, padding='valid'),
Conv2D(filters=16, kernel_size=(5,5), padding='valid'),
Activation('sigmoid'),
MaxPooling2D(pool_size=(2,2), strides=2, padding='valid'),
Flatten(),
Dense(120, activation='sigmoid'),
Dense(84, activation='sigmoid'),
Dense(6, activation='softmax', name='Output')
])
module.compile(optimizer='adam',
loss=SparseCategoricalCrossentropy(from_logits=False),
metrics=['sparse_categorical_accuracy'])
ckpt_path = './checkpoint/leNet.ckpt'
if os.path.exists(ckpt_path + '.index'):
print('----load----')
module.load_weights(ckpt_path)
cp_callback = tf.keras.callbacks.ModelCheckpoint(filepath=ckpt_path,
save_weights_only=True,
save_best_only=True)
history = module.fit(train_Data, train_Labels, batch_size=32, epochs=30, validation_data=(test_Data, test_Labels),
validation_freq=1, callbacks=cp_callback)
module.summary()
tf.keras.models.save_model(module,'./Module/leNet2')
acc = history.history['sparse_categorical_accuracy']
val_acc = history.history['val_sparse_categorical_accuracy']
loss = history.history['loss']
val_loss = history.history['val_loss']
plt.subplot(1, 2, 1)
plt.plot(acc, label='Training Accuracy')
plt.plot(val_acc, label='Validation Accuracy')
plt.title('Training and Validation Accuracy')
plt.legend()
plt.subplot(1, 2, 2)
plt.plot(loss, label='Training Loss')
plt.plot(val_loss, label='Validation Loss')
plt.title('Training and Validation Loss')
plt.show()
Original: https://blog.csdn.net/qq_52967097/article/details/123708406
Author: 啊!山鬃
Title: Tensorflow训练数字识别数据集并部署在OpenCV上
原创文章受到原创版权保护。转载请注明出处:https://www.johngo689.com/692165/
转载文章受原作者版权保护。转载请注明原作者出处!