Update app.py

app.py

@@ -1,32 +1,29 @@
+# ✅ Import required libraries
 from PIL import Image
 import torch
 import torchvision.transforms as transforms
 from torchvision import models
 import gradio as gr
-from rembg import remove  # 🆕 Background removal
-import io
-import datetime
+import webbrowser
 
 # 🔧 Set device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# 📦 Load your fine-tuned model
+# 📦 Load the trained model
 model = models.resnet50(pretrained=False)
 model.fc = torch.nn.Linear(model.fc.in_features, 2)  # 2 classes: Edible, Poisonous
 model.load_state_dict(torch.load("resnet_mushroom_classifier.pth", map_location=device))
 model = model.to(device)
 model.eval()
 
-# 🏷️ Class names
+# 🏷️ Class names and species mapping
 class_names = ['Edible', 'Poisonous']
-
-# 🍄 Mapping for more detailed species
 mushroom_species = {
     "Edible": "Possible species:\n• Amanita citrina\n• Russula delica\n• Phaeogyroporus portentosus",
     "Poisonous": "Possible species:\n• Amanita phalloides\n• Inocybe rimosa"
 }
 
-# 🎨 Image preprocessing (must match training)
+# 🎨 Image preprocessing
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
@@ -34,16 +31,13 @@ transform = transforms.Compose([
                          [0.229, 0.224, 0.225])
 ])
 
-# 🧠 Prediction function
+# 🧠 Classification function with validation
+CONFIDENCE_THRESHOLD = 85.0  # Minimum confidence considered safe enough to show suggestion
+
 def classify_mushroom(image: Image.Image):
     try:
         image = image.convert("RGB")
-
-        # 🆕 Remove background
-        bg_removed = remove(image)
-
-        # 🆗 Prepare cleaned image for model
-        tensor = transform(bg_removed).unsqueeze(0).to(device)
+        tensor = transform(image).unsqueeze(0).to(device)
 
         with torch.no_grad():
             outputs = model(tensor)
@@ -51,42 +45,54 @@ def classify_mushroom(image: Image.Image):
             label = class_names[predicted.item()]
             score = torch.softmax(outputs, dim=1)[0][predicted.item()].item() * 100
 
-        suggestion = mushroom_species[label]
+        suggestion = mushroom_species[label] if score >= CONFIDENCE_THRESHOLD else "Confidence too low to suggest species."
 
-        return label, "กินได้" if label == "Edible" else "พิษ", f"{score:.2f}%", suggestion, bg_removed
+        # Check confidence and warn
+        if score < CONFIDENCE_THRESHOLD:
+            label = "Uncertain"
+
+        return label, ("กินได้" if label == "Edible" else ("พิษ" if label == "Poisonous" else "ไม่มั่นใจ")), f"{score:.2f}%", suggestion
 
     except Exception as e:
         print(f"❌ Error: {e}")
-        return "Error", "ผิดพลาด", "N/A", "N/A", image
+        return "Error", "ผิดพลาด", "N/A", "Invalid image. Please upload a valid mushroom photo."
 
-# 📆 Version Info
-version_info = "**Version:** 1.0.1  \n**Last updated:** " + datetime.datetime.now().strftime("%Y-%m-%d")
+# 🔗 Open user manual
+MANUAL_URL = "https://drive.google.com/drive/folders/19lUCEaLstrRjCzqpDlWErhRd1EXWUGbf?usp=sharing"
+def open_manual():
+    webbrowser.open(MANUAL_URL)
+    return "", "", "", "Opening user manual..."
 
 # 🎛️ Gradio UI
 if __name__ == "__main__":
     with gr.Blocks() as demo:
         gr.Markdown("## 🍄 Mushroom Safety Classifier")
-        gr.Markdown("Upload a mushroom photo to check if it’s edible or poisonous.\nอัปโหลดรูปเห็ดเพื่อทำนายว่าเห็ดกินได้หรือมีพิษ")
+        gr.Markdown("Upload a mushroom photo or use your camera to check if it’s edible or poisonous.\nอัปโหลดรูปเห็ดหรือใช้กล้องเพื่อตรวจสอบว่าเห็ดกินได้หรือมีพิษ")
 
         with gr.Row():
+            image_input = gr.Image(type="pil", label="📷 Upload or Capture Mushroom Image", source="upload", tool="editor")
             with gr.Column():
-                image_input = gr.Image(type="pil", label="📷 Upload Mushroom Image")
-                cleaned_image = gr.Image(label="🧼 Background Removed Image")  # 🆕 Show cleaned image
-            with gr.Column():
-                label_en = gr.Textbox(label="🧠 Prediction (English)")
-                label_th = gr.Textbox(label="🗣️ คำทำนาย (ภาษาไทย)")
+                label_en = gr.Textbox(label="🧬 Prediction (English)")
+                label_th = gr.Textbox(label="🔁 คำทำนาย (ภาษาไทย)")
                 confidence = gr.Textbox(label="📶 Confidence Score")
                 label_hint = gr.Textbox(label="🏷️ Likely Species (Based on Training Data)")
 
         classify_btn = gr.Button("🔍 Classify")
+        manual_btn = gr.Button("Open User Manual 📄")
 
         classify_btn.click(
             fn=classify_mushroom,
             inputs=image_input,
-            outputs=[label_en, label_th, confidence, label_hint, cleaned_image]
+            outputs=[label_en, label_th, confidence, label_hint]
+        )
+
+        manual_btn.click(
+            fn=open_manual,
+            inputs=[],
+            outputs=[label_en, label_th, confidence, label_hint]
         )
 
         gr.Markdown("---")
-        gr.Markdown(version_info)  # 📅 Version and update info
+        gr.Markdown("App version: 1.0.1 | Updated: August 2025")
 
     demo.launch()