Skip to content

3.2.2 手写数字识别系统交互流程设计

1.导入库

python
import onnxruntime
import numpy as np
from PIL import Image
  • onnxruntime : 用于加载和运⾏ ONNX 模型
  • numpy : 用于数值计算和数组操作
  • PIL.Image : Python 图像处理库

2.加载 ONNX 模型

python
ort_session = onnxruntime.InferenceSession('mnist.onnx')
  • 创建 ONNX 运⾏时会话
  • 加载名为 'mnist.onnx' 的预训练模型⽂件
  • 返回的 ort_session 对象将⽤于后续的推理

3.加载测试图像

python
image = Image.open('img_test.png').convert('L')
  • 打开名为 'img_test.png' 的图像⽂件
  • .convert('L') 将图像转换为灰度图(单通道)
  • MNIST 数据集是黑白手写数字,所以需要灰度图

4.图像预处理

python
image = image.resize((28, 28))
  • 调整图像⼤⼩为 28×28 像素
  • MNIST 模型的输入要求就是这个尺⼨
python
image_array = np.array(image, dtype=np.float32)
  • 将 PIL 图像转换为 numpy 数组
  • 指定数据类型为 float32 (模型要求的输入类型)
python
image_array = np.expand_dims(image_array, axis=0)
  • 在第⼀个维度添加 batch 维度
  • 形状从 (28,28) 变为 (1,28,28)
  • 因为模型预期输入是批量形式的
image_array = np.expand_dims(image_array, axis=0)
  • 再次添加通道维度
  • 形状从 (1,28,28) 变为 (1,1,28,28)
  • 虽然 MNIST 是单通道,但模型可能期望 4D 输入 (batch,channel,height,width)

5.准备模型输入

python
ort_inputs = {ort_session.get_inputs()[0].name: image_array}
  • 获取模型输入节点的名称
  • 创建输入字典: { 输入节点名称 : 处理后的图像数据 }
  • ONNX Runtime 需要以这种格式接收输入

6.执行预测

python
ort_outs = ort_session.run(None, ort_inputs)
  • 运⾏模型推理
  • 第⼀个参数 None 表⽰我们希望获取所有输出
  • 第⼆个参数是输入数据字典
  • 返回包含所有输出的列表

7.获取预测结果

python
predicted_class = np.argmax(ort_outs[0])
  • ort_outs[0] 获取第⼀个输出(通常是分类 logits )
  • np.argmax() 找到概率最⼤的类别索引
  • MNIST 输出是 0-9 的数字概率,取最⼤值即为预测数字

8.输出预测结果

python
print(f"Predicted class: {predicted_class}")
  • 打印预测结果
  • 输出格式⽰例: "Predicted class: 8"

人机交互的最优方式

(1)用户界面设计:提供一个简洁直观的用户界面,方便用户上传图片和查看识别结果。

(2)加载模型“mnist.onnx”;

(3)加载一张本地手写数字图片“img_test.png”,并进行预处理图像以符合模型输入要求;

(4)使用mnist模型对手写数字图片进行识别;

(5)输出识别后的数字。