a

app.py

@@ -1,25 +1,30 @@
+# app.py
+import json
+
 import gradio as gr
 import torch
 from torchvision import transforms
 from PIL import Image
-import json
 import os
-import datetime
 from model import get_mobilenet_model
 
-# 推理设置
+# 自动检测是否使用GPU
 device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# 加载模型
+weights_path = "best_model.pth"
 model, _, _ = get_mobilenet_model(num_classes=16)
-checkpoint = torch.load("best_model.pth", map_location=device, weights_only=False)
-model.load_state_dict(checkpoint['state_dict'])
+checkpoint = torch.load(weights_path, map_location=device, weights_only=False)
+model.load_state_dict(checkpoint['state_dict'])  # 注意这里是 checkpoint['state_dict']
 model.to(device)
 model.eval()
 
-# 加载类别信息
 with open("class_eg_to_cn.json", "r", encoding="utf-8") as f:
     class_eg_to_cn = json.load(f)
+
 with open("num_to_class.json", "r", encoding="utf-8") as f:
     idx2label = json.load(f)
+
 with open("class_inf.json", "r", encoding="utf-8") as f:
     plant_info = json.load(f)
 
@@ -29,88 +34,29 @@ transform = transforms.Compose([
     transforms.ToTensor(),
 ])
 
-# 存储检测记录（内存中）
-history = []
-
-
-def predict(image):
-    image_tensor = transform(image).unsqueeze(0).to(device)
+# 推理函数
+def predict(img):
+    image = transform(img).unsqueeze(0).to(device)
     with torch.no_grad():
-        output = model(image_tensor)
+        output = model(image)
         probs = torch.softmax(output, dim=1)
         conf, pred = torch.max(probs, 1)
-        label_en = idx2label[str(pred.item())]
-        label_cn = class_eg_to_cn[label_en]
-        intro = plant_info.get(label_en, "暂无介绍")
-
-        # 保存图像快照到本地（临时）
-        os.makedirs("history_imgs", exist_ok=True)
-        timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
-        image_path = f"history_imgs/{timestamp}.png"
-        image.save(image_path)
+        cn_label = class_eg_to_cn[idx2label[str(pred.item())]]
+        intro = plant_info.get(idx2label[str(pred.item())], "暂无介绍")
+        return cn_label, float(conf.item()), intro
 
-        # 添加到历史记录
-        history.append({
-            "image": image_path,
-            "label": label_cn,
-            "confidence": f"{conf.item():.2f}",
-            "time": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-        })
-
-        return label_cn, float(conf.item()), intro
-
-
-def get_history():
-    if not history:
-        return [], [], [], []
-    imgs = [entry["image"] for entry in history]
-    labels = [entry["label"] for entry in history]
-    confs = [entry["confidence"] for entry in history]
-    times = [entry["time"] for entry in history]
-    return imgs, labels, confs, times
-
-
-# UI
-with gr.Blocks(css="body { font-family: sans-serif; display: flex; }") as demo:
+# Gradio界面
+with gr.Blocks(css="body { background-color: #90caf9; font-family: sans-serif; }") as demo:
+    gr.Markdown("## 🌱 基于深度学习的植物识别系统", elem_id="title")
     with gr.Row():
-        with gr.Column(scale=1, min_width=200):  # Left navigation bar
-            gr.Markdown("## 🌱 功能导航", elem_id="nav")
-            btn1 = gr.Button("图像识别", elem_id="btn1")
-            btn2 = gr.Button("检测记录", elem_id="btn2")
-
+        with gr.Column(scale=2):
+            image_input = gr.Image(type="pil", label="选择图像")
+            btn = gr.Button("开始检测", elem_id="detect")
         with gr.Column(scale=3):
-            # 图像识别界面
-            with gr.Box() as img_recognition:
-                gr.Markdown("## 植物识别系统")
-                with gr.Row():
-                    with gr.Column(scale=2):
-                        img_input = gr.Image(type="pil", label="上传植物图像")
-                        btn = gr.Button("识别")
-                    with gr.Column(scale=3):
-                        out_label = gr.Textbox(label="识别类别")
-                        out_conf = gr.Textbox(label="置信度")
-                        out_desc = gr.Textbox(label="植物介绍", lines=5)
-                btn.click(fn=predict, inputs=img_input, outputs=[out_label, out_conf, out_desc])
-
-            # 检测记录界面
-            with gr.Box() as detection_history:
-                gr.Markdown("## 最近识别记录")
-                show_btn = gr.Button("刷新记录")
-                history_gallery = gr.Gallery(label="历史图像", elem_id="history_gallery", columns=1,  height="auto")
-                label_list = gr.Textbox(label="识别类别")
-                conf_list = gr.Textbox(label="置信度")
-                time_list = gr.Textbox(label="检测时间")
-
-                show_btn.click(fn=get_history, inputs=[], outputs=[history_gallery, label_list, conf_list, time_list])
-
-    # 显示第一个界面
-    img_recognition.visible = True
-    detection_history.visible = False
+            label_output = gr.Textbox(label="类别")
+            confidence_output = gr.Textbox(label="置信度")
+            description_output = gr.Textbox(label="相关介绍", lines=5)
 
-    # 设置按钮点击事件来切换页面
-    btn1.click(fn=lambda: (img_recognition.update(visible=True), detection_history.update(visible=False)), inputs=[],
-               outputs=[])
-    btn2.click(fn=lambda: (img_recognition.update(visible=False), detection_history.update(visible=True)), inputs=[],
-               outputs=[])
+    btn.click(fn=predict, inputs=image_input, outputs=[label_output, confidence_output, description_output])
 
-demo.launch()
+demo.launch()