Update app.py

app.py

@@ -1,28 +1,29 @@
 import gradio as gr
 import torch
 import torch.nn as nn
-from torchvision import transforms, models
+from torchvision import transforms
 import pickle
 from resnest.torch import resnest50
+from rembg import remove
+from PIL import Image
+import io
 
+# 加载类别名称
 with open('class_names.pkl', 'rb') as f:
     class_names = pickle.load(f)
 
-# 加载训练好的模型
+# 初始化模型
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-model = resnest50(pretrained=None)
+model = resnest50(pretrained=False)
 model.fc = nn.Sequential(
     nn.Dropout(0.2),
     nn.Linear(model.fc.in_features, len(class_names))
 )
-
-# 加载模型权重
-model.load_state_dict(torch.load('best_model.pth', map_location=device, weights_only=True))
+model.load_state_dict(torch.load('best_model.pth', map_location=device))
 model = model.to(device)
 model.eval()
 
-# 定义与训练时相同的预处理流程
+# 预处理流程
 preprocess = transforms.Compose([
     transforms.Resize((100, 100)),
     transforms.ToTensor(),
@@ -30,23 +31,42 @@ preprocess = transforms.Compose([
 ])
 
 
-def predict_image(img):
-    img = img.convert('RGB')
+def remove_background(img):
+    """使用rembg去除背景并添加白色背景"""
+    # 转换图像为字节流
+    img_byte_arr = io.BytesIO()
+    img.save(img_byte_arr, format='PNG')
+    img_bytes = img_byte_arr.getvalue()
+
+    # 去除背景
+    removed_bg_bytes = remove(img_bytes)
+
+    # 转换为PIL图像并处理透明度
+    removed_bg_img = Image.open(io.BytesIO(removed_bg_bytes)).convert('RGBA')
+
+    # 创建白色背景
+    white_bg = Image.new('RGBA', removed_bg_img.size, (255, 255, 255, 255))
+    combined = Image.alpha_composite(white_bg, removed_bg_img)
+    return combined.convert('RGB')
 
-    # 应用预处理
-    input_tensor = preprocess(img)
 
-    # 添加批次维度并移动到设备
+def predict_image(img, remove_bg=False):
+    """分类预测主函数"""
+    # 根据选择处理图像
+    if remove_bg:
+        processed_img = remove_background(img)
+    else:
+        processed_img = img.convert('RGB')  # 确保为RGB格式
+
+    # 预处理
+    input_tensor = preprocess(processed_img)
     input_batch = input_tensor.unsqueeze(0).to(device)
 
     # 预测
     with torch.no_grad():
         output = model(input_batch)
 
-    # 计算概率
     probabilities = torch.nn.functional.softmax(output[0], dim=0)
-
-    # 获取前3个预测结果
     top3_probs, top3_indices = torch.topk(probabilities, 3)
 
     results = {
@@ -54,38 +74,66 @@ def predict_image(img):
         for p, i in zip(top3_probs, top3_indices)
     }
 
-    # 获取最佳预测结果
+    # 记录结果
     best_class = class_names[top3_indices[0]]
     best_conf = top3_probs[0].item() * 100
-    
-    return best_class, best_conf, results
 
-# 创建Gradio界面
+    return processed_img, best_class, f"{best_conf:.2f}%", results
+
+
 def create_interface():
     examples = [
         "r0_0_100.jpg",
-        "r0_18_100.jpg"
+        "r0_18_100.jpg",
+        "9_100.jpg",
+        "127_100.jpg",
+        "r0_1_100.jpg",
     ]
 
     with gr.Blocks(title="Fruit Classification", theme=gr.themes.Soft()) as demo:
-        gr.Markdown("# 🍎 水果识别系统")
+        gr.Markdown("""
+        # 🍎 智能水果识别系统
+        """)
 
+        # 新增：模式选择卡片（视觉强化）
+        with gr.Row():
+            with gr.Column(scale=3):
+                with gr.Group():
+                    gr.Markdown("### ⚙️ 处理模式选择")
+                    with gr.Row():
+                        bg_removal = gr.Checkbox(
+                            label="背景去除",
+                            value=False,
+                            interactive=True
+                        )
+
+        # 主操作区域
         with gr.Row():
             with gr.Column():
-                image_input = gr.Image(type="pil", label="上传图像")
-                gr.Examples(examples=examples, inputs=image_input)
-                submit_btn = gr.Button("分类", variant="primary")
+                original_image = gr.Image(label="📤 上传图片", type="pil")
+                gr.Examples(examples=examples, inputs=original_image)
+                submit_btn = gr.Button("🚀 开始识别", variant="primary")
+
+                # 添加模式说明提示
+                gr.Markdown("""
+                <div style="background: #f3f4f6; padding: 15px; border-radius: 8px; margin-top: 10px">
+                <b>💡 使用建议：</b><br>
+                • 上传图片：选择一张图片，点击'开始识别'按钮<br>
+                • 勾选背景去除：适合杂乱背景的图片（识别更准确）<br>
+                • 不勾选：适合纯色背景的图片（速度更快）
+                </div>
+                """)
 
             with gr.Column():
-                best_pred = gr.Textbox(label="预测结果")
-                confidence = gr.Textbox(label="置信度")
-                full_results = gr.Label(label="Top 3", num_top_classes=3)
+                processed_image = gr.Image(label="🖼️ 处理后图片", interactive=False)
+                best_pred = gr.Textbox(label="🔍 识别结果")
+                confidence = gr.Textbox(label="📊 置信度")
+                full_results = gr.Label(label="🏆 Top 3 可能结果", num_top_classes=3)
 
-        # ‘分类’按钮点击事件
         submit_btn.click(
             fn=predict_image,
-            inputs=image_input,
-            outputs=[best_pred, confidence, full_results]
+            inputs=[original_image, bg_removal],
+            outputs=[processed_image, best_pred, confidence, full_results]
         )
 
     return demo
@@ -93,4 +141,4 @@ def create_interface():
 
 if __name__ == "__main__":
     interface = create_interface()
-    interface.launch(share=False)
+    interface.launch(share=True)