Tensorflow训练数字识别数据集并部署在OpenCV上

在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/

转载文章受原作者版权保护。转载请注明原作者出处!

(0)

大家都在看

亲爱的 Coder【最近整理,可免费获取】👉 最新必读书单  | 👏 面试题下载  | 🌎 免费的AI知识星球