mnist数据集利用tensorflow(keras)模型训练,预测(内置数据集),保存,以及.h5转换为.pb再利用opencv-dnn模块的加载和预测(自己手写拍照后的本地数据集)

思路:

思路:

1.先对数据集加载,并进行预处理。

2.建立cnn模型(tensorflow)

3.对模型进行训练,训练次数epoch和学习率learning rate根据最终预测结果可进行参数调整

4.进行预测,这里先使用mnist取样后的测试集进行预测

5.模型保存,再利用opencv-dnn模块只能读取.pb格式,并在本地数据集上预测(注意opencv有python和c++两个版本,如果在c++部署平台上,自行下载c++版opencv进行加载,本文使用opencv-python)

环境:

python3.7  tensorflow-cpu==1.14.0  opencv-python=4.5.3(环境配置可自行下载)

代码如下:(我将代码分为还几个.py,所以再导入模块时要自己注意路径)

1.数据处理(cnn_mnist.py)

from tensorflow import keras
import os
import matplotlib.pyplot as plt
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
def get_data():
    mnist = keras.datasets.mnist
    class_names = ['0','1','2','3','4','5','5','7','8','9']
    (train_images, train_labels), (test_images, test_labels) = mnist.load_data()

    plt.figure(figsize=(10,10)) #显示前25张图像
    for i in range(25):
        plt.subplot(5,5,i+1)
        plt.grid(False)
        plt.imshow(train_images[i],cmap=plt.cm.binary)

    plt.show()

    print(train_images.shape,train_labels.shape,)

    # 将图像的数据类型转换成浮点型,再将像素值缩小到0-1,完成数据的预处理
    train_images = train_images.reshape([-1,28,28,1]) / 255.0
    test_images = test_images.reshape([-1,28,28,1]) / 255.0
    # print(test_images)
    return train_images,test_images,train_labels,test_labels,class_names
get_data()

 2.建立模型(model.py)

import tensorflow as tf
from tensorflow import keras

def get_model():
    M=keras.Sequential([
    #(-1,28,28,1)->(-1,28,28,32)
    keras.layers.Conv2D(input_shape=(28, 28, 1),filters=32,kernel_size=5,strides=1,padding='same'),     # Padding method),
    #(-1,28,28,32)->(-1,14,14,32)
    keras.layers.MaxPool2D(pool_size=2,strides=2,padding='same'),
    #(-1,14,14,32)->(-1,14,14,64)
    keras.layers.Conv2D(filters=64,kernel_size=3,strides=1,padding='same'),
    #(-1,14,14,64)->(-1,7,7,64)
    keras.layers.MaxPool2D(pool_size=2,strides=2,padding='same'),
    #(-1,7,7,64)->(-1,7*7*64)
    keras.layers.Flatten(),
    #(-1,7*7*64)->(-1,256)
    keras.layers.Dense(256, activation=tf.nn.relu),
    #(-1,256)->(-1,10)
    keras.layers.Dense(10, activation=tf.nn.softmax)
    ])
    return M
print(get_model().summary())

3.模型训练和mnist取样后的数据进行预测(train.py)[注意:这里训练模型的时候使用的keras,则保存为.h5的格式,不同同框架和方法保存的模型格式不一样,参考https://blog.csdn.net/sunflower_sara/article/details/99364051]

from model import get_model
import cnn_mnsit
import tensorflow as tf
import matplotlib.pyplot as plt
import numpy as np
def go_train(model):
    # pb_file_path = os.getcwd()
    train_images,test_images,train_labels,test_labels,class_names=cnn_mnsit.get_data()
    lr = 0.005
    epochs = 10
    # 编译模型
    model.compile(optimizer=tf.train.AdamOptimizer(lr),
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])

    # 拟合数据
    model.fit(train_images, train_labels, epochs=epochs,shuffle=True,verbose=2, validation_split=0.3)
    # 模型评测
    test_loss, test_acc =model.evaluate(test_images, test_labels)

    print('the model\'s test_loss is {} and test_acc is {}'.format(test_loss, test_acc))

    # 部分预测结果展示
    show_images = test_images[10:100:10]
    print(show_images.shape)
    predictions = model.predict(show_images)
    predict_labels = np.argmax(predictions, 1)

    plt.figure(figsize=(10, 5))  # 显示前10张图像,并在图像上显示类别
    for i in range(9):
        plt.subplot(2, 5, i + 1)
        plt.grid(False)
        plt.imshow(show_images[i, :, :, 0], cmap=plt.cm.binary)
        plt.title(class_names[predict_labels[i]])

    plt.show()
    model.save('model.h5')
    tf.saved_model.save(model, '/home/tjk/project/tf_doc/model_save/fashion_10/')
