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

# Load your Hugging Face model
model_id = "KabeerAmjad/food_classification_model"  # Replace with your actual model ID
model = AutoModelForImageClassification.from_pretrained(model_id)
feature_extractor = AutoFeatureExtractor.from_pretrained(model_id)

# Define the prediction function
def classify_image(img):
    inputs = feature_extractor(images=img, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.softmax(outputs.logits, dim=-1)
    
    # Get the label with the highest probability
    top_label = model.config.id2label[probs.argmax().item()]
    return top_label

# Create the Gradio interface
iface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="Food Image Classification",
    description="Upload an image to classify if it’s an apple pie, etc."
)

# Launch the app
iface.launch()