import gradio as gr
from transformers import AutoModelForImageClassification, AutoFeatureExtractor
import torch

# Food-specific model
model_name = "nateraw/food101"

# Model aur feature extractor load karo
model = AutoModelForImageClassification.from_pretrained(model_name)
feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)

# Prediction function
def classify_image(image):
    inputs = feature_extractor(images=image, return_tensors="pt")
    outputs = model(**inputs)
    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
    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
demo = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs="text",
    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()