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

leNet训练自制数据集并部署在OpenCV上

• leNet训练自制数据集并部署在OpenCV上
  - 0. 简介
  - 1. 数据集介绍
  - 2. 数据集读取
  - 3. 网络搭建
  - 4. 模型配置与部分超参数设置
  - 5. 训练结果
  - 6. 模型转换与部署
  - 7. 模型训练源码

0. 简介

在leNet的基础上修改网络部分超参数，训练自制数据集并保存模型为pb文件，最后部署在OpenCV的dnn模块上用于目标分类。

1. 数据集介绍

这是大疆举办的RoboMaster比赛中视觉组所要识别的一些装甲板。其中带数字的装甲板是比赛中要识别主要的目标，对应的标签为其数字；而一些比赛中可能误识别的目标，需要进行排除，其标签均设为0。故总共6类标签。

我随机挑选了1463张图片作为训练集，559张图片作为测试集。将训练好的模型保存

2. 数据集读取

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))     # leNet 的输入尺寸为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

3. 网络搭建

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种标签

4. 模型配置与部分超参数设置

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文件进行保存。

5. 训练结果

网络在训练集上的正确率结果非常接近于1，但是在测试集上的效果一般，正确率约为0.9

6. 模型转换与部署

通过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))

# Get frozen graph def
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<opencv2/opencv.hpp>
#include<opencv2/dnn.hpp>

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);
}

7. 模型训练源码

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()