Update app.py

version 7/8/2025


# ✅ Training Script Summary:
# - Use torchvision.datasets.ImageFolder with folders "Edible/" and "Poisonous/"
# - Apply data augmentation (flip, jitter, rotate, etc.)
# - Split into train/val and load with DataLoader
# - Use pretrained ResNet50 with replaced final FC layer
# - Train using CrossEntropyLoss and Adam optimizer
# - Save model using torch.save(...)

app.py

@@ -2,15 +2,14 @@
 from PIL import Image
 import torch
 import torchvision.transforms as transforms
-from torchvision import models
-import gradio as gr
-import webbrowser
-from rembg import remove  # 🆕 Background removal
+from torchvision import models\import gradio as gr
+from rembg import remove  # Background removal
+from transformers import pipeline  # For non-mushroom detection
 
 # 🔧 Set device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# 📦 Load the trained model
+# 📦 Load your fine-tuned 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))
@@ -32,13 +31,28 @@ transform = transforms.Compose([
                          [0.229, 0.224, 0.225])
 ])
 
+# 🔍 Pretrained image classifier for screening non-mushroom
+label_detector = pipeline("image-classification", model="microsoft/resnet-50")
+mushroom_keywords = ["mushroom", "agaric", "amanita", "fungus", "earthstar", "toadstool"]
+
+def is_mushroom_image(image):
+    try:
+        result = label_detector(image)
+        top_label = result[0]["label"].lower()
+        return any(keyword in top_label for keyword in mushroom_keywords)
+    except:
+        return False
+
 # 🧠 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("RGBA")  # Required format for rembg
-        image_no_bg = remove(image).convert("RGB")  # 🆕 Remove background
+        if not is_mushroom_image(image):
+            return "Not a mushroom", "ไม่ใช่เห็ด", "0.00%", "Please upload a mushroom photo."
+
+        image = image.convert("RGBA")
+        image_no_bg = remove(image).convert("RGB")
         tensor = transform(image_no_bg).unsqueeze(0).to(device)
 
         with torch.no_grad():
@@ -49,7 +63,6 @@ def classify_mushroom(image: Image.Image):
 
         suggestion = mushroom_species[label] if score >= CONFIDENCE_THRESHOLD else "Confidence too low to suggest species."
 
-        # Check confidence and warn
         if score < CONFIDENCE_THRESHOLD:
             label = "Uncertain"
 
@@ -59,9 +72,6 @@ def classify_mushroom(image: Image.Image):
         print(f"❌ Error: {e}")
         return "Error", "ผิดพลาด", "N/A", "Invalid image. Please upload a valid mushroom photo."
 
-# 🔗 Open user manual (link version) 🆕
-MANUAL_URL = "https://drive.google.com/drive/folders/19lUCEaLstrRjCzqpDlWErhRd1EXWUGbf?usp=sharing"
-manual_link = f"<a href=\"{MANUAL_URL}\" target=\"_blank\">📄 Open User Manual</a>"
 # 🎛️ Gradio UI
 if __name__ == "__main__":
     with gr.Blocks() as demo:
@@ -69,7 +79,7 @@ if __name__ == "__main__":
         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")  # 🆕 safer source usage
+            image_input = gr.Image(type="pil", label="📷 Upload or Capture Mushroom Image")
             with gr.Column():
                 label_en = gr.Textbox(label="🧬 Prediction (English)")
                 label_th = gr.Textbox(label="🔁 คำทำนาย (ภาษาไทย)")
@@ -85,7 +95,8 @@ if __name__ == "__main__":
         )
 
         gr.Markdown("---")
-        gr.Markdown(manual_link)  # 🆕 Display clickable user manual
-        gr.Markdown("App version: 1.0.2 | Updated: August 2025")  # 🆕 Version bump
+        gr.Markdown("App version: 1.1.0 | Updated: August 2025")
 
     demo.launch()
+
+