Update app.py

app.py

@@ -0,0 +1,58 @@
+import gradio as gr
+import onnxruntime as ort
+import numpy as np
+from PIL import Image
+
+# 替换成你的 ONNX 模型文件路径
+ONNX_MODEL_PATH = 'Car_Bike_Classification/end2end.onnx'
+
+# 加载模型
+ort_session = ort.InferenceSession(ONNX_MODEL_PATH)
+
+# 定义预测函数
+def classify_image(image):
+    # 确保image是一个PIL图像对象
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
+    
+    # 预处理输入图片
+    image = image.resize((224, 224))
+    image = np.array(image).astype('float32')
+    image = np.transpose(image, (2, 0, 1))  # Change data layout from HWC to CHW
+    image = np.expand_dims(image, axis=0)
+
+    # 使用 ONNX 运行推理
+    inputs = {ort_session.get_inputs()[0].name: image}
+    outputs = ort_session.run(None, inputs)
+    
+    # 获取预测结果
+    predictions = outputs[0]
+    
+    # 创建一个矩形图展示可能性
+    pred_probs = predictions[0].tolist()  # 假设模型输出是一个概率列表
+    pred_probs = [float(i)/sum(pred_probs) for i in pred_probs]  # 归一化概率值
+    class_names = ['Bike', 'Car']
+    result = {class_name: prob for class_name, prob in zip(class_names, pred_probs)}
+    
+    # 生成结果文本
+    pred_text = f"该图片里是{'汽车' if np.argmax(predictions) else '单车'}。"
+
+    return result, pred_text
+
+# 根据Gradio新的API更新了导入和组件
+from gradio import Interface, components
+
+# 创建 Gradio 界面
+iface = Interface(
+    fn=classify_image,
+    inputs=components.Image(shape=(224, 224)),
+    outputs=[
+        components.Label(num_top_classes=2),
+        components.Textbox(label="分类结果")
+    ],
+    title="单车与汽车图像分类",
+    description="上传一张图片，模型将预测图片是单车还是汽车。"
+)
+
+# 启动界面
+iface.launch()