Spaces:
Sleeping
Sleeping
import gradio as gr | |
from transformers import AutoModelForImageClassification, AutoProcessor | |
import torch | |
from PIL import Image | |
# Load your model and processor outside the function to avoid reloading them on each function call | |
model_name = "Khadidja22/my_awesome_food_model" | |
model = AutoModelForImageClassification.from_pretrained(model_name) | |
processor = AutoProcessor.from_pretrained(model_name) | |
def classify_image(uploaded_image): | |
# Process the uploaded image | |
inputs = processor(images=uploaded_image, return_tensors="pt") | |
# Predict | |
with torch.no_grad(): | |
outputs = model(**inputs) | |
logits = outputs.logits | |
# Get the highest probability label | |
predicted_label_idx = logits.argmax(-1).item() | |
predicted_label = model.config.id2label[predicted_label_idx] | |
return predicted_label | |
iface = gr.Interface(fn=classify_image, | |
inputs=gr.Image(), | |
outputs="text", | |
title="Food Classification", | |
description="Upload an image of food, and the model will classify it.") | |
iface.launch() |