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左右。

[En]

The correct rate of the network on the training set is very close to 1, but the effect on the test set is general, the correct rate is about 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/508847/

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

(0)

大家都在看

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