如何在本地Gradio应用中使用自制TensorFlow CNN模型并解决报错

报错原因分析 #

1. 第一个TypeError：Gradio的gr.Image组件的shape参数仅接受二维（宽、高）数值，不能传入包含通道数的三维元组（比如(28,28,1)），这会导致PIL库在处理图片尺寸时出错。
2. 第二个ValueError：Gradio默认会将输入图片转为RGB三通道格式（形状(28,28,3)），但你的模型要求输入是单通道灰度图（形状(None,28,28,1)，其中None代表batch维度），同时缺少模型需要的batch维度。

解决方案步骤 #

• 配置Gradio的Image组件，指定输入为灰度图，且仅设置宽高维度
• 编写预处理函数，将Gradio传入的图片转换为模型要求的格式：
  • 确保图片是单通道灰度图
  • 增加batch维度（模型接收的输入是(batch_size, 28,28,1)格式）
  • 保持和训练数据一致的数值范围（MNIST数据是0-255，无需额外归一化，若训练时做了归一化则同步处理）

修改后的完整Gradio代码 #

import tensorflow as tf
import gradio as gr
import numpy as np

# 加载模型
model = tf.keras.models.load_model("mnist_model.h5")

def predict_digit(img):
    # 将Gradio传入的图片转为numpy数组，此时形状为(28,28)
    img_array = np.array(img)
    # 增加通道维度和batch维度，转为模型需要的(1,28,28,1)格式
    img_array = img_array[np.newaxis, ..., np.newaxis]
    # 执行预测
    predictions = model.predict(img_array)
    # 整理为Gradio Label组件需要的字典格式
    return {str(i): float(predictions[0][i]) for i in range(10)}

# 创建Gradio接口
app = gr.Interface(
    fn=predict_digit,
    inputs=gr.Image(shape=(28, 28), image_mode="L"),  # "L"指定输入为灰度图
    outputs=gr.Label(num_top_classes=1)
)

app.launch()

补充：训练代码的优化建议 #

你的训练代码中，MNIST加载的train_data形状是(60000,28,28)，而模型输入要求是(28,28,1)，虽然TensorFlow会自动兼容，但建议显式增加通道维度，避免潜在问题：

import tensorflow as tf
from keras.datasets import mnist

(train_data, train_labels), (test_data, test_labels) = mnist.load_data()

# 显式增加通道维度
train_data = train_data[..., tf.newaxis]
test_data = test_data[..., tf.newaxis]

model_2 = tf.keras.Sequential([
    tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
    tf.keras.layers.MaxPooling2D((2, 2)),
    tf.keras.layers.Conv2D(64, (3, 3), activation='relu'),
    tf.keras.layers.MaxPooling2D((2, 2)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Dense(10, activation='softmax')
])

model_2.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# 训练模型
model_2.fit(train_data, train_labels, epochs=15)

# 评估模型
model_2.evaluate(test_data, test_labels)

# 保存模型
model_2.save("mnist_model.h5")

内容的提问来源于stack exchange，提问作者Darkkindess