Add gradio ui

gradio_ui.py

@@ -0,0 +1,71 @@
+import gradio as gr
+import numpy as np
+from PIL import Image
+import torch
+import test  # 假设 test 模块包含预测逻辑
+
+# 加载模型 (与 Qt 版本保持一致)
+model = test.load_trained_model()
+
+
+def predict_interface(sketch_image):
+    """处理绘制图像的预测逻辑"""
+    if sketch_image is None:
+        return "请先绘制数字", {}
+
+    # 将 sketchpad 的 numpy 数组转换为模型需要的格式
+    img = Image.fromarray(sketch_image).convert('L')  # 转换为灰度图
+
+    # 可能需要添加预处理步骤（根据 test.predict_user_image 的接口调整）
+    # 如果用原始 Qt 的预处理逻辑，这里可以复用 test 模块的函数
+    pred_class, probabilities = test.predict_user_image(img, model)
+
+    # 转换概率为字典供 Label 组件显示
+    prob_dict = {str(i): float(prob) for i, prob in enumerate(probabilities)}
+    return f"识别结果: {pred_class}", prob_dict
+
+
+def clear_canvas():
+    """清空画布的函数"""
+    return None, "识别结果: ", {}
+
+
+# 构建 Gradio 界面
+with gr.Blocks(title="手写数字识别") as demo:
+    gr.Markdown("# 手写数字识别系统")
+
+    with gr.Row():
+        # 手写板组件 (调整尺寸匹配原 Qt 设计)
+        sketch = gr.Sketchpad(
+            label="绘制区域",
+            shape=(750, 750),
+            brush_radius=15,  # 根据原 Qt 的笔刷大小调整
+            image_mode="L",  # 灰度模式
+            invert_colors=True  # 反转颜色（白底黑字）
+        )
+
+        # 结果显示区域
+        with gr.Column():
+            result_label = gr.Label(label="概率分布", num_top_classes=5)
+            output_text = gr.Markdown("识别结果: ")
+
+    # 按钮行
+    with gr.Row():
+        clear_btn = gr.Button("清除", variant="secondary")
+        submit_btn = gr.Button("识别", variant="primary")
+
+    # 绑定交互事件
+    submit_btn.click(
+        fn=predict_interface,
+        inputs=sketch,
+        outputs=[output_text, result_label]
+    )
+
+    clear_btn.click(
+        fn=lambda: [None, "识别结果: ", None],  # 清空所有输出
+        outputs=[sketch, output_text, result_label]
+    )
+
+# 启动应用
+if __name__ == "__main__":
+    demo.launch()