Update app.py

# ✅ 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

@@ -5,8 +5,6 @@ import torchvision.transforms as transforms
 from torchvision import models
 import gradio as gr
 from rembg import remove  # Background removal
-from transformers import pipeline  # For non-mushroom detection
-from torchvision.models import ResNet50_Weights
 
 # 🔧 Set device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -14,8 +12,8 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 # 📦 Load your fine-tuned model
 model = models.resnet50(weights=None)  # No pretrained weights
 model.fc = torch.nn.Sequential(
-    torch.nn.Dropout(0.5),
-    torch.nn.Linear(model.fc.in_features, 2)
+    torch.nn.Dropout(0.5),  # 🔧 Dropout to reduce overfitting
+    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)
@@ -36,26 +34,11 @@ 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
+# 🧠 Classification function with confidence filtering
 CONFIDENCE_THRESHOLD = 85.0  # Minimum confidence considered safe enough to show suggestion
 
 def classify_mushroom(image: Image.Image):
     try:
-        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)
@@ -108,7 +91,7 @@ if __name__ == "__main__":
         gr.Markdown("---")
         gr.Markdown(manual_link)  # 🆕 Display clickable user manual
         gr.Markdown(feedback_link)  # 🆕 Display clickable user manual
-        gr.Markdown("App version: 1.1.0 | Updated: August 2025")
+        gr.Markdown("App version: 1.1.1 | Updated: August 2025")
 
     demo.launch()