基于 TensorFlow2.7（Keras） + Flask 的 Mnist 手写数字集识别

基于 TensorFlow2.7（Keras） + Flask 的 Mnist 手写数字集识别

首先展示我的flask服务：

import os
import re
import base64
import socket

import numpy as np
import tensorflow.keras as keras
from flask import Flask, jsonify, request

from tensorflow.keras.preprocessing.image import img_to_array, load_img

import tensorflow as tf

app = Flask(__name__)
# 配置 TensorFlow 运行时
# os.environ["CUDA_VISIBLE_DEVICES"] = "0"
# gpus = tf.config.list_physical_devices('GPU')
# if gpus:
#     try:
#         # 设置仅使用第一个 GPU 设备
#         tf.config.experimental.set_visible_devices(gpus[0], 'GPU')
#         logical_gpus = tf.config.experimental.list_logical_devices('GPU')
#         print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
#     except RuntimeError as e:
#         # 异常处理
#         print(e)
print(tf.test.is_gpu_available())
print(tf.test.is_built_with_cuda())
# tf.debugging.set_log_device_placement(True)

def host_ip():
    """
    查询本机ip地址
    :return: ip
    """
    ip = '0.0.0.0'
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        s.connect(('8.8.8.8', 80))
        ip = s.getsockname()[0]
    except OSError as ex:
        hostname = socket.gethostname()
        ip = socket.gethostbyname(hostname)
    finally:
        s.close()

    return ip



model_file = './model/model.h5'
global model
model = keras.models.load_model(model_file)


@app.route('/predict_mnist', methods=["POST"])
def predict():

    data = request.form['file_base64']
    file_data = base64.b64decode(data)
    with open('./image/cache.png', 'wb') as f:
        f.write(file_data)
    # parseImage(request.data)
    # # parseImage(request.get_data(as_text=True))
    img = img_to_array(load_img('./image/cache.png', target_size=(28, 28), color_mode="grayscale")) / 255.
    img = np.expand_dims(img, axis=0)
    code = model.predict_classes(img)[0]
    # # 使用模型预测图像
    # preds = model.predict(img)
    #
    # # 找到最高概率对应的类别
    # code = np.argmax(preds, axis=1)[0]

    response = {'my_code': int(code)}
    return jsonify(response)


def parseImage(imgData):
    imgStr = re.search('base64,(.*)', imgData).group(1)
    print(imgData)
    with open('./output.png', 'wb') as output:
        output.write(base64.decodebytes(imgStr.encode()))



if __name__ == '__main__':
    hostip = host_ip()
    print(hostip)
    app.run(port=5553, host=hostip)

从flask服务中去解析这个程序的运行：

重点代码端的注释：：：：
print(tf.test.is_gpu_available())  # 打印是否存在可用的GPU
print(tf.test.is_built_with_cuda())  # 打印是否使用CUDA进行构建

model_file = './model/model.h5'  # 模型文件路径
global model  # 定义全局变量model
model = keras.models.load_model(model_file)  # 加载模型文件



    img = img_to_array(load_img('./image/cache.png', target_size=(28, 28), color_mode="grayscale")) / 255.  # 处理图像数据
    img = np.expand_dims(img, axis=0)  # 在第0维度上扩展数组形状
    code = model.predict_classes(img)[0]  # 使用模型预测图像并获取结果
    response = {'my_code': int(code)}  # 构建响应数据

源码使用的是训练好的模型 model.h5,测试集成功率95以上，如果自己训练的话：python model/train.py即可。