Update app.py

app.py

@@ -1,11 +1,9 @@
 import gradio as gr
 from transformers import AutoModelForImageClassification, AutoFeatureExtractor
-from PIL import Image
-import requests
 import torch
 
-# Hugging Face se model ka repo name
-model_name = "google/vit-base-patch16-224"
+# Food-specific model
+model_name = "nateraw/food101"
 
 # Model aur feature extractor load karo
 model = AutoModelForImageClassification.from_pretrained(model_name)
@@ -19,6 +17,11 @@ def classify_image(image):
     pred_idx = torch.argmax(probs)
     pred_label = model.config.id2label[pred_idx.item()]
     confidence = probs[0][pred_idx].item()
+
+    # Agar confidence bahut low hai to reject kar do
+    if confidence < 0.30:
+        return "Not a food/vegetable"
+
     return f"{pred_label} ({confidence*100:.2f}%)"
 
 # Gradio interface
@@ -26,8 +29,8 @@ demo = gr.Interface(
     fn=classify_image,
     inputs=gr.Image(type="pil"),
     outputs="text",
-    title="Food Classifier",
-    description="Upload any food image and get the prediction!"
+    title="Food/Vegetable Classifier",
+    description="Upload any food or vegetable image. If it's not food, you'll get a 'Not a food/vegetable' message."
 )
 
 demo.launch()