if __name__ == '__main__':
    go_train(get_model())

 4.模型预测本地数据集(提示:这里我的数据集是放在一个文件夹number里面的,然后图片的名称和图片的数字一样,你也可以输入一张本地图片进行检测,改一下路径就可以) 

import tensorflow as tf
from tensorflow import keras
import cv2
import numpy as np
import os
model = tf.keras.models.load_model("model.h5")
pred_x=os.listdir('number')
print(pred_x)
result_sum=[]
result_one={}
for img in pred_x:
    # print(img)
    pred_img = cv2.imread('number/'+img)
    # print(pred_img)
    pred_img = cv2.resize(pred_img,(28,28))
    pred_img = cv2.cvtColor(pred_img, cv2.COLOR_BGR2GRAY)
    # cv2.imshow('0',pred_img)
    pred_img = np.array(pred_img)
    pred_img = 1-pred_img.reshape(-1,28,28,1)
    predict = model.predict(pred_img)
    Final_prediction = [result.argmax() for result in predict][0]
    result_one={img:Final_prediction}
    result_sum.append(result_one)
    # for i in predict[0]:
    #     print('Percent:{:.10f}'.format(i))

print(result_sum)

5.h5转化为.pb格式(这里使用是tensorflow1.x  ,如果使用tensorflow2.x需要使用2.x的tf.function功能并转化为需要保存为frozen graph格式,参考需要保存为frozen graph格式

再利用opencv-dnn模块加载.pb并进行本地图片预测(会生成一个.pb文件,读取该文件后进行预测即可)

import tensorflow.keras as keras
import tensorflow as tf
import os
import cv2
import numpy as np
# 这个函数参考自网上
def freeze_session(session, keep_var_names=None, output_names=None, clear_devices=True):

    graph = session.graph
    with graph.as_default():
        freeze_var_names = list(set(v.op.name for v in tf.global_variables()).difference(keep_var_names or []))
        output_names = output_names or []
        output_names += [v.op.name for v in tf.global_variables()]
        input_graph_def = graph.as_graph_def()
        if clear_devices:
            for node in input_graph_def.node:
                node.device = ''
        frozen_graph = tf.graph_util.convert_variables_to_constants(
            session, input_graph_def, output_names, freeze_var_names)
        return frozen_graph


if __name__ == '__main__':
    input_path = "D:\\teach\\Handwriting\\"
    # keras训练保存的h5文件
    input_file = 'model.h5'
    weight_file_path = os.path.join(input_path, input_file)
    output_graph_name = weight_file_path[:-3] + '.pb'

    #  加载模型
    keras.backend.set_learning_phase(0)
    h5_model = keras.models.load_model(weight_file_path)
    frozen_graph = freeze_session(keras.backend.get_session(), output_names=[out.op.name for out in h5_model.outputs])
    tf.train.write_graph(frozen_graph, input_path, output_graph_name, as_text=False)

    print('Finished')


    model = cv2.dnn.readNetFromTensorflow("model.pb")
    print('Load')
    pred_x = os.listdir('number')
    print(pred_x)
    result_sum = []
    result_one = {}
    for img in pred_x:
        # print(img)
        pred_img = cv2.imread('number/' + img)
        # print(pred_img)
        pred_img = cv2.resize(pred_img, (28, 28))
        pred_img = cv2.cvtColor(pred_img, cv2.COLOR_BGR2GRAY)
        # cv2.imshow('0',pred_img)
        pred_img = np.array(pred_img)
        pred_img = pred_img.reshape(-1, 28, 28, 1)
        model.setInput(pred_img)
        out = model.forward()
        out = out.flatten()
        Final_prediction = np.argmax(out)
        result_one = {img: Final_prediction}
        result_sum.append(result_one)
print(result_sum)
总结:1.代码有点冗余,代码能力强的可以自行修改,大概思路是这样的     2.菜鸟一位,第一次写博客,很多不足,希望多和大家交流。

 

上一篇:EF的双表联查、分页、添加


下一篇:更新Appium中的WebDriverAgent